2 files changed, 162 insertions, 135 deletions

diff --git a/absl/flags/internal/flag.cc b/absl/flags/internal/flag.cc
index a45e883e..56a5ed2b 100644
--- a/absl/flags/internal/flag.cc
+++ b/absl/flags/internal/flag.cc
@@ -74,6 +74,59 @@ class MutexRelock {
 
 }  // namespace
 
+///////////////////////////////////////////////////////////////////////////////
+// Persistent state of the flag data.
+
+class FlagImpl;
+
+class FlagState : public flags_internal::FlagStateInterface {
+ public:
+  template <typename V>
+  FlagState(FlagImpl* flag_impl, const V& v, bool modified,
+            bool on_command_line, int64_t counter)
+      : flag_impl_(flag_impl),
+        value_(v),
+        modified_(modified),
+        on_command_line_(on_command_line),
+        counter_(counter) {}
+
+  ~FlagState() override {
+    if (flag_impl_->ValueStorageKind() != FlagValueStorageKind::kHeapAllocated)
+      return;
+    flags_internal::Delete(flag_impl_->op_, value_.dynamic);
+  }
+
+ private:
+  friend class FlagImpl;
+
+  // Restores the flag to the saved state.
+  void Restore() const override {
+    if (!flag_impl_->RestoreState(*this)) return;
+
+    ABSL_INTERNAL_LOG(
+        INFO, absl::StrCat("Restore saved value of ", flag_impl_->Name(),
+                           " to: ", flag_impl_->CurrentValue()));
+  }
+
+  // Flag and saved flag data.
+  FlagImpl* flag_impl_;
+  union SavedValue {
+    explicit SavedValue(void* v) : dynamic(v) {}
+    explicit SavedValue(int64_t v) : one_word(v) {}
+    explicit SavedValue(flags_internal::AlignedTwoWords v) : two_words(v) {}
+
+    void* dynamic;
+    int64_t one_word;
+    flags_internal::AlignedTwoWords two_words;
+  } value_;
+  bool modified_;
+  bool on_command_line_;
+  int64_t counter_;
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Flag implementation, which does not depend on flag value type.
+
 void FlagImpl::Init() {
   new (&data_guard_) absl::Mutex;
 
@@ -86,13 +139,13 @@ void FlagImpl::Init() {
       break;
     case FlagValueStorageKind::kOneWordAtomic: {
       int64_t atomic_value;
-      std::memcpy(&atomic_value, init_value.get(), Sizeof(op_));
+      std::memcpy(&atomic_value, init_value.get(), flags_internal::Sizeof(op_));
       value_.one_word_atomic.store(atomic_value, std::memory_order_release);
       break;
     }
     case FlagValueStorageKind::kTwoWordsAtomic: {
       AlignedTwoWords atomic_value{0, 0};
-      std::memcpy(&atomic_value, init_value.get(), Sizeof(op_));
+      std::memcpy(&atomic_value, init_value.get(), flags_internal::Sizeof(op_));
       value_.two_words_atomic.store(atomic_value, std::memory_order_release);
       break;
     }
@@ -145,17 +198,17 @@ std::unique_ptr<void, DynValueDeleter> FlagImpl::MakeInitValue() const {
 void FlagImpl::StoreValue(const void* src) {
   switch (ValueStorageKind()) {
     case FlagValueStorageKind::kHeapAllocated:
-      Copy(op_, src, value_.dynamic);
+      flags_internal::Copy(op_, src, value_.dynamic);
       break;
     case FlagValueStorageKind::kOneWordAtomic: {
       int64_t one_word_val;
-      std::memcpy(&one_word_val, src, Sizeof(op_));
+      std::memcpy(&one_word_val, src, flags_internal::Sizeof(op_));
       value_.one_word_atomic.store(one_word_val, std::memory_order_release);
       break;
     }
     case FlagValueStorageKind::kTwoWordsAtomic: {
       AlignedTwoWords two_words_val{0, 0};
-      std::memcpy(&two_words_val, src, Sizeof(op_));
+      std::memcpy(&two_words_val, src, flags_internal::Sizeof(op_));
       value_.two_words_atomic.store(two_words_val, std::memory_order_release);
       break;
     }
@@ -172,11 +225,17 @@ std::string FlagImpl::Filename() const {
   return flags_internal::GetUsageConfig().normalize_filename(filename_);
 }
 
+absl::string_view FlagImpl::Typename() const { return ""; }
+
 std::string FlagImpl::Help() const {
   return HelpSourceKind() == FlagHelpKind::kLiteral ? help_.literal
                                                     : help_.gen_func();
 }
 
+FlagStaticTypeId FlagImpl::TypeId() const {
+  return flags_internal::StaticTypeId(op_);
+}
+
 bool FlagImpl::IsModified() const {
   absl::MutexLock l(DataGuard());
   return modified_;
@@ -195,10 +254,10 @@ std::string FlagImpl::DefaultValue() const {
 }
 
 std::string FlagImpl::CurrentValue() const {
-  DataGuard();  // Make sure flag initialized
+  auto* guard = DataGuard();  // Make sure flag initialized
   switch (ValueStorageKind()) {
     case FlagValueStorageKind::kHeapAllocated: {
-      absl::MutexLock l(DataGuard());
+      absl::MutexLock l(guard);
       return flags_internal::Unparse(op_, value_.dynamic);
     }
     case FlagValueStorageKind::kOneWordAtomic: {
@@ -250,23 +309,53 @@ void FlagImpl::InvokeCallback() const {
   cb();
 }
 
-bool FlagImpl::RestoreState(const void* value, bool modified,
-                            bool on_command_line, int64_t counter) {
-  {
-    absl::MutexLock l(DataGuard());
+std::unique_ptr<FlagStateInterface> FlagImpl::SaveState() {
+  absl::MutexLock l(DataGuard());
 
-    if (counter_ == counter) return false;
+  bool modified = modified_;
+  bool on_command_line = on_command_line_;
+  switch (ValueStorageKind()) {
+    case FlagValueStorageKind::kHeapAllocated: {
+      return absl::make_unique<FlagState>(
+          this, flags_internal::Clone(op_, value_.dynamic), modified,
+          on_command_line, counter_);
+    }
+    case FlagValueStorageKind::kOneWordAtomic: {
+      return absl::make_unique<FlagState>(
+          this, value_.one_word_atomic.load(std::memory_order_acquire),
+          modified, on_command_line, counter_);
+    }
+    case FlagValueStorageKind::kTwoWordsAtomic: {
+      return absl::make_unique<FlagState>(
+          this, value_.two_words_atomic.load(std::memory_order_acquire),
+          modified, on_command_line, counter_);
+    }
   }
+  return nullptr;
+}
 
-  Write(value);
+bool FlagImpl::RestoreState(const FlagState& flag_state) {
+  absl::MutexLock l(DataGuard());
 
-  {
-    absl::MutexLock l(DataGuard());
+  if (flag_state.counter_ == counter_) {
+    return false;
+  }
 
-    modified_ = modified;
-    on_command_line_ = on_command_line;
+  switch (ValueStorageKind()) {
+    case FlagValueStorageKind::kHeapAllocated:
+      StoreValue(flag_state.value_.dynamic);
+      break;
+    case FlagValueStorageKind::kOneWordAtomic:
+      StoreValue(&flag_state.value_.one_word);
+      break;
+    case FlagValueStorageKind::kTwoWordsAtomic:
+      StoreValue(&flag_state.value_.two_words);
+      break;
   }
 
+  modified_ = flag_state.modified_;
+  on_command_line_ = flag_state.on_command_line_;
+
   return true;
 }
 
@@ -290,10 +379,10 @@ std::unique_ptr<void, DynValueDeleter> FlagImpl::TryParse(
 }
 
 void FlagImpl::Read(void* dst) const {
-  DataGuard();  // Make sure flag initialized
+  auto* guard = DataGuard();  // Make sure flag initialized
   switch (ValueStorageKind()) {
     case FlagValueStorageKind::kHeapAllocated: {
-      absl::MutexLock l(DataGuard());
+      absl::MutexLock l(guard);
 
       flags_internal::CopyConstruct(op_, value_.dynamic, dst);
       break;
@@ -301,13 +390,13 @@ void FlagImpl::Read(void* dst) const {
     case FlagValueStorageKind::kOneWordAtomic: {
       const auto one_word_val =
           value_.one_word_atomic.load(std::memory_order_acquire);
-      std::memcpy(dst, &one_word_val, Sizeof(op_));
+      std::memcpy(dst, &one_word_val, flags_internal::Sizeof(op_));
       break;
     }
     case FlagValueStorageKind::kTwoWordsAtomic: {
       const auto two_words_val =
           value_.two_words_atomic.load(std::memory_order_acquire);
-      std::memcpy(dst, &two_words_val, Sizeof(op_));
+      std::memcpy(dst, &two_words_val, flags_internal::Sizeof(op_));
       break;
     }
   }
diff --git a/absl/flags/internal/flag.h b/absl/flags/internal/flag.h
index 2c4aba68..f27e558b 100644
--- a/absl/flags/internal/flag.h
+++ b/absl/flags/internal/flag.h
@@ -40,9 +40,6 @@ namespace absl {
 ABSL_NAMESPACE_BEGIN
 namespace flags_internal {
 
-template <typename T>
-class Flag;
-
 ///////////////////////////////////////////////////////////////////////////////
 // Flag value type operations, eg., parsing, copying, etc. are provided
 // by function specific to that type with a signature matching FlagOpFn.
@@ -149,36 +146,6 @@ inline FlagStaticTypeId StaticTypeId(FlagOpFn op) {
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-// Persistent state of the flag data.
-
-template <typename T>
-class FlagState : public flags_internal::FlagStateInterface {
- public:
-  FlagState(Flag<T>* flag, T&& cur, bool modified, bool on_command_line,
-            int64_t counter)
-      : flag_(flag),
-        cur_value_(std::move(cur)),
-        modified_(modified),
-        on_command_line_(on_command_line),
-        counter_(counter) {}
-
-  ~FlagState() override = default;
-
- private:
-  friend class Flag<T>;
-
-  // Restores the flag to the saved state.
-  void Restore() const override;
-
-  // Flag and saved flag data.
-  Flag<T>* flag_;
-  T cur_value_;
-  bool modified_;
-  bool on_command_line_;
-  int64_t counter_;
-};
-
-///////////////////////////////////////////////////////////////////////////////
 // Flag help auxiliary structs.
 
 // This is help argument for absl::Flag encapsulating the string literal pointer
@@ -341,13 +308,15 @@ struct FlagCallback {
 struct DynValueDeleter {
   explicit DynValueDeleter(FlagOpFn op_arg = nullptr) : op(op_arg) {}
   void operator()(void* ptr) const {
-    if (op != nullptr) Delete(op, ptr);
+    if (op != nullptr) flags_internal::Delete(op, ptr);
   }
 
   FlagOpFn op;
 };
 
-class FlagImpl {
+class FlagState;
+
+class FlagImpl final : public flags_internal::CommandLineFlag {
  public:
   constexpr FlagImpl(const char* name, const char* filename, FlagOpFn op,
                      FlagHelpArg help, FlagValueStorageKind value_kind,
@@ -368,15 +337,7 @@ class FlagImpl {
         data_guard_{} {}
 
   // Constant access methods
-  absl::string_view Name() const;
-  std::string Filename() const;
-  std::string Help() const;
-  bool IsModified() const ABSL_LOCKS_EXCLUDED(*DataGuard());
-  bool IsSpecifiedOnCommandLine() const ABSL_LOCKS_EXCLUDED(*DataGuard());
-  std::string DefaultValue() const ABSL_LOCKS_EXCLUDED(*DataGuard());
-  std::string CurrentValue() const ABSL_LOCKS_EXCLUDED(*DataGuard());
-  void Read(void* dst) const ABSL_LOCKS_EXCLUDED(*DataGuard());
-
+  void Read(void* dst) const override ABSL_LOCKS_EXCLUDED(*DataGuard());
   template <typename T, typename std::enable_if<FlagUseHeapStorage<T>::value,
                                                 int>::type = 0>
   void Get(T* dst) const {
@@ -404,34 +365,12 @@ class FlagImpl {
 
   // Mutating access methods
   void Write(const void* src) ABSL_LOCKS_EXCLUDED(*DataGuard());
-  bool ParseFrom(absl::string_view value, FlagSettingMode set_mode,
-                 ValueSource source, std::string* err)
-      ABSL_LOCKS_EXCLUDED(*DataGuard());
 
   // Interfaces to operate on callbacks.
   void SetCallback(const FlagCallbackFunc mutation_callback)
       ABSL_LOCKS_EXCLUDED(*DataGuard());
   void InvokeCallback() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(*DataGuard());
 
-  // Interfaces to save/restore mutable flag data
-  template <typename T>
-  std::unique_ptr<FlagStateInterface> SaveState(Flag<T>* flag) const
-      ABSL_LOCKS_EXCLUDED(*DataGuard()) {
-    T&& cur_value = flag->Get();
-    absl::MutexLock l(DataGuard());
-
-    return absl::make_unique<FlagState<T>>(
-        flag, std::move(cur_value), modified_, on_command_line_, counter_);
-  }
-  bool RestoreState(const void* value, bool modified, bool on_command_line,
-                    int64_t counter) ABSL_LOCKS_EXCLUDED(*DataGuard());
-
-  // Value validation interfaces.
-  void CheckDefaultValueParsingRoundtrip() const
-      ABSL_LOCKS_EXCLUDED(*DataGuard());
-  bool ValidateInputValue(absl::string_view value) const
-      ABSL_LOCKS_EXCLUDED(*DataGuard());
-
   // Used in read/write operations to validate source/target has correct type.
   // For example if flag is declared as absl::Flag<int> FLAGS_foo, a call to
   // absl::GetFlag(FLAGS_foo) validates that the type of FLAGS_foo is indeed
@@ -441,6 +380,10 @@ class FlagImpl {
   void AssertValidType(FlagStaticTypeId type_id) const;
 
  private:
+  template <typename T>
+  friend class Flag;
+  friend class FlagState;
+
   // Ensures that `data_guard_` is initialized and returns it.
   absl::Mutex* DataGuard() const ABSL_LOCK_RETURNED((absl::Mutex*)&data_guard_);
   // Returns heap allocated value of type T initialized with default value.
@@ -467,6 +410,37 @@ class FlagImpl {
     return static_cast<FlagDefaultKind>(def_kind_);
   }
 
+  // CommandLineFlag interface implementation
+  absl::string_view Name() const override;
+  std::string Filename() const override;
+  absl::string_view Typename() const override;
+  std::string Help() const override;
+  FlagStaticTypeId TypeId() const override;
+  bool IsModified() const override ABSL_LOCKS_EXCLUDED(*DataGuard());
+  bool IsSpecifiedOnCommandLine() const override
+      ABSL_LOCKS_EXCLUDED(*DataGuard());
+  std::string DefaultValue() const override ABSL_LOCKS_EXCLUDED(*DataGuard());
+  std::string CurrentValue() const override ABSL_LOCKS_EXCLUDED(*DataGuard());
+  bool ValidateInputValue(absl::string_view value) const override
+      ABSL_LOCKS_EXCLUDED(*DataGuard());
+  void CheckDefaultValueParsingRoundtrip() const override
+      ABSL_LOCKS_EXCLUDED(*DataGuard());
+
+  // Interfaces to save and restore flags to/from persistent state.
+  // Returns current flag state or nullptr if flag does not support
+  // saving and restoring a state.
+  std::unique_ptr<FlagStateInterface> SaveState() override
+      ABSL_LOCKS_EXCLUDED(*DataGuard());
+
+  // Restores the flag state to the supplied state object. If there is
+  // nothing to restore returns false. Otherwise returns true.
+  bool RestoreState(const FlagState& flag_state)
+      ABSL_LOCKS_EXCLUDED(*DataGuard());
+
+  bool ParseFrom(absl::string_view value, FlagSettingMode set_mode,
+                 ValueSource source, std::string* error) override
+      ABSL_LOCKS_EXCLUDED(*DataGuard());
+
   // Immutable flag's state.
 
   // Flags name passed to ABSL_FLAG as second arg.
@@ -536,7 +510,7 @@ class FlagImpl {
 // flag reflection handle interface.
 
 template <typename T>
-class Flag final : public flags_internal::CommandLineFlag {
+class Flag {
  public:
   constexpr Flag(const char* name, const char* filename, const FlagHelpArg help,
                  const FlagDfltGenFunc default_value_gen)
@@ -568,60 +542,24 @@ class Flag final : public flags_internal::CommandLineFlag {
   }
 
   // CommandLineFlag interface
-  absl::string_view Name() const override { return impl_.Name(); }
-  std::string Filename() const override { return impl_.Filename(); }
-  absl::string_view Typename() const override { return ""; }
-  std::string Help() const override { return impl_.Help(); }
-  bool IsModified() const override { return impl_.IsModified(); }
-  bool IsSpecifiedOnCommandLine() const override {
+  absl::string_view Name() const { return impl_.Name(); }
+  std::string Filename() const { return impl_.Filename(); }
+  absl::string_view Typename() const { return ""; }
+  std::string Help() const { return impl_.Help(); }
+  bool IsModified() const { return impl_.IsModified(); }
+  bool IsSpecifiedOnCommandLine() const {
     return impl_.IsSpecifiedOnCommandLine();
   }
-  std::string DefaultValue() const override { return impl_.DefaultValue(); }
-  std::string CurrentValue() const override { return impl_.CurrentValue(); }
-  bool ValidateInputValue(absl::string_view value) const override {
-    return impl_.ValidateInputValue(value);
-  }
-
-  // Interfaces to save and restore flags to/from persistent state.
-  // Returns current flag state or nullptr if flag does not support
-  // saving and restoring a state.
-  std::unique_ptr<FlagStateInterface> SaveState() override {
-    return impl_.SaveState(this);
-  }
-
-  // Restores the flag state to the supplied state object. If there is
-  // nothing to restore returns false. Otherwise returns true.
-  bool RestoreState(const FlagState<T>& flag_state) {
-    return impl_.RestoreState(&flag_state.cur_value_, flag_state.modified_,
-                              flag_state.on_command_line_, flag_state.counter_);
-  }
-  bool ParseFrom(absl::string_view value, FlagSettingMode set_mode,
-                 ValueSource source, std::string* error) override {
-    return impl_.ParseFrom(value, set_mode, source, error);
-  }
-  void CheckDefaultValueParsingRoundtrip() const override {
-    impl_.CheckDefaultValueParsingRoundtrip();
-  }
+  std::string DefaultValue() const { return impl_.DefaultValue(); }
+  std::string CurrentValue() const { return impl_.CurrentValue(); }
 
  private:
-  friend class FlagState<T>;
-
-  void Read(void* dst) const override { impl_.Read(dst); }
-  FlagStaticTypeId TypeId() const override { return &FlagStaticTypeIdGen<T>; }
-
+  template <typename U, bool do_register>
+  friend class FlagRegistrar;
   // Flag's data
   FlagImpl impl_;
 };
 
-template <typename T>
-void FlagState<T>::Restore() const {
-  if (flag_->RestoreState(*this)) {
-    ABSL_INTERNAL_LOG(INFO,
-                      absl::StrCat("Restore saved value of ", flag_->Name(),
-                                   " to: ", flag_->CurrentValue()));
-  }
-}
-
 // This class facilitates Flag object registration and tail expression-based
 // flag definition, for example:
 // ABSL_FLAG(int, foo, 42, "Foo help").OnUpdate(NotifyFooWatcher);
@@ -629,7 +567,7 @@ template <typename T, bool do_register>
 class FlagRegistrar {
  public:
   explicit FlagRegistrar(Flag<T>* flag) : flag_(flag) {
-    if (do_register) flags_internal::RegisterCommandLineFlag(flag_);
+    if (do_register) flags_internal::RegisterCommandLineFlag(&flag_->impl_);
   }
 
   FlagRegistrar& OnUpdate(FlagCallbackFunc cb) && {