[libSyntax] Make RawSyntax nodes hold a strong reference to their arena

This allows an elegant design in which we can still allocate RawSyntax
nodes using a bump allocator but are able to automatically free that
buffer once the last RawSyntax node within that buffer is freed.

This also resolves a memory leak of RawSyntax nodes that was caused by
ParserUnit not freeing its underlying ASTContext.

Author: ahoppen

include/swift/AST/ASTContext.h

@@ -392,7 +392,7 @@ class ASTContext final {
   }
 
   /// Retrive the syntax node memory manager for this context.
-  syntax::SyntaxArena &getSyntaxArena() const;
+  llvm::IntrusiveRefCntPtr<syntax::SyntaxArena> getSyntaxArena() const;
 
   /// Retrieve the lazy resolver for this context.
   LazyResolver *getLazyResolver() const;

include/swift/Parse/SyntaxParsingContext.h

@@ -18,6 +18,74 @@
 #include "swift/Syntax/Syntax.h"
 #include "swift/Syntax/TokenSyntax.h"
 
+namespace swift {
+
+using namespace swift::syntax;
+
+/// Cache node for RawSyntax.
+class RawSyntaxCacheNode : public llvm::FoldingSetNode {
+
+  friend llvm::FoldingSetTrait<RawSyntaxCacheNode>;
+
+  /// Associated RawSyntax.
+  RC<RawSyntax> Obj;
+  /// FoldingSet node identifier of the associated RawSyntax.
+  llvm::FoldingSetNodeIDRef IDRef;
+
+public:
+  RawSyntaxCacheNode(RC<RawSyntax> Obj, const llvm::FoldingSetNodeIDRef IDRef)
+      : Obj(Obj), IDRef(IDRef) {}
+
+  /// Retrieve assciated RawSyntax.
+  RC<RawSyntax> get() { return Obj; }
+
+  // Only allow allocation of Node using the allocator in SyntaxArena.
+  void *operator new(size_t Bytes, RC<SyntaxArena> Arena,
+                     unsigned Alignment = alignof(RawSyntaxCacheNode)) {
+    return Arena->Allocate(Bytes, Alignment);
+  }
+
+  void *operator new(size_t Bytes) throw() = delete;
+  void operator delete(void *Data) throw() = delete;
+};
+
+class RawSyntaxTokenCache {
+  llvm::FoldingSet<RawSyntaxCacheNode> CachedTokens;
+  std::vector<RawSyntaxCacheNode *> CacheNodes;
+
+public:
+  RC<RawSyntax> getToken(RC<SyntaxArena> Arena, tok TokKind, OwnedString Text,
+                         llvm::ArrayRef<TriviaPiece> LeadingTrivia,
+                         llvm::ArrayRef<TriviaPiece> TrailingTrivia);
+
+  ~RawSyntaxTokenCache();
+};
+
+} // namespace swift
+
+namespace llvm {
+
+using swift::RawSyntaxCacheNode;
+
+/// FoldingSet traits for RawSyntax wrapper.
+template <> struct FoldingSetTrait<RawSyntaxCacheNode> {
+
+  static inline void Profile(RawSyntaxCacheNode &X, FoldingSetNodeID &ID) {
+    ID.AddNodeID(X.IDRef);
+  }
+
+  static inline bool Equals(RawSyntaxCacheNode &X, const FoldingSetNodeID &ID,
+                            unsigned, FoldingSetNodeID &) {
+    return ID == X.IDRef;
+  }
+  static inline unsigned ComputeHash(RawSyntaxCacheNode &X,
+                                     FoldingSetNodeID &) {
+    return X.IDRef.ComputeHash();
+  }
+};
+
+} // namespace llvm
+
 namespace swift {
 class SourceFile;
 class SyntaxParsingCache;
@@ -75,15 +143,19 @@ class alignas(1 << SyntaxAlignInBits) SyntaxParsingContext {
     // Storage for Collected parts.
     std::vector<RC<RawSyntax>> Storage;
 
-    SyntaxArena &Arena;
+    RC<SyntaxArena> Arena;
 
     /// A cache of nodes that can be reused when creating the current syntax
     /// tree
     SyntaxParsingCache *SyntaxCache = nullptr;
 
+    /// Tokens nodes that have already been created and may be reused in other
+    /// parts of the syntax tree.
+    RawSyntaxTokenCache TokenCache;
+
     RootContextData(SourceFile &SF, DiagnosticEngine &Diags,
                     SourceManager &SourceMgr, unsigned BufferID,
-                    SyntaxArena &Arena, SyntaxParsingCache *SyntaxCache)
+                    RC<SyntaxArena> Arena, SyntaxParsingCache *SyntaxCache)
         : SF(SF), Diags(Diags), SourceMgr(SourceMgr), BufferID(BufferID),
           Arena(Arena), SyntaxCache(SyntaxCache) {}
   };
