[libSyntax] Add a reference counted version of OwnedString

We cannot use unowned strings for token texts of incrementally parsed
syntax trees since the source buffer to which reused nodes refer will
have been freed for reused nodes. Always copying the token text whenever
OwnedString is passed is too expensive. A reference counted copy of the
string allows us to keep the token's string alive across incremental
parses while eliminating unnecessary copies.

include/swift/Basic/OwnedString.h

@@ -22,117 +22,89 @@
 
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/TrailingObjects.h"
 
 using llvm::StringRef;
 
 namespace swift {
 
-enum class StringOwnership {
-  /// An OwnedString holds a weak reference to the underlying string storage
-  /// and will never attempt to free it.
-  Unowned,
-
-  /// An OwnedString has its own copy of the underlying string storage and
-  /// will free the storage upon its destruction.
-  Copied,
-};
-
 /// Holds a string - either statically allocated or dynamically allocated
 /// and owned by this type.
 class OwnedString {
-  const char *Data;
-  size_t Length;
-  StringOwnership Ownership = StringOwnership::Unowned;
-
-  void release() {
-    if (Ownership == StringOwnership::Copied)
-      free(const_cast<char *>(Data));
-  }
-
-  void initialize(const char* Data, size_t Length, StringOwnership Ownership) {
-    this->Length = Length;
-    this->Ownership = Ownership;
-    if (Ownership == StringOwnership::Copied && Data) {
-      char *substring = static_cast<char *>(malloc(Length + 1));
-      assert(substring && "expected successful malloc of copy");
-
-      memcpy(substring, Data, Length);
-      substring[Length] = '\0';
-
-      this->Data = substring;
+  /// An owner that keeps the buffer of a ref counted \c OwnedString alive.
+  class TextOwner final : public llvm::ThreadSafeRefCountedBase<TextOwner>,
+                          public llvm::TrailingObjects<TextOwner, char> {
+    TextOwner(StringRef Text) {
+      std::uninitialized_copy(Text.begin(), Text.end(),
+                              getTrailingObjects<char>());
     }
-    else
-      this->Data = Data;
-  }
-  OwnedString(const char* Data, size_t Length, StringOwnership Ownership) {
-    initialize(Data, Length, Ownership);
-  }
+
+  public:
+    static TextOwner *make(StringRef Text) {
+      auto size = totalSizeToAlloc<char>(Text.size());
+      void *data = ::operator new(size);
+      return new (data) TextOwner(Text);
+    }
+
+    const char *getText() const { return getTrailingObjects<char>(); }
+  };
+
+  /// The text this owned string represents
+  StringRef Text;
+
+  /// In case of a ref counted string an owner that keeps the buffer \c Text
+  /// references alive.
+  llvm::IntrusiveRefCntPtr<TextOwner> OwnedPtr;
+
+  OwnedString(StringRef Text, llvm::IntrusiveRefCntPtr<TextOwner> OwnedPtr)
+      : Text(Text), OwnedPtr(OwnedPtr) {}
+
 public:
-  OwnedString(): OwnedString(nullptr, 0, StringOwnership::Unowned) {}
+  OwnedString() : OwnedString(/*Text=*/StringRef(), /*OwnedPtr=*/nullptr) {}
 
-  OwnedString(const char *Data, size_t Length):
-    OwnedString(Data, Length, StringOwnership::Copied) {}
+  /// Create a ref counted \c OwnedString that is initialized with the text of
+  /// the given \c StringRef.
+  OwnedString(StringRef Str) : OwnedString(makeRefCounted(Str)) {}
 
-  OwnedString(StringRef Str) : OwnedString(Str.data(), Str.size()) {}
+  /// Create a ref counted \c OwnedString that is initialized with the text of
+  /// the given buffer.
+  OwnedString(const char *Str) : OwnedString(StringRef(Str)) {}
 
-  OwnedString(const char *Data) : OwnedString(StringRef(Data)) {}
-
-  OwnedString(const OwnedString &Other):
-    OwnedString(Other.Data, Other.Length, Other.Ownership) {}
-
-  OwnedString(OwnedString &&Other): Data(Other.Data), Length(Other.Length),
-      Ownership(Other.Ownership) {
-    Other.Data = nullptr;
-    Other.Ownership = StringOwnership::Unowned;
+  /// Create an \c OwnedString that references the given string. The
+  /// \c OwnedString will not take ownership of that buffer and will assume that
+  /// the buffer outlives its lifetime.
+  static OwnedString makeUnowned(StringRef Str) {
+    return OwnedString(Str, /*OwnedPtr=*/nullptr);
   }
 
-  OwnedString& operator=(const OwnedString &Other) {
-    if (&Other != this) {
-      release();
-      initialize(Other.Data, Other.Length, Other.Ownership);
+  /// Create an \c OwnedString that keeps its contents in a reference counted
+  /// buffer. The contents of \p Str will be copied initially and are allowed to
+  /// be disposed after the \c OwnedString has been created.
+  static OwnedString makeRefCounted(StringRef Str) {
+    if (Str.empty()) {
+      // Copying an empty string doesn't make sense. Just create an unowned
+      // string that points to the empty string.
+      return makeUnowned(Str);
+    } else {
+      llvm::IntrusiveRefCntPtr<TextOwner> OwnedPtr(TextOwner::make(Str));
+      return OwnedString(StringRef(OwnedPtr->getText(), Str.size()),
+                         std::move(OwnedPtr));
     }
-    return *this;
-  }
-
-  OwnedString& operator=(OwnedString &&Other) {
-    if (&Other != this) {
-      release();
-      this->Data = Other.Data;
-      this->Length = Other.Length;
-      this->Ownership = Other.Ownership;
-      Other.Ownership = StringOwnership::Unowned;
-      Other.Data = nullptr;
-    }
-    return *this;
-  }
-
-  OwnedString copy() const {
-    return OwnedString(Data, Length, StringOwnership::Copied);
   }
 
   /// Returns the length of the string in bytes.
-  size_t size() const {
-    return Length;
-  }
+  size_t size() const { return Text.size(); }
 
   /// Returns true if the length is 0.
-  bool empty() const {
-    return Length == 0;
-  }
+  bool empty() const { return size() == 0; }
 
   /// Returns a StringRef to the underlying data. No copy is made and no
   /// ownership changes take place.
-  StringRef str() const {
-    return StringRef { Data, Length };
-  }
+  StringRef str() const { return Text; }
 
   bool operator==(const OwnedString &Right) const {
     return str() == Right.str();
   }
-
-  ~OwnedString() {
-    release();
-  }
 };
 
 } // end namespace swift