Mangler: Fix substitution ordering when mangling opaque return types.

When a declaration has a structural opaque return type like:

  func foo() -> Bar<some P>

then to mangle that return type `Bar<some P>`, we have to mangle the `some P`
part by referencing its defining declaration `foo()`, which in turn includes
its return type `Bar<some P>` again (this time using a special mangling for
`some P` that prevents infinite recursion). Since we mangle `Bar<some P>`
once as part of mangling the declaration, and we register substitutions for
bound generic types when they're complete, we end up registering the
substitution for `Bar<some P>` twice, once as the return type of the
declaration name, and again as the actual type. This would be fine, except
that the mangler doesn't check for key collisions, and it picks
substitution indexes based on the number of entries in its hash map, so
the duplicated substitution ends up corrupting the substitution sequence,
causing the mangler to produce an invalid mangled name.

Fixing that exposes us to another problem in the remangler: the AST
mangler keys substitutions by type identity, but the remangler
uses the value of the demangled nodes to recognize substitutions.
The mangling for `Bar<current declaration's opaque return type>` can
appear multiple times in a demangled tree, but referring to different
declarations' opaque return types, and the remangler would reconstruct
an incorrect mangled name when this happens. To avoid this, change the
way the demangler represents `OpaqueReturnType` nodes so that they
contain a backreference to the declaration they represent, so that
substitutions involving different declarations' opaque return types
don't get confused.

Author: jckarter

include/swift/Basic/Mangler.h

