Add LifetimeSymbolTable::LookupLifetimeAndMaybeDeclare().
PiperOrigin-RevId: 413108197
diff --git a/lifetime_annotations/lifetime_symbol_table.cc b/lifetime_annotations/lifetime_symbol_table.cc
index c5b6b31..040a0a7 100644
--- a/lifetime_annotations/lifetime_symbol_table.cc
+++ b/lifetime_annotations/lifetime_symbol_table.cc
@@ -5,6 +5,7 @@
#include "lifetime_annotations/lifetime_symbol_table.h"
#include <optional>
+#include <string>
namespace devtools_rust {
@@ -50,4 +51,44 @@
return iter->second;
}
+static std::string NameFromIndex(int index) {
+ int num_chars = 1;
+ // Number of combinations that are possible with `num_chars` characters.
+ int num_combinations = 26;
+ while (index >= num_combinations) {
+ index -= num_combinations;
+ ++num_chars;
+ num_combinations *= 26;
+ }
+ std::string name;
+ name.reserve(num_chars);
+ for (int i = 0; i < num_chars; ++i) {
+ name.insert(0, 1, 'a' + index % 26);
+ index /= 26;
+ }
+ return name;
+}
+
+llvm::StringRef LifetimeSymbolTable::LookupLifetimeAndMaybeDeclare(
+ Lifetime lifetime) {
+ if (lifetime == Lifetime::Static()) {
+ return "static";
+ }
+
+ auto lifetime_to_name_iter = lifetime_to_name_.find(lifetime);
+ if (lifetime_to_name_iter != lifetime_to_name_.end()) {
+ return lifetime_to_name_iter->second;
+ }
+
+ while (true) {
+ std::string name = NameFromIndex(next_name_index_++);
+ auto [name_to_lifetime_iter, inserted] =
+ name_to_lifetime_.try_emplace(name, lifetime);
+ if (inserted) {
+ lifetime_to_name_[lifetime] = name;
+ return name_to_lifetime_iter->first();
+ }
+ }
+}
+
} // namespace devtools_rust
diff --git a/lifetime_annotations/lifetime_symbol_table.h b/lifetime_annotations/lifetime_symbol_table.h
index c9599f7..d53ca4d 100644
--- a/lifetime_annotations/lifetime_symbol_table.h
+++ b/lifetime_annotations/lifetime_symbol_table.h
@@ -35,9 +35,18 @@
// If `lifetime` is `Lifetime::Static()`, returns "static".
std::optional<llvm::StringRef> LookupLifetime(Lifetime lifetime) const;
+ // Looks up a lifetime in the symbol table and inserts a new autogenerated
+ // name for it if the lifetime was not found.
+ // If `lifetime` is `Lifetime::Static()`, returns "static".
+ // The autogenerated name will be the first name from the sequence
+ // "a", ..., "z", "aa", "ab", "ac", ... that is not yet contained in the
+ // symbol table.
+ llvm::StringRef LookupLifetimeAndMaybeDeclare(Lifetime lifetime);
+
private:
llvm::StringMap<Lifetime> name_to_lifetime_;
llvm::DenseMap<Lifetime, std::string> lifetime_to_name_;
+ int next_name_index_ = 0;
};
} // namespace devtools_rust
diff --git a/lifetime_annotations/lifetime_symbol_table_test.cc b/lifetime_annotations/lifetime_symbol_table_test.cc
index 9453778..9e2f6b9 100644
--- a/lifetime_annotations/lifetime_symbol_table_test.cc
+++ b/lifetime_annotations/lifetime_symbol_table_test.cc
@@ -40,5 +40,30 @@
EXPECT_EQ(table.LookupLifetime(b), "b");
}
+TEST(LifetimeSymbolTableTest, LookupLifetimeAndMaybeDeclare) {
+ {
+ LifetimeSymbolTable table;
+
+ table.LookupNameAndMaybeDeclare("a");
+ table.LookupNameAndMaybeDeclare("c");
+
+ EXPECT_EQ(table.LookupLifetimeAndMaybeDeclare(Lifetime::CreateVariable()),
+ "b");
+ EXPECT_EQ(table.LookupLifetimeAndMaybeDeclare(Lifetime::CreateVariable()),
+ "d");
+ }
+
+ {
+ LifetimeSymbolTable table;
+ for (int i = 0; i < 26; ++i) {
+ table.LookupLifetimeAndMaybeDeclare(Lifetime::CreateVariable());
+ }
+ EXPECT_EQ(table.LookupLifetimeAndMaybeDeclare(Lifetime::CreateVariable()),
+ "aa");
+ EXPECT_EQ(table.LookupLifetimeAndMaybeDeclare(Lifetime::CreateVariable()),
+ "ab");
+ }
+}
+
} // namespace
} // namespace devtools_rust
