Add Arena enabled constructors to ArenaStringPtr

This optimizes non const and copy initialization of arena string values. String hardening is automatically performed on as needed by the ArenaStringPtr class, removing the need for extra clutter or set operations added by the protocol compiler in generated code.

PiperOrigin-RevId: 549480061

src/google/protobuf/BUILD.bazel

@@ -973,6 +973,7 @@ cc_test(
     name = "arenastring_unittest",
     srcs = ["arenastring_unittest.cc"],
     deps = [
+        ":port_def",
         ":protobuf",
         "//src/google/protobuf/io",
         "//src/google/protobuf/stubs",

src/google/protobuf/arenastring.cc

@@ -115,6 +115,11 @@ TaggedStringPtr CreateArenaString(Arena& arena, absl::string_view s) {
 
 }  // namespace
 
+TaggedStringPtr TaggedStringPtr::ForceCopy(Arena* arena) const {
+  return arena != nullptr ? CreateArenaString(*arena, *Get())
+                          : CreateString(*Get());
+}
+
 void ArenaStringPtr::Set(absl::string_view value, Arena* arena) {
   ScopedCheckPtrInvariants check(&tagged_ptr_);
   if (IsDefault()) {

src/google/protobuf/arenastring.h

@@ -32,15 +32,16 @@
 #define GOOGLE_PROTOBUF_ARENASTRING_H__
 
 #include <algorithm>
+#include <cstdint>
 #include <string>
 #include <type_traits>
 #include <utility>
 
 #include "absl/log/absl_check.h"
 #include "google/protobuf/arena.h"
-#include "google/protobuf/port.h"
 #include "absl/strings/string_view.h"
 #include "google/protobuf/explicitly_constructed.h"
+#include "google/protobuf/port.h"
 
 // must be last:
 #include "google/protobuf/port_def.inc"
@@ -94,7 +95,7 @@ class PROTOBUF_EXPORT LazyString {
   const std::string& Init() const;
 };
 
-class TaggedStringPtr {
+class PROTOBUF_EXPORT TaggedStringPtr {
  public:
   // Bit flags qualifying string properties. We can use 2 bits as
   // ptr_ is guaranteed and enforced to be aligned on 4 byte boundaries.
@@ -193,13 +194,25 @@ class TaggedStringPtr {
   // Returns true if the contained pointer is null, indicating some error.
   // The Null value is only used during parsing for temporary values.
   // A persisted ArenaStringPtr value is never null.
-  inline bool IsNull() { return ptr_ == nullptr; }
+  inline bool IsNull() const { return ptr_ == nullptr; }
+
+  // Returns a copy of this instance. In debug builds, the returned value may be
+  // a forced copy regardless if the current instance is a compile time default.
+  TaggedStringPtr Copy(Arena* arena) const;
+
+  // Identical to the above `Copy` function except that in debug builds,
+  // `default_value` can be used to substitute an empty default with a
+  // hardened copy of the default value.
+  TaggedStringPtr Copy(Arena* arena, const LazyString& default_value) const;
 
  private:
   static inline void assert_aligned(const void* p) {
     ABSL_DCHECK_EQ(reinterpret_cast<uintptr_t>(p) & kMask, 0UL);
   }
 
+  // Creates a heap or arena allocated copy of this instance.
+  TaggedStringPtr ForceCopy(Arena* arena) const;
+
   inline std::string* TagAs(Type type, std::string* p) {
     ABSL_DCHECK(p != nullptr);
     assert_aligned(p);
@@ -234,11 +247,54 @@ static_assert(std::is_trivial<TaggedStringPtr>::value,
 // See TaggedStringPtr for more information about the types of string values
 // being held, and the mutable and ownership invariants for each type.
 struct PROTOBUF_EXPORT ArenaStringPtr {
+  // Default constructor, leaves current instance uninitialized (does nothing)
   ArenaStringPtr() = default;
+
+  // Constexpr constructor, initializes to a constexpr, empty string value.
   constexpr ArenaStringPtr(ExplicitlyConstructedArenaString* default_value,
                            ConstantInitialized)
       : tagged_ptr_(default_value) {}
 
+  // Arena enabled constructor for strings without a default value.
+  // Initializes this instance to a constexpr, empty string value, unless debug
+  // hardening is enabled, in which case this instance will hold a forced copy.
+  explicit ArenaStringPtr(Arena* arena)
+      : tagged_ptr_(&fixed_address_empty_string) {
+    if (DebugHardenStringValues()) {
+      Set(absl::string_view(""), arena);
+    }
+  }
+
+  // Arena enabled constructor for strings with a non-empty default value.
+  // Initializes this instance to a constexpr, empty string value, unless debug
+  // hardening is enabled, in which case this instance will be forced to hold a
+  // forced copy of the value in `default_value`.
+  ArenaStringPtr(Arena* arena, const LazyString& default_value)
+      : tagged_ptr_(&fixed_address_empty_string) {
+    if (DebugHardenStringValues()) {
+      Set(absl::string_view(default_value.get()), arena);
+    }
+  }
+
+  // Arena enabled copy constructor for strings without a default value.
+  // This instance will be initialized with a copy of the value in `rhs`.
+  // If `rhs` holds a default (empty) value, then this instance will also be
+  // initialized with the default empty value, unless debug hardening is
+  // enabled, in which case this instance will be forced to hold a copy of
+  // an empty default value.
+  ArenaStringPtr(Arena* arena, const ArenaStringPtr& rhs)
+      : tagged_ptr_(rhs.tagged_ptr_.Copy(arena)) {}
+
+  // Arena enabled copy constructor for strings with a non-empty default value.
+  // This instance will be initialized with a copy of the value in `rhs`.
+  // If `rhs` holds a default (empty) value, then this instance will also be
+  // initialized with the default empty value, unless debug hardening is
+  // enabled, in which case this instance will be forced to hold forced copy
+  // of the value in `default_value`.
+  ArenaStringPtr(Arena* arena, const ArenaStringPtr& rhs,
+                 const LazyString& default_value)
+      : tagged_ptr_(rhs.tagged_ptr_.Copy(arena, default_value)) {}
+
   // Called from generated code / reflection runtime only. Resets value to point
   // to a default string pointer, with the semantics that this ArenaStringPtr
   // does not own the pointed-to memory. Disregards initial value of ptr_ (so
@@ -392,6 +448,27 @@ struct PROTOBUF_EXPORT ArenaStringPtr {
   friend class EpsCopyInputStream;
 };
 
+inline TaggedStringPtr TaggedStringPtr::Copy(Arena* arena) const {
+  if (DebugHardenStringValues()) {
+    // Harden by forcing an allocated string value.
+    return IsNull() ? *this : ForceCopy(arena);
+  }
+  return IsDefault() ? *this : ForceCopy(arena);
+}
+
+inline TaggedStringPtr TaggedStringPtr::Copy(
+    Arena* arena, const LazyString& default_value) const {
+  if (DebugHardenStringValues()) {
+    // Harden by forcing an allocated string value.
+    TaggedStringPtr hardened(*this);
+    if (IsDefault()) {
+      hardened.SetDefault(&default_value.get());
+    }
+    return hardened.ForceCopy(arena);
+  }
+  return IsDefault() ? *this : ForceCopy(arena);
+}
+
 inline void ArenaStringPtr::InitDefault() {
   tagged_ptr_ = TaggedStringPtr(&fixed_address_empty_string);
 }

src/google/protobuf/arenastring_unittest.cc

@@ -42,8 +42,11 @@
 #include <gtest/gtest.h>
 #include "absl/log/absl_check.h"
 #include "absl/strings/string_view.h"
+#include "google/protobuf/explicitly_constructed.h"
 #include "google/protobuf/io/coded_stream.h"
 #include "google/protobuf/io/zero_copy_stream_impl.h"
+#include "google/protobuf/message_lite.h"
+#include "google/protobuf/port.h"
 
 
 // Must be included last.
@@ -146,6 +149,123 @@ TEST_P(DualArena, Swap) {
   rhs.Destroy();
 }
 
+TEST(ArenaStringPtrTest, ConstInit) {
+  // Verify that we can constinit construct an ArenaStringPtr from an arbitrary
+  // ExplicitlyConstructed<std::string>*.
+  static internal::ExplicitlyConstructedArenaString str;
+  PROTOBUF_CONSTINIT static ArenaStringPtr ptr(&str,
+                                               internal::ConstantInitialized{});
+  EXPECT_EQ(&ptr.Get(), str.get_mutable());
+
+  PROTOBUF_CONSTINIT static const ArenaStringPtr ptr2(
+      &internal::fixed_address_empty_string, internal::ConstantInitialized{});
+  EXPECT_EQ(&ptr2.Get(), &internal::GetEmptyStringAlreadyInited());
+}
+
+TEST_P(SingleArena, ConstructEmpty) {
+  auto arena = GetArena();
+  ArenaStringPtr field(arena.get());
+
+  EXPECT_EQ(field.Get(), "");
+  if (internal::DebugHardenStringValues()) {
+    EXPECT_FALSE(field.IsDefault());
+  } else {
+    EXPECT_TRUE(field.IsDefault());
+  }
+  if (arena == nullptr) field.Destroy();
+}
+
+TEST_P(SingleArena, ConstructEmptyWithDefault) {
+  auto arena = GetArena();
+  internal::LazyString default_value{{{"Hello default", 13}}, {nullptr}};
+  ArenaStringPtr field(arena.get(), default_value);
+
+  if (internal::DebugHardenStringValues()) {
+    EXPECT_EQ(field.Get(), "Hello default");
+    EXPECT_FALSE(field.IsDefault());
+  } else {
+    EXPECT_EQ(field.Get(), "");
+    EXPECT_TRUE(field.IsDefault());
+  }
+  if (arena == nullptr) field.Destroy();
+}
+
+TEST_P(SingleArena, CopyConstructEmpty) {
+  std::string empty;
+  auto arena = GetArena();
+  ArenaStringPtr field;
+  field.InitExternal(&empty);
+
+  ArenaStringPtr dst(arena.get(), field);
+  EXPECT_EQ(dst.Get(), "");
+  if (internal::DebugHardenStringValues()) {
+    EXPECT_FALSE(dst.IsDefault());
+  } else {
+    EXPECT_TRUE(dst.IsDefault());
+  }
+  if (arena == nullptr) dst.Destroy();
+  field.Destroy();
+}
+
+TEST_P(SingleArena, CopyConstructEmptyWithDefault) {
+  std::string empty;
+  auto arena = GetArena();
+  ArenaStringPtr field;
+  field.InitExternal(&empty);
+
+  internal::LazyString default_value{{{"Hello default", 13}}, {nullptr}};
+  ArenaStringPtr dst(arena.get(), field, default_value);
+  if (internal::DebugHardenStringValues()) {
+    EXPECT_EQ(dst.Get(), "Hello default");
+    EXPECT_FALSE(dst.IsDefault());
+  } else {
+    EXPECT_EQ(dst.Get(), "");
+    EXPECT_TRUE(dst.IsDefault());
+  }
+  if (arena == nullptr) dst.Destroy();
+  field.Destroy();
+}
+
+TEST_P(SingleArena, CopyConstructValueWithDefault) {
+  std::string empty;
+  auto arena = GetArena();
+  ArenaStringPtr field;
+  field.InitExternal(&empty);
+  field.Set("Hello world", nullptr);
+
+  internal::LazyString default_value{{{"Hello default", 13}}, {nullptr}};
+  ArenaStringPtr dst(arena.get(), field, default_value);
+  EXPECT_EQ(dst.Get(), "Hello world");
+  if (arena == nullptr) dst.Destroy();
+  field.Destroy();
+}
+
+TEST_P(SingleArena, CopyConstructSSO) {
+  std::string empty;
+  auto arena = GetArena();
+  ArenaStringPtr field;
+  field.InitExternal(&empty);
+  field.Set("Hello world", nullptr);
+
+  ArenaStringPtr dst(arena.get(), field);
+  EXPECT_EQ(dst.Get(), "Hello world");
+  if (arena == nullptr) dst.Destroy();
+  field.Destroy();
+}
+
+TEST_P(SingleArena, CopyConstructLong) {
+  std::string empty;
+  auto arena = GetArena();
+  ArenaStringPtr field;
+  field.InitExternal(&empty);
+  field.Set("A string long enough to not be inlined", nullptr);
+
+  ArenaStringPtr dst(arena.get(), field);
+  EXPECT_EQ(dst.Get(), "A string long enough to not be inlined");
+  if (arena == nullptr) dst.Destroy();
+  field.Destroy();
+}
+
 
 }  // namespace protobuf
 }  // namespace google

src/google/protobuf/port.h

@@ -215,6 +215,15 @@ enum { kCacheAlignment = alignof(max_align_t) };  // do the best we can
 // The maximum byte alignment we support.
 enum { kMaxMessageAlignment = 8 };
 
+// Returns true if debug string hardening is required
+inline constexpr bool DebugHardenStringValues() {
+#ifdef PROTOBUF_FORCE_COPY_DEFAULT_STRING
+  return true;
+#else
+  return false;
+#endif
+}
+
 }  // namespace internal
 }  // namespace protobuf
 }  // namespace google