Merge pull request #64674 from mikeash/remotemirror-nodefactory-clear-fix

mikeash · web-flow · commit dff98f9e6812 · 2023-03-30T10:12:24.000-04:00
[RemoteMirror] Fix demangle node corruption caused by excessively eager clearing of the NodeFactory.
diff --git a/include/swift/Demangling/Demangler.h b/include/swift/Demangling/Demangler.h
@@ -53,7 +53,8 @@ class NodeFactory {
   /// The head of the single-linked slab list.
   Slab *CurrentSlab = nullptr;
 
-  /// The size of the previously allocated slab.
+  /// The size of the previously allocated slab. This may NOT be the size of
+  /// CurrentSlab, in the case where a checkpoint has been popped.
   ///
   /// The slab size can only grow, even clear() does not reset the slab size.
   /// This initial size is good enough to fit most de-manglings.
@@ -220,6 +221,26 @@ class NodeFactory {
     return {copiedString, str.size()};
   }
 
+  /// A checkpoint which captures the allocator's state at any given time. A
+  /// checkpoint can be popped to free all allocations made since it was made.
+  struct Checkpoint {
+    Slab *Slab;
+    char *CurPtr;
+    char *End;
+  };
+
+  /// Create a new checkpoint with the current state of the allocator.
+  Checkpoint pushCheckpoint() const;
+
+  /// Clear all allocations made since the given checkpoint. It is
+  /// undefined behavior to pop checkpoints in an order other than the
+  /// order in which they were pushed, or to pop a checkpoint when
+  /// clear() was called after creating it. The implementation attempts
+  /// to raise a fatal error in that case, but does not guarantee it. It
+  /// is allowed to pop outer checkpoints without popping inner ones, or
+  /// to call clear() without popping existing checkpoints.
+  void popCheckpoint(Checkpoint checkpoint);
+
   /// Creates a node of kind \p K.
   NodePointer createNode(Node::Kind K);
 
diff --git a/include/swift/RemoteInspection/TypeRefBuilder.h b/include/swift/RemoteInspection/TypeRefBuilder.h
@@ -456,7 +456,26 @@ class TypeRefBuilder {
 
   Demangle::NodeFactory &getNodeFactory() { return Dem; }
 
-  void clearNodeFactory() { Dem.clear(); }
+  NodeFactory::Checkpoint pushNodeFactoryCheckpoint() const {
+    return Dem.pushCheckpoint();
+  }
+
+  void popNodeFactoryCheckpoint(NodeFactory::Checkpoint checkpoint) {
+    Dem.popCheckpoint(checkpoint);
+  }
+
+  class ScopedNodeFactoryCheckpoint {
+    TypeRefBuilder *Builder;
+    NodeFactory::Checkpoint Checkpoint;
+
+  public:
+    ScopedNodeFactoryCheckpoint(TypeRefBuilder *Builder)
+        : Builder(Builder), Checkpoint(Builder->pushNodeFactoryCheckpoint()) {}
+
+    ~ScopedNodeFactoryCheckpoint() {
+      Builder->popNodeFactoryCheckpoint(Checkpoint);
+    }
+  };
 
   BuiltType decodeMangledType(Node *node, bool forRequirement = true);
 
@@ -1190,13 +1209,19 @@ class TypeRefBuilder {
       for (auto descriptor : sections.AssociatedType) {
         // Read out the relevant info from the associated type descriptor:
         // The type's name and which protocol conformance it corresponds to
-        auto typeRef = readTypeRef(descriptor, descriptor->ConformingTypeName);
-        auto typeName = nodeToString(demangleTypeRef(typeRef));
-        auto optionalMangledTypeName = normalizeReflectionName(typeRef);
-        auto protocolNode = demangleTypeRef(
-            readTypeRef(descriptor, descriptor->ProtocolTypeName));
-        auto protocolName = nodeToString(protocolNode);
-        clearNodeFactory();
+        llvm::Optional<std::string> optionalMangledTypeName;
+        std::string typeName;
+        std::string protocolName;
+        {
+          ScopedNodeFactoryCheckpoint checkpoint(this);
+          auto typeRef =
+              readTypeRef(descriptor, descriptor->ConformingTypeName);
+          typeName = nodeToString(demangleTypeRef(typeRef));
+          optionalMangledTypeName = normalizeReflectionName(typeRef);
+          auto protocolNode = demangleTypeRef(
+              readTypeRef(descriptor, descriptor->ProtocolTypeName));
+          protocolName = nodeToString(protocolNode);
+        }
         if (optionalMangledTypeName.has_value()) {
           auto mangledTypeName = optionalMangledTypeName.value();
           if (forMangledTypeName.has_value()) {
@@ -2011,10 +2036,10 @@ class TypeRefBuilder {
     std::unordered_map<std::string, std::string> typeNameToManglingMap;
     for (const auto &section : ReflectionInfos) {
       for (auto descriptor : section.Field) {
+        ScopedNodeFactoryCheckpoint checkpoint(this);
         auto TypeRef = readTypeRef(descriptor, descriptor->MangledTypeName);
         auto OptionalMangledTypeName = normalizeReflectionName(TypeRef);
         auto TypeName = nodeToString(demangleTypeRef(TypeRef));
-        clearNodeFactory();
         if (OptionalMangledTypeName.has_value()) {
           typeNameToManglingMap[TypeName] = OptionalMangledTypeName.value();
         }
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -505,13 +505,91 @@ void NodeFactory::clear() {
     // cycles.
     CurrentSlab->Previous = nullptr;
     CurPtr = (char *)(CurrentSlab + 1);
-    assert(End == CurPtr + SlabSize);
 #ifdef NODE_FACTORY_DEBUGGING
     allocatedMemory = 0;
 #endif
   }
 }
 
+NodeFactory::Checkpoint NodeFactory::pushCheckpoint() const {
+  return {CurrentSlab, CurPtr, End};
+}
+
+void NodeFactory::popCheckpoint(NodeFactory::Checkpoint checkpoint) {
+  if (checkpoint.Slab == CurrentSlab) {
+    if (checkpoint.CurPtr > CurPtr) {
+      fatal(0,
+            "Popping checkpoint {%p, %p, %p} that is after the current "
+            "pointer.\n",
+            checkpoint.Slab, checkpoint.CurPtr, checkpoint.End);
+    }
+    if (checkpoint.End != End) {
+      fatal(0,
+            "Popping checkpoint {%p, %p, %p} with End that does not match "
+            "current End %p.\n",
+            checkpoint.Slab, checkpoint.CurPtr, checkpoint.End, End);
+    }
+#ifndef NDEBUG
+    // Scribble the newly freed area.
+    memset(checkpoint.CurPtr, 0xaa, CurPtr - checkpoint.CurPtr);
+#endif
+    CurPtr = checkpoint.CurPtr;
+  } else {
+    // We may want to keep using the current slab rather than destroying
+    // it, since over time this allows us to converge on a steady state
+    // with no allocation activity (see the comment above in
+    // NodeFactory::clear). Keep using the current slab if the free space in the
+    // checkpoint's slab is less than 1/16th of the current slab's space. We
+    // won't repeatedly allocate and deallocate the current slab. The size
+    // doubles each time so we'll quickly pass the threshold.
+    Slab *savedSlab = nullptr;
+    if (CurrentSlab) {
+      size_t checkpointSlabFreeSpace = checkpoint.End - checkpoint.CurPtr;
+      size_t currentSlabSize = End - (char *)(CurrentSlab + 1);
+      if (currentSlabSize / 16 > checkpointSlabFreeSpace) {
+        savedSlab = CurrentSlab;
+        CurrentSlab = CurrentSlab->Previous;
+        // No need to save End, as it will still be in place later.
+      }
+    }
+
+    // Free all slabs (possibly except the one we saved) until we find the end
+    // or we find the checkpoint.
+    while (CurrentSlab && checkpoint.Slab != CurrentSlab) {
+      auto Slab = CurrentSlab;
+      CurrentSlab = CurrentSlab->Previous;
+      free(Slab);
+    }
+
+    // If we ran to the end and the checkpoint actually has a slab pointer, then
+    // the checkpoint is invalid.
+    if (!CurrentSlab && checkpoint.Slab) {
+      fatal(0,
+            "Popping checkpoint {%p, %p, %p} with slab that is not within "
+            "the allocator's slab chain.\n",
+            checkpoint.Slab, checkpoint.CurPtr, checkpoint.End);
+    }
+
+    if (savedSlab) {
+      // Reinstall the saved slab.
+      savedSlab->Previous = CurrentSlab;
+      CurrentSlab = savedSlab;
+      CurPtr = (char *)(CurrentSlab + 1);
+      // End is still valid from before.
+    } else {
+      // Set CurPtr and End to match the checkpoint's position.
+      CurPtr = checkpoint.CurPtr;
+      End = checkpoint.End;
+    }
+
+#ifndef NDEBUG
+    // Scribble the now freed end of the slab.
+    if (CurPtr)
+      memset(CurPtr, 0xaa, End - CurPtr);
+#endif
+  }
+}
+
 NodePointer NodeFactory::createNode(Node::Kind K) {
   return new (Allocate<Node>()) Node(K);
 }
diff --git a/stdlib/public/RemoteInspection/TypeRefBuilder.cpp b/stdlib/public/RemoteInspection/TypeRefBuilder.cpp
@@ -93,6 +93,7 @@ RemoteRef<char> TypeRefBuilder::readTypeRef(uint64_t remoteAddr) {
 /// Load and normalize a mangled name so it can be matched with string equality.
 llvm::Optional<std::string>
 TypeRefBuilder::normalizeReflectionName(RemoteRef<char> reflectionName) {
+  ScopedNodeFactoryCheckpoint checkpoint(this);
   // Remangle the reflection name to resolve symbolic references.
   if (auto node = demangleTypeRef(reflectionName,
                                   /*useOpaqueTypeSymbolicReferences*/ false)) {
@@ -104,7 +105,6 @@ TypeRefBuilder::normalizeReflectionName(RemoteRef<char> reflectionName) {
       return {};
     default:
       auto mangling = mangleNode(node);
-      clearNodeFactory();
       if (!mangling.isSuccess()) {
         return {};
       }
@@ -159,11 +159,11 @@ lookupTypeWitness(const std::string &MangledTypeName,
             0)
           continue;
 
+        ScopedNodeFactoryCheckpoint checkpoint(this);
         auto SubstitutedTypeName = readTypeRef(AssocTy,
                                                AssocTy->SubstitutedTypeName);
         auto Demangled = demangleTypeRef(SubstitutedTypeName);
         auto *TypeWitness = decodeMangledType(Demangled);
-        clearNodeFactory();
 
         AssociatedTypeCache.insert(std::make_pair(key, TypeWitness));
         return TypeWitness;
@@ -181,9 +181,9 @@ const TypeRef *TypeRefBuilder::lookupSuperclass(const TypeRef *TR) {
   if (!FD->hasSuperclass())
     return nullptr;
 
+  ScopedNodeFactoryCheckpoint checkpoint(this);
   auto Demangled = demangleTypeRef(readTypeRef(FD, FD->Superclass));
   auto Unsubstituted = decodeMangledType(Demangled);
-  clearNodeFactory();
   if (!Unsubstituted)
     return nullptr;
 
@@ -367,9 +367,9 @@ bool TypeRefBuilder::getFieldTypeRefs(
       continue;
     }
 
+    ScopedNodeFactoryCheckpoint checkpoint(this);
     auto Demangled = demangleTypeRef(readTypeRef(Field,Field->MangledTypeName));
     auto Unsubstituted = decodeMangledType(Demangled);
-    clearNodeFactory();
     if (!Unsubstituted)
       return false;
 
@@ -486,10 +486,10 @@ TypeRefBuilder::getClosureContextInfo(RemoteRef<CaptureDescriptor> CD) {
     auto CR = CD.getField(*i);
     
     if (CR->hasMangledTypeName()) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto MangledName = readTypeRef(CR, CR->MangledTypeName);
       auto DemangleTree = demangleTypeRef(MangledName);
       TR = decodeMangledType(DemangleTree);
-      clearNodeFactory();
     }
     Info.CaptureTypes.push_back(TR);
   }
@@ -499,10 +499,10 @@ TypeRefBuilder::getClosureContextInfo(RemoteRef<CaptureDescriptor> CD) {
     auto MSR = CD.getField(*i);
     
     if (MSR->hasMangledTypeName()) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto MangledName = readTypeRef(MSR, MSR->MangledTypeName);
       auto DemangleTree = demangleTypeRef(MangledName);
       TR = decodeMangledType(DemangleTree);
-      clearNodeFactory();
     }
 
     const MetadataSource *MS = nullptr;
@@ -526,11 +526,11 @@ TypeRefBuilder::getClosureContextInfo(RemoteRef<CaptureDescriptor> CD) {
 
 void TypeRefBuilder::dumpTypeRef(RemoteRef<char> MangledName,
                                  std::ostream &stream, bool printTypeName) {
+  ScopedNodeFactoryCheckpoint checkpoint(this);
   auto DemangleTree = demangleTypeRef(MangledName);
   auto TypeName = nodeToString(DemangleTree);
   stream << TypeName << "\n";
   auto Result = swift::Demangle::decodeMangledType(*this, DemangleTree);
-  clearNodeFactory();
   if (Result.isError()) {
     auto *Error = Result.getError();
     char *ErrorStr = Error->copyErrorString();
@@ -549,10 +549,14 @@ FieldTypeCollectionResult TypeRefBuilder::collectFieldTypes(
   FieldTypeCollectionResult result;
   for (const auto &sections : ReflectionInfos) {
     for (auto descriptor : sections.Field) {
-      auto typeRef = readTypeRef(descriptor, descriptor->MangledTypeName);
-      auto typeName = nodeToString(demangleTypeRef(typeRef));
-      auto optionalMangledTypeName = normalizeReflectionName(typeRef);
-      clearNodeFactory();
+      llvm::Optional<std::string> optionalMangledTypeName;
+      std::string typeName;
+      {
+        ScopedNodeFactoryCheckpoint checkpoint(this);
+        auto typeRef = readTypeRef(descriptor, descriptor->MangledTypeName);
+        typeName = nodeToString(demangleTypeRef(typeRef));
+        optionalMangledTypeName = normalizeReflectionName(typeRef);
+      }
       if (optionalMangledTypeName.has_value()) {
         auto mangledTypeName =
           optionalMangledTypeName.value();
@@ -618,10 +622,10 @@ void TypeRefBuilder::dumpFieldSection(std::ostream &stream) {
 void TypeRefBuilder::dumpBuiltinTypeSection(std::ostream &stream) {
   for (const auto &sections : ReflectionInfos) {
     for (auto descriptor : sections.Builtin) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto typeNode =
           demangleTypeRef(readTypeRef(descriptor, descriptor->TypeName));
       auto typeName = nodeToString(typeNode);
-      clearNodeFactory();
 
       stream << "\n- " << typeName << ":\n";
       stream << "Size: " << descriptor->Size << "\n";
@@ -670,10 +674,10 @@ void TypeRefBuilder::dumpCaptureSection(std::ostream &stream) {
 void TypeRefBuilder::dumpMultiPayloadEnumSection(std::ostream &stream) {
   for (const auto &sections : ReflectionInfos) {
     for (const auto descriptor : sections.MultiPayloadEnum) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto typeNode =
           demangleTypeRef(readTypeRef(descriptor, descriptor->TypeName));
       auto typeName = nodeToString(typeNode);
-      clearNodeFactory();
 
       stream << "\n- " << typeName << ":\n";
       stream << "  Descriptor Size: " << descriptor->getSizeInBytes() << "\n";