@@ -54,6 +54,12 @@ class Mangler {
 
   /// Identifier substitutions.
   llvm::StringMap<unsigned> StringSubstitutions;
+  
+  /// Index to use for the next added substitution.
+  /// Note that this is not simply the sum of the size of the \c Substitutions
+  /// and \c StringSubstitutions maps above, since in some circumstances the
+  /// same entity may be registered for multiple substitution indexes.
+  unsigned NextSubstitutionIndex = 0;
 
   /// Word substitutions in mangled identifiers.
   llvm::SmallVector<SubstitutionWord, 26> Words;
@@ -102,7 +108,6 @@ class Mangler {
   }
 
 protected:
-
   Mangler() : Buffer(Storage) { }
 
   /// Begins a new mangling but does not add the mangling prefix.
@@ -134,15 +139,15 @@ class Mangler {
   void addSubstitution(const void *ptr) {
     if (!UseSubstitutions)
       return;
-
-    auto value = Substitutions.size() + StringSubstitutions.size();
+    
+    auto value = NextSubstitutionIndex++;
     Substitutions[ptr] = value;
   }
   void addSubstitution(StringRef Str) {
     if (!UseSubstitutions)
       return;
 
-    auto value = Substitutions.size() + StringSubstitutions.size();
+    auto value = NextSubstitutionIndex++;
     StringSubstitutions[Str] = value;
   }
 

include/swift/Demangling/DemangleNodes.def

@@ -339,7 +339,6 @@ NODE(AccessibleFunctionRecord)
 NODE(CompileTimeConst)
 
 // Added in Swift 5.7
-NODE(OpaqueReturnTypeIndexed)
 NODE(BackDeploymentThunk)
 NODE(BackDeploymentFallback)
 NODE(ExtendedExistentialTypeShape)
@@ -351,6 +350,8 @@ NODE(SymbolicExtendedExistentialType)
 // Added in Swift 5.8
 NODE(MetatypeParamsRemoved)
 NODE(HasSymbolQuery)
+NODE(OpaqueReturnTypeIndex)
+NODE(OpaqueReturnTypeParent)
 
 #undef CONTEXT_NODE
 #undef NODE

include/swift/Demangling/Demangler.h

@@ -475,8 +475,9 @@ class Demangler : public NodeFactory {
   friend DemangleInitRAII;
 
   void addSubstitution(NodePointer Nd) {
-    if (Nd)
+    if (Nd) {
       Substitutions.push_back(Nd, *this);
+    }
   }
 
   NodePointer addChild(NodePointer Parent, NodePointer Child);

lib/AST/ASTMangler.cpp

@@ -3237,7 +3237,7 @@ void ASTMangler::appendDeclType(const ValueDecl *decl,
   } else {
     appendType(type, sig, decl);
   }
-
+  
   // Mangle the generic signature, if any.
   if (genericSig && appendGenericSignature(genericSig, parentGenericSig)) {
     // The 'F' function mangling doesn't need a 'u' for its generic signature.

lib/Basic/Mangler.cpp

@@ -114,6 +114,7 @@ void Mangler::beginManglingWithoutPrefix() {
   Storage.clear();
   Substitutions.clear();
   StringSubstitutions.clear();
+  NextSubstitutionIndex = 0;
   Words.clear();
   SubstMerging.clear();
 }

lib/Demangling/Demangler.cpp

@@ -1328,6 +1328,39 @@ NodePointer Demangler::demangleExtensionContext() {
   return Ext;
 }
 
+/// Associate any \c OpaqueReturnType nodes with the declaration whose opaque
+/// return type they refer back to.
+static Node *setParentForOpaqueReturnTypeNodes(Demangler &D,
+                                              Node *parent,
+                                              Node *visitedNode) {
+  if (!parent || !visitedNode)
+    return nullptr;
+  if (visitedNode->getKind() == Node::Kind::OpaqueReturnType) {
+    // If this node is not already parented, parent it.
+    if (visitedNode->hasChildren()
+        && visitedNode->getLastChild()->getKind() == Node::Kind::OpaqueReturnTypeParent) {
+      return parent;
+    }
+    visitedNode->addChild(D.createNode(Node::Kind::OpaqueReturnTypeParent,
+                                       (Node::IndexType)parent), D);
+    return parent;
+  }
+  
+  // If this node is one that may in turn define its own opaque return type,
+  // stop recursion, since any opaque return type nodes underneath would refer
+  // to the nested declaration rather than the one we're looking at.
+  if (visitedNode->getKind() == Node::Kind::Function
+      || visitedNode->getKind() == Node::Kind::Variable
+      || visitedNode->getKind() == Node::Kind::Subscript) {
+    return parent;
+  }
+  
+  for (size_t i = 0, e = visitedNode->getNumChildren(); i < e; ++i) {
+    setParentForOpaqueReturnTypeNodes(D, parent, visitedNode->getChild(i));
+  }
+  return parent;
+}
+
 NodePointer Demangler::demanglePlainFunction() {
   NodePointer GenSig = popNode(Node::Kind::DependentGenericSignature);
   NodePointer Type = popFunctionType(Node::Kind::FunctionType);
@@ -1341,10 +1374,12 @@ NodePointer Demangler::demanglePlainFunction() {
   auto Name = popNode(isDeclName);
   auto Ctx = popContext();
 
-  if (LabelList)
-    return createWithChildren(Node::Kind::Function, Ctx, Name, LabelList, Type);
-
-  return createWithChildren(Node::Kind::Function, Ctx, Name, Type);
+  NodePointer result = LabelList
+    ? createWithChildren(Node::Kind::Function, Ctx, Name, LabelList, Type)
+    : createWithChildren(Node::Kind::Function, Ctx, Name, Type);
+    
+  result = setParentForOpaqueReturnTypeNodes(*this, result, Type);
+  return result;
 }
 
 NodePointer Demangler::popFunctionType(Node::Kind kind, bool hasClangType) {
@@ -2270,7 +2305,8 @@ NodePointer Demangler::demangleArchetype() {
     int ordinal = demangleIndex();
     if (ordinal < 0)
       return NULL;
-    return createType(createNode(Node::Kind::OpaqueReturnTypeIndexed, ordinal));
+    return createType(createWithChild(Node::Kind::OpaqueReturnType,
+                      createNode(Node::Kind::OpaqueReturnTypeIndex, ordinal)));
   }
 
   case 'x': {
@@ -3360,7 +3396,7 @@ NodePointer Demangler::demangleSpecialType() {
         return nullptr;
       // Build layout.
       auto layout = createNode(Node::Kind::SILBoxLayout);
-      for (unsigned i = 0, e = fieldTypes->getNumChildren(); i < e; ++i) {
+      for (size_t i = 0, e = fieldTypes->getNumChildren(); i < e; ++i) {
         auto fieldType = fieldTypes->getChild(i);
         assert(fieldType->getKind() == Node::Kind::Type);
         bool isMutable = false;
@@ -3600,8 +3636,11 @@ NodePointer Demangler::demangleEntity(Node::Kind Kind) {
   NodePointer LabelList = popFunctionParamLabels(Type);
   NodePointer Name = popNode(isDeclName);
   NodePointer Context = popContext();
-  return LabelList ? createWithChildren(Kind, Context, Name, LabelList, Type)
-                   : createWithChildren(Kind, Context, Name, Type);
+  auto result = LabelList ? createWithChildren(Kind, Context, Name, LabelList, Type)
+                          : createWithChildren(Kind, Context, Name, Type);
+                          
+  result = setParentForOpaqueReturnTypeNodes(*this, result, Type);
+  return result;
 }
 
 NodePointer Demangler::demangleVariable() {
@@ -3623,6 +3662,8 @@ NodePointer Demangler::demangleSubscript() {
   addChild(Subscript, LabelList);
   Subscript = addChild(Subscript, Type);
   addChild(Subscript, PrivateName);
+  
+  Subscript = setParentForOpaqueReturnTypeNodes(*this, Subscript, Type);
 
   return demangleAccessor(Subscript);
 }

lib/Demangling/NodePrinter.cpp

@@ -579,7 +579,8 @@ class NodePrinter {
     case Node::Kind::OpaqueType:
     case Node::Kind::OpaqueTypeDescriptorSymbolicReference:
     case Node::Kind::OpaqueReturnType:
-    case Node::Kind::OpaqueReturnTypeIndexed:
+    case Node::Kind::OpaqueReturnTypeIndex:
+    case Node::Kind::OpaqueReturnTypeParent:
     case Node::Kind::OpaqueReturnTypeOf:
     case Node::Kind::CanonicalSpecializedGenericMetaclass:
     case Node::Kind::CanonicalSpecializedGenericTypeMetadataAccessFunction:
@@ -2915,7 +2916,6 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     Printer << ")";
     return nullptr;
   case Node::Kind::OpaqueReturnType:
-  case Node::Kind::OpaqueReturnTypeIndexed:
     Printer << "some";
     return nullptr;
   case Node::Kind::OpaqueReturnTypeOf:
@@ -3080,6 +3080,9 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
   case Node::Kind::HasSymbolQuery:
     Printer << "#_hasSymbol query for ";
     return nullptr;
+  case Node::Kind::OpaqueReturnTypeIndex:
+  case Node::Kind::OpaqueReturnTypeParent:
+    return nullptr;
   }
 
   printer_unreachable("bad node kind!");

lib/Demangling/OldDemangler.cpp

@@ -2089,7 +2089,10 @@ class OldDemangler {
         Node::IndexType ordinal;
         if (!demangleIndex(ordinal, depth))
           return nullptr;
-        return Factory.createNode(Node::Kind::OpaqueReturnTypeIndexed, ordinal);
+        auto result = Factory.createNode(Node::Kind::OpaqueReturnType);
+        result->addChild(
+          Factory.createNode(Node::Kind::OpaqueReturnTypeIndex, ordinal), Factory);
+        return result;
       }
       return demangleArchetypeType(depth + 1);
     }

lib/Demangling/OldRemangler.cpp

@@ -2674,13 +2674,21 @@ ManglingError Remangler::mangleSugaredParen(Node *node, unsigned depth) {
 }
 
 ManglingError Remangler::mangleOpaqueReturnType(Node *node, unsigned depth) {
+  if (node->hasChildren()
+      && node->getFirstChild()->getKind() == Node::Kind::OpaqueReturnTypeIndex){
+    Buffer << "QU";
+    mangleIndex(node->getFirstChild()->getIndex());
+    return ManglingError::Success;
+  }
+
   Buffer << "Qu";
   return ManglingError::Success;
 }
-ManglingError Remangler::mangleOpaqueReturnTypeIndexed(Node *node, unsigned depth) {
-  Buffer << "QU";
-  mangleIndex(node->getIndex());
-  return ManglingError::Success;
+ManglingError Remangler::mangleOpaqueReturnTypeIndex(Node *node, unsigned depth) {
+  return ManglingError::WrongNodeType;
+}
+ManglingError Remangler::mangleOpaqueReturnTypeParent(Node *node, unsigned depth) {
+  return ManglingError::WrongNodeType;
 }
 ManglingError Remangler::mangleOpaqueReturnTypeOf(Node *node,
                                                   EntityContext &ctx,

lib/Demangling/Remangler.cpp

@@ -3387,13 +3387,22 @@ ManglingError Remangler::mangleSugaredParen(Node *node, unsigned depth) {
 }
 
 ManglingError Remangler::mangleOpaqueReturnType(Node *node, unsigned depth) {
+  if (node->hasChildren()
+      && node->getFirstChild()->getKind() == Node::Kind::OpaqueReturnTypeIndex) {
+    Buffer << "QR";
+    mangleIndex(node->getFirstChild()->getIndex());
+    return ManglingError::Success;
+  }
   Buffer << "Qr";
   return ManglingError::Success;
 }
-ManglingError Remangler::mangleOpaqueReturnTypeIndexed(Node *node, unsigned depth) {
-  Buffer << "QR";
-  mangleIndex(node->getIndex());
-  return ManglingError::Success;
+ManglingError Remangler::mangleOpaqueReturnTypeIndex(Node *node, unsigned depth) {
+  // Cannot appear unparented to an OpaqueReturnType.
+  return ManglingError::WrongNodeType;
+}
+ManglingError Remangler::mangleOpaqueReturnTypeParent(Node *node, unsigned depth) {
+  // Cannot appear unparented to an OpaqueReturnType.
+  return ManglingError::WrongNodeType;
 }
 ManglingError Remangler::mangleOpaqueReturnTypeOf(Node *node, unsigned depth) {
   RETURN_IF_ERROR(mangle(node->getChild(0), depth + 1));