Merge pull request llvm#22655 from eeckstein/stack-allocated-demangler

Make the demangler in the runtime use stack allocated memory.

Author: eeckstein

include/swift/Demangling/Demangle.h

@@ -142,32 +142,38 @@ class Node {
   friend class NodeFactory;
 
 private:
-  Kind NodeKind;
 
-  enum class PayloadKind : uint8_t {
-    None, Text, Index
+  struct NodeVector {
+    NodePointer *Nodes;
+    uint32_t Number = 0;
+    uint32_t Capacity = 0;
   };
-  PayloadKind NodePayloadKind;
 
   union {
-    llvm::StringRef TextPayload;
-    IndexType IndexPayload;
+    llvm::StringRef Text;
+    IndexType Index;
+    NodePointer InlineChildren[2];
+    NodeVector Children;
   };
 
-  NodePointer *Children = nullptr;
-  size_t NumChildren = 0;
-  size_t ReservedChildren = 0;
+
+  Kind NodeKind;
+
+  enum class PayloadKind : uint8_t {
+    None, Text, Index, OneChild, TwoChildren, ManyChildren
+  };
+  PayloadKind NodePayloadKind;
 
   Node(Kind k)
       : NodeKind(k), NodePayloadKind(PayloadKind::None) {
   }
   Node(Kind k, llvm::StringRef t)
       : NodeKind(k), NodePayloadKind(PayloadKind::Text) {
-    TextPayload = t;
+    Text = t;
   }
   Node(Kind k, IndexType index)
       : NodeKind(k), NodePayloadKind(PayloadKind::Index) {
-    IndexPayload = index;
+    Index = index;
   }
   Node(const Node &) = delete;
   Node &operator=(const Node &) = delete;
@@ -178,37 +184,33 @@ class Node {
   bool hasText() const { return NodePayloadKind == PayloadKind::Text; }
   llvm::StringRef getText() const {
     assert(hasText());
-    return TextPayload;
+    return Text;
   }
 
   bool hasIndex() const { return NodePayloadKind == PayloadKind::Index; }
   uint64_t getIndex() const {
     assert(hasIndex());
-    return IndexPayload;
+    return Index;
   }
 
-  using iterator = NodePointer *;
-  using const_iterator = const NodePointer *;
-  using size_type = size_t;
+  using iterator = const NodePointer *;
+
+  size_t getNumChildren() const;
 
-  bool hasChildren() const { return NumChildren != 0; }
-  size_t getNumChildren() const { return NumChildren; }
-  iterator begin() { return Children; }
-  iterator end() { return Children + NumChildren; }
-  const_iterator begin() const { return Children; }
-  const_iterator end() const { return Children + NumChildren; }
+  bool hasChildren() const { return getNumChildren() != 0; }
+
+  iterator begin() const;
+
+  iterator end() const;
 
   NodePointer getFirstChild() const {
-    assert(NumChildren >= 1);
-    return Children[0];
+    return getChild(0);
   }
   NodePointer getChild(size_t index) const {
-    assert(NumChildren > index);
-    return Children[index];
+    assert(getNumChildren() > index);
+    return begin()[index];
   }
 
-//  inline void addChild(NodePointer Child, Context &Ctx);
-
   // Only to be used by the demangler parsers.
   void addChild(NodePointer Child, NodeFactory &Factory);
   // Only to be used by the demangler parsers.
@@ -472,8 +474,10 @@ void mangleIdentifier(const char *data, size_t length,
 
 /// Remangle a demangled parse tree.
 ///
-/// This should always round-trip perfectly with demangleSymbolAsNode.
-std::string mangleNode(NodePointer root);
+/// If \p BorrowFrom is specified, the initial bump pointer memory is
+/// borrowed from the free memory of BorrowFrom.
+std::string mangleNode(NodePointer root,
+                       NodeFactory *BorrowFrom = nullptr);
 
 using SymbolicResolver =
   llvm::function_ref<Demangle::NodePointer (SymbolicReferenceKind,
@@ -482,14 +486,17 @@ using SymbolicResolver =
 /// Remangle a demangled parse tree, using a callback to resolve
 /// symbolic references.
 ///
-/// This should always round-trip perfectly with demangleSymbolAsNode.
-std::string mangleNode(NodePointer root, SymbolicResolver resolver);
+/// If \p BorrowFrom is specified, the initial bump pointer memory is
+/// borrowed from the free memory of BorrowFrom.
+std::string mangleNode(NodePointer root, SymbolicResolver resolver,
+                       NodeFactory *BorrowFrom = nullptr);
 
 /// Remangle in the old mangling scheme.
 ///
-/// This is only used for objc-runtime names and should be removed as soon as
-/// we switch to the new mangling for those names as well.
-std::string mangleNodeOld(NodePointer root);
+/// This is only used for objc-runtime names.
+/// If \p BorrowFrom is specified, the initial bump pointer memory is
+/// borrowed from the free memory of BorrowFrom.
+std::string mangleNodeOld(NodePointer root, NodeFactory *BorrowFrom = nullptr);
 
 /// Transform the node structure to a string.
 ///

include/swift/Demangling/DemangleNodes.def

@@ -88,6 +88,7 @@ NODE(FullTypeMetadata)
 CONTEXT_NODE(Function)
 NODE(FunctionSignatureSpecialization)
 NODE(FunctionSignatureSpecializationParam)
+NODE(FunctionSignatureSpecializationReturn)
 NODE(FunctionSignatureSpecializationParamKind)
 NODE(FunctionSignatureSpecializationParamPayload)
 NODE(FunctionType)

include/swift/Demangling/Demangler.h

@@ -69,31 +69,81 @@ class NodeFactory {
   }
 
   static void freeSlabs(Slab *slab);
-  
+
+  /// If not null, the NodeFactory from which this factory borrowed free memory.
+  NodeFactory *BorrowedFrom = nullptr;
+
+  /// True if some other NodeFactory borrowed free memory from this factory.
+  bool isBorrowed = false;
+
+#ifdef NODE_FACTORY_DEBUGGING
+  size_t allocatedMemory = 0;
+  static int nestingLevel;
+  std::string indent() { return std::string(nestingLevel * 2, ' '); }
+#endif
+
 public:
 
   NodeFactory() {
 #ifdef NODE_FACTORY_DEBUGGING
-    std::cerr << "## New NodeFactory " << this << "\n";
+    std::cerr << indent() << "## New NodeFactory\n";
+    nestingLevel++;
 #endif
   }
-  
+
+  /// Provide pre-allocated memory, e.g. memory on the stack.
+  ///
+  /// Only if this memory overflows, the factory begins to malloc.
+  void providePreallocatedMemory(char *Memory, size_t Size) {
+#ifdef NODE_FACTORY_DEBUGGING
+    std::cerr << indent() << "++ provide preallocated memory, size = "
+                          << Size << '\n';
+#endif
+    assert(!CurPtr && !End && !CurrentSlab);
+    CurPtr = Memory;
+    End = CurPtr + Size;
+  }
+
+  /// Borrow free memory from another factory \p BorrowFrom.
+  ///
+  /// While this factory is alive, no allocations can be done in the
+  /// \p BorrowFrom factory.
+  void providePreallocatedMemory(NodeFactory &BorrowFrom) {
+    assert(!CurPtr && !End && !CurrentSlab);
+    assert(!BorrowFrom.isBorrowed && !BorrowedFrom);
+    BorrowFrom.isBorrowed = true;
+    BorrowedFrom = &BorrowFrom;
+    CurPtr = BorrowFrom.CurPtr;
+    End = BorrowFrom.End;
+#ifdef NODE_FACTORY_DEBUGGING
+    std::cerr << indent() << "++ borrow memory, size = "
+                          << (End - CurPtr) << '\n';
+#endif
+  }
+
   virtual ~NodeFactory() {
     freeSlabs(CurrentSlab);
 #ifdef NODE_FACTORY_DEBUGGING
-    std::cerr << "Delete NodeFactory " << this << "\n";
+    nestingLevel--;
+    std::cerr << indent() << "## Delete NodeFactory: allocated memory = "
+                          << allocatedMemory  << '\n';
 #endif
+    if (BorrowedFrom) {
+      BorrowedFrom->isBorrowed = false;
+    }
   }
 
   virtual void clear();
 
   /// Allocates an object of type T or an array of objects of type T.
   template<typename T> T *Allocate(size_t NumObjects = 1) {
+    assert(!isBorrowed);
     size_t ObjectSize = NumObjects * sizeof(T);
     CurPtr = align(CurPtr, alignof(T));
 #ifdef NODE_FACTORY_DEBUGGING
-    std::cerr << "  alloc " << ObjectSize << ", CurPtr = "
+    std::cerr << indent() << "alloc " << ObjectSize << ", CurPtr = "
               << (void *)CurPtr << "\n";
+    allocatedMemory += ObjectSize;
 #endif
 
     // Do we have enough space in the current slab?
@@ -113,7 +163,7 @@ class NodeFactory {
       End = (char *)newSlab + AllocSize;
       assert(CurPtr + ObjectSize <= End);
 #ifdef NODE_FACTORY_DEBUGGING
-      std::cerr << "    ** new slab " << newSlab << ", allocsize = "
+      std::cerr << indent() << "** new slab " << newSlab << ", allocsize = "
                 << AllocSize << ", CurPtr = " << (void *)CurPtr
                 << ", End = " << (void *)End << "\n";
 #endif
@@ -132,14 +182,15 @@ class NodeFactory {
   /// new memory address.
   /// The \p Capacity is enlarged at least by \p MinGrowth, but can also be
   /// enlarged by a bigger value.
-  template<typename T> void Reallocate(T *&Objects, size_t &Capacity,
+  template<typename T> void Reallocate(T *&Objects, uint32_t &Capacity,
                                        size_t MinGrowth) {
+    assert(!isBorrowed);
     size_t OldAllocSize = Capacity * sizeof(T);
     size_t AdditionalAlloc = MinGrowth * sizeof(T);
 
 #ifdef NODE_FACTORY_DEBUGGING
-    std::cerr << "  realloc " << Objects << ", num = " << NumObjects
-              << " (size = " << OldAllocSize << "), Growth = " << Growth
+    std::cerr << indent() << "realloc: capacity = " << Capacity
+              << " (size = " << OldAllocSize << "), growth = " << MinGrowth
               << " (size = " << AdditionalAlloc << ")\n";
 #endif
     if ((char *)Objects + OldAllocSize == CurPtr
@@ -149,7 +200,8 @@ class NodeFactory {
       CurPtr += AdditionalAlloc;
       Capacity += MinGrowth;
 #ifdef NODE_FACTORY_DEBUGGING
-      std::cerr << "    ** can grow: CurPtr = " << (void *)CurPtr << "\n";
+      std::cerr << indent() << "** can grow: " << (void *)CurPtr << '\n';
+      allocatedMemory += AdditionalAlloc;
 #endif
       return;
     }
@@ -203,8 +255,8 @@ template<typename T> class Vector {
 
 protected:
   T *Elems = nullptr;
-  size_t NumElems = 0;
-  size_t Capacity = 0;
+  uint32_t NumElems = 0;
+  uint32_t Capacity = 0;
 
 public:
   using iterator = T *;
@@ -469,7 +521,7 @@ class Demangler : public NodeFactory {
   NodePointer demangleThunkOrSpecialization();
   NodePointer demangleGenericSpecialization(Node::Kind SpecKind);
   NodePointer demangleFunctionSpecialization();
-  NodePointer demangleFuncSpecParam(Node::IndexType ParamIdx);
+  NodePointer demangleFuncSpecParam(Node::Kind Kind);
   NodePointer addFuncSpecParamNumber(NodePointer Param,
                               FunctionSigSpecializationParamKind Kind);
 
@@ -536,7 +588,19 @@ class Demangler : public NodeFactory {
   /// Demangler or with a call of clear().
   NodePointer demangleType(StringRef MangledName);
 };
-  
+
+/// A demangler which uses stack space for its initial memory.
+///
+/// The \p Size paramter specifies the size of the stack space.
+template <size_t Size> class StackAllocatedDemangler : public Demangler {
+  char StackSpace[Size];
+
+public:
+  StackAllocatedDemangler() {
+    providePreallocatedMemory(StackSpace, Size);
+  }
+};
+
 NodePointer demangleOldSymbolAsNode(StringRef MangledName,
                                     NodeFactory &Factory);
 } // end namespace Demangle

include/swift/Demangling/TypeDecoder.h

@@ -676,21 +676,21 @@ class TypeDecoder {
       auto base = decodeMangledType(Node->getChild(0));
       if (!base)
         return BuiltType();
-      auto member = Node->getChild(1)->getText();
       auto assocTypeChild = Node->getChild(1);
-      if (assocTypeChild->getNumChildren() < 1)
+      auto member = assocTypeChild->getFirstChild()->getText();
+      if (assocTypeChild->getNumChildren() < 2)
         return Builder.createDependentMemberType(member, base);
 
-      auto protocol = decodeMangledProtocolType(assocTypeChild->getChild(0));
+      auto protocol = decodeMangledProtocolType(assocTypeChild->getChild(1));
       if (!protocol)
         return BuiltType();
       return Builder.createDependentMemberType(member, base, protocol);
     }
     case NodeKind::DependentAssociatedTypeRef: {
-      if (Node->getNumChildren() < 1)
+      if (Node->getNumChildren() < 2)
         return BuiltType();
 
-      return decodeMangledType(Node->getChild(0));
+      return decodeMangledType(Node->getChild(1));
     }
     case NodeKind::Unowned: {
       if (Node->getNumChildren() < 1)

include/swift/Reflection/TypeRefBuilder.h

@@ -210,12 +210,12 @@ class TypeRefBuilder {
 
   Optional<std::string>
   createTypeDecl(Node *node, bool &typeAlias) {
-    return Demangle::mangleNode(node);
+    return Demangle::mangleNode(node, &Dem);
   }
 
   BuiltProtocolDecl
   createProtocolDecl(Node *node) {
-    return std::make_pair(Demangle::mangleNode(node), false);
+    return std::make_pair(Demangle::mangleNode(node, &Dem), false);
   }
 
   BuiltProtocolDecl

include/swift/Remote/MetadataReader.h

@@ -2109,9 +2109,10 @@ class MetadataReader {
     }
 
     if (importInfo && !importInfo->RelatedEntityName.empty()) {
-      auto relatedNode =
-        dem.createNode(Node::Kind::RelatedEntityDeclName,
-                               std::move(importInfo->RelatedEntityName));
+      auto kindNode = dem.createNode(Node::Kind::Identifier,
+                                 std::move(importInfo->RelatedEntityName));
+      auto relatedNode = dem.createNode(Node::Kind::RelatedEntityDeclName);
+      relatedNode->addChild(kindNode, dem);
       relatedNode->addChild(nameNode, dem);
       nameNode = relatedNode;
     }

lib/AST/ASTDemangler.cpp

@@ -820,8 +820,8 @@ ASTBuilder::findDeclContext(NodePointer node) {
 
     } else if (declNameNode->getKind() ==
                  Demangle::Node::Kind::RelatedEntityDeclName) {
-      name = declNameNode->getChild(0)->getText();
-      relatedEntityKind = declNameNode->getText();
+      name = declNameNode->getChild(1)->getText();
+      relatedEntityKind = declNameNode->getFirstChild()->getText();
 
     // Ignore any other decl-name productions for now.
     } else {