[libSyntax] Always copy leading and trailing trivia strings into a SyntaxArena buffer

Referencing a string in arbitrary memory is not safe since the source
buffer to which it points may have been freed. Instead copy all strings
into the SyntaxArena. Since RawSyntax nodes retain their arena, they can
be sure that the string won't disappear if it lives in their arena.

To avoid lots of small copies, we copy the entire source buffer once
into the syntax arena and make StringRefs point into that buffer.

Author: ahoppen

include/swift/Syntax/SyntaxArena.h

@@ -40,6 +40,10 @@ class SyntaxArena : public llvm::ThreadSafeRefCountedBase<SyntaxArena> {
   void *Allocate(size_t size, size_t alignment) {
     return Allocator.Allocate(size, alignment);
   }
+
+  bool containsPointer(const void *Ptr) {
+    return getAllocator().identifyObject(Ptr) != llvm::None;
+  }
 };
 
 } // namespace syntax

include/swift/SyntaxParse/SyntaxTreeCreator.h

@@ -15,6 +15,7 @@
 
 #include "swift/Parse/SyntaxParseActions.h"
 #include "swift/Syntax/References.h"
+#include "llvm/ADT/StringRef.h"
 
 namespace swift {
   class RawSyntaxTokenCache;
@@ -38,6 +39,12 @@ class SyntaxTreeCreator: public SyntaxParseActions {
   unsigned BufferID;
   RC<syntax::SyntaxArena> Arena;
 
+  /// A string allocated in \c Arena that contains an exact copy of the source
+  /// file for which this \c SyntaxTreeCreator creates a syntax tree. \c
+  /// RawSyntax nodes can safely reference text inside this buffer since they
+  /// retain the \c SyntaxArena which holds the buffer.
+  StringRef ArenaSourceBuffer;
+
   /// A cache of nodes that can be reused when creating the current syntax
   /// tree.
   SyntaxParsingCache *SyntaxCache;

lib/Syntax/RawSyntax.cpp

@@ -87,6 +87,23 @@ Trivia lexTrivia(StringRef TriviaStr) {
   return SyntaxTrivia;
 }
 
+/// If the \p Str is not allocated in \p Arena, copy it to \p Arena and adjust
+/// \p Str to point to the string's copy in \p Arena.
+void copyToArenaIfNecessary(StringRef &Str, const RC<SyntaxArena> Arena) {
+  if (Str.empty()) {
+    // Empty strings can live wherever they want. Nothing to do.
+    return;
+  }
+  if (Arena->containsPointer(Str.data())) {
+    // String already in arena. Nothing to do.
+    return;
+  }
+  // Copy string to arena
+  char *Data = (char *)Arena->Allocate(Str.size(), alignof(char *));
+  std::uninitialized_copy(Str.begin(), Str.end(), Data);
+  Str = StringRef(Data, Str.size());
+}
+
 // FIXME: If we want thread-safety for tree creation, this needs to be atomic.
 unsigned RawSyntax::NextFreeNodeId = 1;
 
@@ -129,6 +146,8 @@ RawSyntax::RawSyntax(tok TokKind, OwnedString Text, size_t TextLength,
     : Arena(Arena), Bits({{unsigned(TextLength), unsigned(Presence), true}}),
       RefCount(0) {
   assert(Arena && "RawSyntax nodes must always be allocated in an arena");
+  copyToArenaIfNecessary(LeadingTrivia, Arena);
+  copyToArenaIfNecessary(TrailingTrivia, Arena);
 
   if (NodeId.hasValue()) {
     this->NodeId = NodeId.getValue();

lib/SyntaxParse/SyntaxTreeCreator.cpp

@@ -42,6 +42,11 @@ SyntaxTreeCreator::SyntaxTreeCreator(SourceManager &SM, unsigned bufferID,
       Arena(std::move(arena)),
       SyntaxCache(syntaxCache),
       TokenCache(new RawSyntaxTokenCache()) {
+  StringRef BufferContent = SM.getEntireTextForBuffer(BufferID);
+  char *Data = (char *)Arena->Allocate(BufferContent.size(), alignof(char *));
+  std::uninitialized_copy(BufferContent.begin(), BufferContent.end(), Data);
+  ArenaSourceBuffer = StringRef(Data, BufferContent.size());
+  assert(ArenaSourceBuffer == BufferContent);
 }
 
 SyntaxTreeCreator::~SyntaxTreeCreator() = default;
@@ -113,14 +118,25 @@ OpaqueSyntaxNode SyntaxTreeCreator::recordToken(tok tokenKind,
                                                 StringRef leadingTrivia,
                                                 StringRef trailingTrivia,
                                                 CharSourceRange range) {
-  SourceLoc tokLoc = range.getStart().getAdvancedLoc(leadingTrivia.size());
   unsigned tokLength =
       range.getByteLength() - leadingTrivia.size() - trailingTrivia.size();
-  CharSourceRange tokRange = CharSourceRange{tokLoc, tokLength};
-  StringRef tokenText = SM.extractText(tokRange, BufferID);
+  auto leadingTriviaStartOffset =
+      SM.getLocOffsetInBuffer(range.getStart(), BufferID);
+  auto tokStartOffset = leadingTriviaStartOffset + leadingTrivia.size();
+  auto trailingTriviaStartOffset = tokStartOffset + tokLength;
+
+  // Get StringRefs of the token's texts that point into the syntax arena's
+  // buffer.
+  StringRef leadingTriviaText =
+      ArenaSourceBuffer.substr(leadingTriviaStartOffset, leadingTrivia.size());
+  StringRef tokenText = ArenaSourceBuffer.substr(tokStartOffset, tokLength);
+  StringRef trailingTriviaText = ArenaSourceBuffer.substr(
+      trailingTriviaStartOffset, trailingTrivia.size());
+
   auto ownedText = OwnedString::makeRefCounted(tokenText);
-  auto raw = TokenCache->getToken(Arena, tokenKind, range.getByteLength(),
-                                  ownedText, leadingTrivia, trailingTrivia);
+  auto raw =
+      TokenCache->getToken(Arena, tokenKind, range.getByteLength(), ownedText,
+                           leadingTriviaText, trailingTriviaText);
   OpaqueSyntaxNode opaqueN = raw.get();
   raw.resetWithoutRelease();
   return opaqueN;