@@ -217,7 +289,9 @@ class alignas(1 << SyntaxAlignInBits) SyntaxParsingContext {
     return getRootData()->SyntaxCache;
   }
 
-  SyntaxArena &getArena() const { return getRootData()->Arena; }
+  RawSyntaxTokenCache &getTokenCache() { return getRootData()->TokenCache; }
+
+  RC<SyntaxArena> getArena() const { return getRootData()->Arena; }
 
   const SyntaxParsingContext *getRoot() const;
 

include/swift/Syntax/RawSyntax.h

@@ -31,6 +31,7 @@
 
 #include "swift/Basic/InlineBitfield.h"
 #include "swift/Syntax/References.h"
+#include "swift/Syntax/SyntaxArena.h"
 #include "swift/Syntax/SyntaxKind.h"
 #include "swift/Syntax/TokenKinds.h"
 #include "swift/Syntax/Trivia.h"
@@ -218,8 +219,7 @@ typedef unsigned SyntaxNodeId;
 ///
 /// This is implementation detail - do not expose it in public API.
 class RawSyntax final
-    : public llvm::ThreadSafeRefCountedBase<RawSyntax>,
-      private llvm::TrailingObjects<RawSyntax, RC<RawSyntax>, OwnedString,
+    : private llvm::TrailingObjects<RawSyntax, RC<RawSyntax>, OwnedString,
                                     TriviaPiece> {
   friend TrailingObjects;
 
@@ -230,18 +230,20 @@ class RawSyntax final
   /// An ID of this node that is stable across incremental parses
   SyntaxNodeId NodeId;
 
+  /// If this node was allocated using a \c SyntaxArena's bump allocator, a
+  /// reference to the arena to keep the underlying memory buffer of this node
+  /// alive. If this is a \c nullptr, the node owns its own memory buffer.
+  RC<SyntaxArena> Arena;
+
   union {
     uint64_t OpaqueBits;
     struct {
       /// The kind of syntax this node represents.
       unsigned Kind : bitmax(NumSyntaxKindBits, 8);
       /// Whether this piece of syntax was actually present in the source.
       unsigned Presence : 1;
-      /// Whether this piece of syntax was constructed with manually managed
-      /// memory.
-      unsigned ManualMemory : 1;
     } Common;
-    enum { NumRawSyntaxBits = bitmax(NumSyntaxKindBits, 8) + 1 + 1 };
+    enum { NumRawSyntaxBits = bitmax(NumSyntaxKindBits, 8) + 1 };
 
     // For "layout" nodes.
     struct {
@@ -282,18 +284,23 @@ class RawSyntax final
              : 0;
   }
 
-  /// Constructor for creating layout nodes
+  /// Constructor for creating layout nodes.
+  /// If the node has been allocated inside the bump allocator of a
+  /// \c SyntaxArena, that arena must be passed as \p Arena to retain the node's
+  /// underlying storage.
   /// If \p NodeId is \c None, the next free NodeId is used, if it is passed,
   /// the caller needs to assure that the node ID has not been used yet.
   RawSyntax(SyntaxKind Kind, ArrayRef<RC<RawSyntax>> Layout,
-            SourcePresence Presence, bool ManualMemory,
+            SourcePresence Presence, RC<SyntaxArena> Arena,
             llvm::Optional<SyntaxNodeId> NodeId);
   /// Constructor for creating token nodes
+  /// \c SyntaxArena, that arena must be passed as \p Arena to retain the node's
+  /// underlying storage.
   /// If \p NodeId is \c None, the next free NodeId is used, if it is passed,
   /// the caller needs to assure that the NodeId has not been used yet.
   RawSyntax(tok TokKind, OwnedString Text, ArrayRef<TriviaPiece> LeadingTrivia,
             ArrayRef<TriviaPiece> TrailingTrivia, SourcePresence Presence,
-            bool ManualMemory, llvm::Optional<SyntaxNodeId> NodeId);
+            RC<SyntaxArena> Arena, llvm::Optional<SyntaxNodeId> NodeId);
 
   /// Compute the node's text length by summing up the length of its childern
   size_t computeTextLength() {
@@ -307,18 +314,30 @@ class RawSyntax final
     return TextLength;
   }
 
+  mutable std::atomic<int> RefCount;
+
 public:
   ~RawSyntax();
 
+  // This is a copy-pased implementation of llvm::ThreadSafeRefCountedBase with
+  // the difference that we do not delete the RawSyntax node's memory if the
+  // node was allocated within a SyntaxArena and thus doesn't own its memory.
+  void Retain() const { RefCount.fetch_add(1, std::memory_order_relaxed); }
+
   void Release() const {
-    if (Bits.Common.ManualMemory)
-      return;
-    return llvm::ThreadSafeRefCountedBase<RawSyntax>::Release();
-  }
-  void Retain() const {
-    if (Bits.Common.ManualMemory)
-      return;
-    return llvm::ThreadSafeRefCountedBase<RawSyntax>::Retain();
+    int NewRefCount = RefCount.fetch_sub(1, std::memory_order_acq_rel) - 1;
+    assert(NewRefCount >= 0 && "Reference count was already zero.");
+    if (NewRefCount == 0) {
+      if (Arena) {
+        // The node was allocated inside a SyntaxArena and thus doesn't own its
+        // own memory region. Hence we cannot free it. It will be deleted once
+        // the last RawSyntax node allocated with it will release its reference
+        // to the arena.
+        this->~RawSyntax();
+      } else {
+        delete this;
+      }
+    }
   }
 
   /// \name Factory methods.
@@ -327,33 +346,29 @@ class RawSyntax final
   /// Make a raw "layout" syntax node.
   static RC<RawSyntax> make(SyntaxKind Kind, ArrayRef<RC<RawSyntax>> Layout,
                             SourcePresence Presence,
-                            SyntaxArena *Arena = nullptr,
+                            RC<SyntaxArena> Arena = nullptr,
                             llvm::Optional<SyntaxNodeId> NodeId = llvm::None);
 
   /// Make a raw "token" syntax node.
   static RC<RawSyntax> make(tok TokKind, OwnedString Text,
                             ArrayRef<TriviaPiece> LeadingTrivia,
                             ArrayRef<TriviaPiece> TrailingTrivia,
                             SourcePresence Presence,
-                            SyntaxArena *Arena = nullptr,
+                            RC<SyntaxArena> Arena = nullptr,
                             llvm::Optional<SyntaxNodeId> NodeId = llvm::None);
 
   /// Make a missing raw "layout" syntax node.
-  static RC<RawSyntax> missing(SyntaxKind Kind, SyntaxArena *Arena = nullptr) {
+  static RC<RawSyntax> missing(SyntaxKind Kind,
+                               RC<SyntaxArena> Arena = nullptr) {
     return make(Kind, {}, SourcePresence::Missing, Arena);
   }
 
   /// Make a missing raw "token" syntax node.
   static RC<RawSyntax> missing(tok TokKind, OwnedString Text,
-                               SyntaxArena *Arena = nullptr) {
+                               RC<SyntaxArena> Arena = nullptr) {
     return make(TokKind, Text, {}, {}, SourcePresence::Missing, Arena);
   }
 
-  static RC<RawSyntax> getToken(SyntaxArena &Arena, tok TokKind,
-                                OwnedString Text,
-                                ArrayRef<TriviaPiece> LeadingTrivia,
-                                ArrayRef<TriviaPiece> TrailingTrivia);
-
   /// @}
 
   SourcePresence getPresence() const {

include/swift/Syntax/SyntaxArena.h

@@ -17,26 +17,26 @@
 #ifndef SWIFT_SYNTAX_SYNTAXARENA_H
 #define SWIFT_SYNTAX_SYNTAXARENA_H
 
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/Support/Allocator.h"
 
 namespace swift {
 namespace syntax {
 
 /// Memory manager for Syntax nodes.
-class SyntaxArena {
+class SyntaxArena : public llvm::ThreadSafeRefCountedBase<SyntaxArena> {
   SyntaxArena(const SyntaxArena &) = delete;
   void operator=(const SyntaxArena &) = delete;
 
-public:
-  struct Implementation;
-  Implementation &Impl;
+  llvm::BumpPtrAllocator Allocator;
 
-  SyntaxArena();
-  ~SyntaxArena();
+public:
+  SyntaxArena() {}
 
-  llvm::BumpPtrAllocator &getAllocator() const;
-  void *Allocate(size_t size, size_t alignment);
-  void *AllocateRawSyntax(size_t size, size_t alignment);
+  llvm::BumpPtrAllocator &getAllocator() { return Allocator; }
+  void *Allocate(size_t size, size_t alignment) {
+    return Allocator.Allocate(size, alignment);
+  }
 };
 
 } // namespace syntax

include/swift/Syntax/SyntaxBuilders.h.gyb

@@ -33,7 +33,7 @@ class SyntaxArena;
 %   if node.is_buildable():
 %     child_count = len(node.children)
 class ${node.name}Builder {
-  SyntaxArena *Arena = nullptr;
+  RC<SyntaxArena> Arena = nullptr;
   RC<RawSyntax> Layout[${child_count}] = {
 %     for child in node.children:
     nullptr,
@@ -42,7 +42,7 @@ class ${node.name}Builder {
 
 public:
   ${node.name}Builder() = default;
-  ${node.name}Builder(SyntaxArena &Arena) : Arena(&Arena) {}
+  ${node.name}Builder(RC<SyntaxArena> Arena) : Arena(Arena) {}
 
 %     for child in node.children:
   ${node.name}Builder &use${child.name}(${child.type_name} ${child.name});