1 files changed, 63 insertions, 30 deletions

diff --git a/absl/flags/reflection.cc b/absl/flags/reflection.cc
index d7060221..0c761101 100644
--- a/absl/flags/reflection.cc
+++ b/absl/flags/reflection.cc
@@ -17,6 +17,7 @@
 
 #include <assert.h>
 
+#include <atomic>
 #include <map>
 #include <string>
 
@@ -49,28 +50,30 @@ class FlagRegistry {
   ~FlagRegistry() = default;
 
   // Store a flag in this registry. Takes ownership of *flag.
-  void RegisterFlag(CommandLineFlag& flag);
+  void RegisterFlag(CommandLineFlag& flag, const char* filename);
 
   void Lock() ABSL_EXCLUSIVE_LOCK_FUNCTION(lock_) { lock_.Lock(); }
   void Unlock() ABSL_UNLOCK_FUNCTION(lock_) { lock_.Unlock(); }
 
   // Returns the flag object for the specified name, or nullptr if not found.
   // Will emit a warning if a 'retired' flag is specified.
-  CommandLineFlag* FindFlagLocked(absl::string_view name);
+  CommandLineFlag* FindFlag(absl::string_view name);
 
   static FlagRegistry& GlobalRegistry();  // returns a singleton registry
 
  private:
   friend class flags_internal::FlagSaverImpl;  // reads all the flags in order
                                                // to copy them
-  friend void ForEachFlagUnlocked(
-      std::function<void(CommandLineFlag&)> visitor);
+  friend void ForEachFlag(std::function<void(CommandLineFlag&)> visitor);
+  friend void FinalizeRegistry();
 
-  // The map from name to flag, for FindFlagLocked().
+  // The map from name to flag, for FindFlag().
   using FlagMap = std::map<absl::string_view, CommandLineFlag*>;
   using FlagIterator = FlagMap::iterator;
   using FlagConstIterator = FlagMap::const_iterator;
   FlagMap flags_;
+  std::vector<CommandLineFlag*> flat_flags_;
+  std::atomic<bool> finalized_flags_{false};
 
   absl::Mutex lock_;
 
@@ -79,15 +82,6 @@ class FlagRegistry {
   FlagRegistry& operator=(const FlagRegistry&);
 };
 
-CommandLineFlag* FlagRegistry::FindFlagLocked(absl::string_view name) {
-  FlagConstIterator i = flags_.find(name);
-  if (i == flags_.end()) {
-    return nullptr;
-  }
-
-  return i->second;
-}
-
 namespace {
 
 class FlagRegistryLock {
@@ -101,8 +95,37 @@ class FlagRegistryLock {
 
 }  // namespace
 
-void FlagRegistry::RegisterFlag(CommandLineFlag& flag) {
+CommandLineFlag* FlagRegistry::FindFlag(absl::string_view name) {
+  if (finalized_flags_.load(std::memory_order_acquire)) {
+    // We could save some gcus here if we make `Name()` be non-virtual.
+    // We could move the `const char*` name to the base class.
+    auto it = std::partition_point(
+        flat_flags_.begin(), flat_flags_.end(),
+        [=](CommandLineFlag* f) { return f->Name() < name; });
+    if (it != flat_flags_.end() && (*it)->Name() == name) return *it;
+  }
+
+  FlagRegistryLock frl(*this);
+  auto it = flags_.find(name);
+  return it != flags_.end() ? it->second : nullptr;
+}
+
+void FlagRegistry::RegisterFlag(CommandLineFlag& flag, const char* filename) {
+  if (filename != nullptr &&
+      flag.Filename() != GetUsageConfig().normalize_filename(filename)) {
+    flags_internal::ReportUsageError(
+        absl::StrCat(
+            "Inconsistency between flag object and registration for flag '",
+            flag.Name(),
+            "', likely due to duplicate flags or an ODR violation. Relevant "
+            "files: ",
+            flag.Filename(), " and ", filename),
+        true);
+    std::exit(1);
+  }
+
   FlagRegistryLock registry_lock(*this);
+
   std::pair<FlagIterator, bool> ins =
       flags_.insert(FlagMap::value_type(flag.Name(), &flag));
   if (ins.second == false) {  // means the name was already in the map
@@ -152,27 +175,39 @@ FlagRegistry& FlagRegistry::GlobalRegistry() {
 
 // --------------------------------------------------------------------
 
-void ForEachFlagUnlocked(std::function<void(CommandLineFlag&)> visitor) {
+void ForEachFlag(std::function<void(CommandLineFlag&)> visitor) {
   FlagRegistry& registry = FlagRegistry::GlobalRegistry();
-  for (FlagRegistry::FlagConstIterator i = registry.flags_.begin();
-       i != registry.flags_.end(); ++i) {
-    visitor(*i->second);
+
+  if (registry.finalized_flags_.load(std::memory_order_acquire)) {
+    for (const auto& i : registry.flat_flags_) visitor(*i);
   }
-}
 
-void ForEachFlag(std::function<void(CommandLineFlag&)> visitor) {
-  FlagRegistry& registry = FlagRegistry::GlobalRegistry();
   FlagRegistryLock frl(registry);
-  ForEachFlagUnlocked(visitor);
+  for (const auto& i : registry.flags_) visitor(*i.second);
 }
 
 // --------------------------------------------------------------------
 
-bool RegisterCommandLineFlag(CommandLineFlag& flag) {
-  FlagRegistry::GlobalRegistry().RegisterFlag(flag);
+bool RegisterCommandLineFlag(CommandLineFlag& flag, const char* filename) {
+  FlagRegistry::GlobalRegistry().RegisterFlag(flag, filename);
   return true;
 }
 
+void FinalizeRegistry() {
+  auto& registry = FlagRegistry::GlobalRegistry();
+  FlagRegistryLock frl(registry);
+  if (registry.finalized_flags_.load(std::memory_order_relaxed)) {
+    // Was already finalized. Ignore the second time.
+    return;
+  }
+  registry.flat_flags_.reserve(registry.flags_.size());
+  for (const auto& f : registry.flags_) {
+    registry.flat_flags_.push_back(f.second);
+  }
+  registry.flags_.clear();
+  registry.finalized_flags_.store(true, std::memory_order_release);
+}
+
 // --------------------------------------------------------------------
 
 namespace {
@@ -244,7 +279,7 @@ void Retire(const char* name, FlagFastTypeId type_id, char* buf) {
   static_assert(alignof(RetiredFlagObj) == kRetiredFlagObjAlignment, "");
   auto* flag = ::new (static_cast<void*>(buf))
       flags_internal::RetiredFlagObj(name, type_id);
-  FlagRegistry::GlobalRegistry().RegisterFlag(*flag);
+  FlagRegistry::GlobalRegistry().RegisterFlag(*flag, nullptr);
 }
 
 // --------------------------------------------------------------------
@@ -298,9 +333,7 @@ CommandLineFlag* FindCommandLineFlag(absl::string_view name) {
   if (name.empty()) return nullptr;
   flags_internal::FlagRegistry& registry =
       flags_internal::FlagRegistry::GlobalRegistry();
-  flags_internal::FlagRegistryLock frl(registry);
-
-  return registry.FindFlagLocked(name);
+  return registry.FindFlag(name);
 }
 
 // --------------------------------------------------------------------
@@ -308,7 +341,7 @@ CommandLineFlag* FindCommandLineFlag(absl::string_view name) {
 absl::flat_hash_map<absl::string_view, absl::CommandLineFlag*> GetAllFlags() {
   absl::flat_hash_map<absl::string_view, absl::CommandLineFlag*> res;
   flags_internal::ForEachFlag([&](CommandLineFlag& flag) {
-    res.insert({flag.Name(), &flag});
+    if (!flag.IsRetired()) res.insert({flag.Name(), &flag});
   });
   return res;
 }