[Syntax] Use TrailingObjects for SyntaxData (#14301)

This should optimize memory usage for SyntaxData.

Author: rintaro

Diff for: include/swift/Syntax/RawSyntax.h

@@ -397,6 +397,12 @@ class RawSyntax final
     return {getTrailingObjects<RC<RawSyntax>>(), Bits.NumChildren};
   }
 
+  size_t getNumChildren() const {
+    if (isToken())
+      return 0;
+    return Bits.NumChildren;
+  }
+
   /// Get a child based on a particular node's "Cursor", indicating
   /// the position of the terms in the production of the Swift grammar.
   const RC<RawSyntax> &getChild(CursorIndex Index) const {

Diff for: include/swift/Syntax/SyntaxData.h

@@ -57,12 +57,13 @@ namespace syntax {
 /// This is only for holding a strong reference to the RawSyntax, a weak
 /// reference to the parent, and, in subclasses, lazily created strong
 /// references to non-terminal child nodes.
-class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
-  using RootDataPair = std::pair<RC<SyntaxData>, RC<SyntaxData>>;
+class SyntaxData final
+    : public llvm::ThreadSafeRefCountedBase<SyntaxData>,
+      private llvm::TrailingObjects<SyntaxData, AtomicCache<SyntaxData>> {
+  friend TrailingObjects;
 
-  llvm::SmallVector<AtomicCache<SyntaxData>, 10> Children;
+  using RootDataPair = std::pair<RC<SyntaxData>, RC<SyntaxData>>;
 
-public:
   /// The shared raw syntax representing this syntax data node.
   const RC<RawSyntax> Raw;
 
@@ -77,18 +78,16 @@ class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
   /// If there is no parent, this is 0.
   const CursorIndex IndexInParent;
 
+  size_t numTrailingObjects(OverloadToken<AtomicCache<SyntaxData>>) const {
+    return Raw->getNumChildren();
+  }
+
   SyntaxData(RC<RawSyntax> Raw, const SyntaxData *Parent = nullptr,
              CursorIndex IndexInParent = 0)
       : Raw(Raw), Parent(Parent), IndexInParent(IndexInParent) {
-    Children.resize(Raw ? Raw->getLayout().size() : 0);
-  }
-
-  /// Constructs a SyntaxNode by replacing `self` and recursively building
-  /// the parent chain up to the root.
-  template <typename SyntaxNode>
-  SyntaxNode replaceSelf(const RC<RawSyntax> NewRaw) const {
-    auto NewRootAndData = replaceSelf(NewRaw);
-    return { NewRootAndData.first, NewRootAndData.second.get() };
+    auto *I = getTrailingObjects<AtomicCache<SyntaxData>>();
+    for (auto *E = I + getNumChildren(); I != E; ++I)
+      ::new (static_cast<void *>(I)) AtomicCache<SyntaxData>();
   }
 
   /// With a new RawSyntax node, create a new node from this one and
@@ -97,8 +96,7 @@ class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
   /// DO NOT expose this as public API.
   RootDataPair replaceSelf(const RC<RawSyntax> NewRaw) const {
     if (hasParent()) {
-      auto NewRootAndParent =
-        getParent().getValue()->replaceChild(NewRaw, IndexInParent);
+      auto NewRootAndParent = Parent->replaceChild(NewRaw, IndexInParent);
       auto NewMe = NewRootAndParent.second->getChild(IndexInParent);
       return { NewRootAndParent.first, NewMe.get() };
     } else {
@@ -116,6 +114,35 @@ class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
     return nullptr;
   }
 
+  /// Replace a child in the raw syntax and recursively rebuild the
+  /// parental chain up to the root.
+  ///
+  /// DO NOT expose this as public API.
+  template <typename CursorType>
+  RootDataPair replaceChild(const RC<RawSyntax> RawChild,
+                            CursorType ChildCursor) const {
+    auto NewRaw = Raw->replaceChild(ChildCursor, RawChild);
+    return replaceSelf(NewRaw);
+  }
+
+  ArrayRef<AtomicCache<SyntaxData>> getChildren() const {
+    return {getTrailingObjects<AtomicCache<SyntaxData>>(), getNumChildren()};
+  }
+
+public:
+  ~SyntaxData() {
+    for (auto &I : getChildren())
+      I.~AtomicCache<SyntaxData>();
+  }
+
+  /// Constructs a SyntaxNode by replacing `self` and recursively building
+  /// the parent chain up to the root.
+  template <typename SyntaxNode>
+  SyntaxNode replaceSelf(const RC<RawSyntax> NewRaw) const {
+    auto NewRootAndData = replaceSelf(NewRaw);
+    return { NewRootAndData.first, NewRootAndData.second.get() };
+  }
+
   /// Replace a child in the raw syntax and recursively rebuild the
   /// parental chain up to the root.
   ///
@@ -130,18 +157,6 @@ class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
     };
   }
 
-  /// Replace a child in the raw syntax and recursively rebuild the
-  /// parental chain up to the root.
-  ///
-  /// DO NOT expose this as public API.
-  template <typename CursorType>
-  RootDataPair replaceChild(const RC<RawSyntax> RawChild,
-                            CursorType ChildCursor) const {
-    auto NewRaw = Raw->replaceChild(ChildCursor, RawChild);
-    return replaceSelf(NewRaw);
-  }
-
-public:
 
   static RC<SyntaxData> make(RC<RawSyntax> Raw,
                              const SyntaxData *Parent = nullptr,
@@ -158,11 +173,8 @@ class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
   }
 
   /// Return the parent syntax if there is one.
-  llvm::Optional<const SyntaxData *> getParent() const {
-    if (Parent != nullptr) {
-      return Parent;
-    }
-    return llvm::None;
+  const SyntaxData * getParent() const {
+    return Parent;
   }
 
   /// Returns true if this syntax node has a parent.
@@ -218,7 +230,7 @@ class SyntaxData final : public llvm::ThreadSafeRefCountedBase<SyntaxData> {
   RC<SyntaxData> getChild(size_t Index) const {
     if (!getRaw()->getChild(Index))
       return nullptr;
-    return Children[Index].getOrCreate([&]() {
+    return getChildren()[Index].getOrCreate([&]() {
       return realizeSyntaxNode(Index);
     });
   }
@@ -262,8 +274,8 @@ namespace llvm {
     }
     static unsigned getHashValue(const RCSD Value) {
       unsigned H = 0;
-      H ^= DenseMapInfo<uintptr_t>::getHashValue(reinterpret_cast<const uintptr_t>(Value->Raw.get()));
-      H ^= DenseMapInfo<uintptr_t>::getHashValue(reinterpret_cast<const uintptr_t>(Value->Parent));
+      H ^= DenseMapInfo<uintptr_t>::getHashValue(reinterpret_cast<const uintptr_t>(Value->getRaw().get()));
+      H ^= DenseMapInfo<uintptr_t>::getHashValue(reinterpret_cast<const uintptr_t>(Value->getParent()));
       H ^= DenseMapInfo<swift::syntax::CursorIndex>::getHashValue(Value->getIndexInParent());
       return H;
     }

Diff for: lib/Syntax/Syntax.cpp

@@ -75,8 +75,8 @@ bool Syntax::isMissing() const {
 }
 
 llvm::Optional<Syntax> Syntax::getParent() const {
-  auto ParentData = getData().Parent;
-  if (ParentData == nullptr) return llvm::None;
+  auto ParentData = getData().getParent();
+  if (!ParentData) return llvm::None;
   return llvm::Optional<Syntax> {
     Syntax { Root, ParentData }
   };

Diff for: lib/Syntax/SyntaxData.cpp

@@ -10,20 +10,17 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "swift/Syntax/UnknownSyntax.h"
-#include "llvm/Support/ErrorHandling.h"
+#include "swift/Syntax/SyntaxData.h"
 
 using namespace swift;
 using namespace swift::syntax;
 
 RC<SyntaxData> SyntaxData::make(RC<RawSyntax> Raw,
                                 const SyntaxData *Parent,
                                 CursorIndex IndexInParent) {
-  if (!Raw)
-    return nullptr;
-  return RC<SyntaxData> {
-    new SyntaxData(Raw, Parent, IndexInParent)
-  };
+  auto size = totalSizeToAlloc<AtomicCache<SyntaxData>>(Raw->getNumChildren());
+  void *data = ::operator new(size);
+  return RC<SyntaxData>{new (data) SyntaxData(Raw, Parent, IndexInParent)};
 }
 
 bool SyntaxData::isType() const {

Diff for: lib/Syntax/SyntaxNodes.cpp.gyb

@@ -28,7 +28,7 @@ using namespace swift::syntax;
 %   if node.requires_validation():
 void ${node.name}::validate() const {
 #ifndef NDEBUG
-  auto raw = Data->Raw;
+  auto raw = Data->getRaw();
   if (isMissing()) return;
   assert(raw->getLayout().size() == ${len(node.children)});
 %     for child in node.children:

Diff for: lib/Syntax/UnknownSyntax.cpp

@@ -17,6 +17,6 @@ using namespace swift;
 using namespace swift::syntax;
 
 void UnknownSyntax::validate() const {
-  assert(Data->Raw->isUnknown());
+  assert(Data->getRaw()->isUnknown());
 }