Account for multiple modules when looking up the DeclContext of a type

When looking up the decl context of a type, ASTDemangler has to take
into account that there are multiple different modules where that type
could've come from. This is due to two facts:

- Thanks to the `-module-abi-name` flag, multiple modules can share the
  same ABI name (which is the module name that is usually used when
  mangling a type).
- In some situations mangling can use the module's real name, for
  example, when mangling for the debugger or USRs coupled with
  @_originallyDefinedIn.

rdar://134095412

Author: augusto2112

include/swift/AST/ASTContext.h

@@ -1200,10 +1200,9 @@ class ASTContext final {
 
   ModuleDecl *getModuleByIdentifier(Identifier ModuleID);
 
-  /// Looks up an already loaded module by its ABI name.
-  ///
-  /// \returns The module if found, nullptr otherwise.
-  ModuleDecl *getLoadedModuleByABIName(StringRef ModuleName);
+  /// Looks up all modules whose real name or ABI name match ModuleName.
+  llvm::SmallVector<ModuleDecl *, 1>
+  getModulesByRealOrABIName(StringRef ModuleName);
 
   /// Returns the standard library module, or null if the library isn't present.
   ///

include/swift/AST/ASTDemangler.h

@@ -254,7 +254,13 @@ class ASTBuilder {
       NominalTypeDecl *nominalDecl,
       NodePointer node);
   DeclContext *findDeclContext(NodePointer node);
-  ModuleDecl *findModule(NodePointer node);
+
+  /// Find all the ModuleDecls that correspond to a module node's identifier.
+  /// The module name encoded in the node is either the module's real or ABI
+  /// name. Multiple modules can share the same ABI name. This function returns
+  /// all modules that contain that name.
+  llvm::SmallVector<ModuleDecl *, 1> findPotentialModules(NodePointer node);
+
   Demangle::NodePointer findModuleNode(NodePointer node);
 
   enum class ForeignModuleKind {

lib/AST/ASTContext.cpp

@@ -2713,12 +2713,16 @@ ModuleDecl *ASTContext::getModuleByIdentifier(Identifier ModuleID) {
   return getModule(builder.get());
 }
 
-ModuleDecl *ASTContext::getLoadedModuleByABIName(StringRef ModuleName) {
-  for (auto &[_, module] : getLoadedModules()) {
-    if (ModuleName == module->getABIName().str())
-      return module;
-  }
-  return nullptr;
+llvm::SmallVector<ModuleDecl *, 1>
+ASTContext::getModulesByRealOrABIName(StringRef ModuleName) {
+  llvm::SmallVector<ModuleDecl *, 1> modules;
+
+  for (auto &[_, module] : getLoadedModules())
+    if (ModuleName == module->getRealName().str() ||
+        ModuleName == module->getABIName().str())
+      modules.push_back(module);
+
+  return modules;
 }
 
 ModuleDecl *ASTContext::getStdlibModule(bool loadIfAbsent) {

lib/AST/ASTDemangler.cpp

@@ -264,25 +264,26 @@ Type ASTBuilder::resolveOpaqueType(NodePointer opaqueDescriptor,
     auto moduleNode = findModuleNode(definingDecl);
     if (!moduleNode)
       return Type();
-    auto parentModule = findModule(moduleNode);
-    if (!parentModule)
+    auto potentialParentModules = findPotentialModules(moduleNode);
+    if (potentialParentModules.empty())
       return Type();
 
-    auto opaqueDecl = parentModule->lookupOpaqueResultType(mangledName);
-    if (!opaqueDecl)
-      return Type();
-    SmallVector<Type, 8> allArgs;
-    for (auto argSet : args) {
-      allArgs.append(argSet.begin(), argSet.end());
-    }
+    for (ModuleDecl *module : potentialParentModules) {
+      auto opaqueDecl = module->lookupOpaqueResultType(mangledName);
+      if (!opaqueDecl)
+        continue;
+      SmallVector<Type, 8> allArgs;
+      for (auto argSet : args) {
+        allArgs.append(argSet.begin(), argSet.end());
+      }
 
-    SubstitutionMap subs = createSubstitutionMapFromGenericArgs(
-        opaqueDecl->getGenericSignature(), allArgs,
-        LookUpConformanceInModule());
-    Type interfaceType = opaqueDecl->getOpaqueGenericParams()[ordinal];
-    return OpaqueTypeArchetypeType::get(opaqueDecl, interfaceType, subs);
+      SubstitutionMap subs = createSubstitutionMapFromGenericArgs(
+          opaqueDecl->getGenericSignature(), allArgs,
+          LookUpConformanceInModule());
+      Type interfaceType = opaqueDecl->getOpaqueGenericParams()[ordinal];
+      return OpaqueTypeArchetypeType::get(opaqueDecl, interfaceType, subs);
+    }
   }
-  
   // TODO: named opaque types
   return Type();
 }
@@ -1123,17 +1124,11 @@ ASTBuilder::createTypeDecl(NodePointer node,
   return dyn_cast<GenericTypeDecl>(DC);
 }
 
-ModuleDecl *ASTBuilder::findModule(NodePointer node) {
+llvm::SmallVector<ModuleDecl *, 1>
+ASTBuilder::findPotentialModules(NodePointer node) {
   assert(node->getKind() == Demangle::Node::Kind::Module);
   const auto moduleName = node->getText();
-  // Respect the module's ABI name when we're trying to resolve
-  // mangled names. But don't touch anything under the Swift stdlib's
-  // umbrella.
-  if (moduleName != STDLIB_NAME)
-    if (auto *Module = Ctx.getLoadedModuleByABIName(moduleName))
-      return Module;
-
-  return Ctx.getModuleByName(moduleName);
+  return Ctx.getModulesByRealOrABIName(moduleName);
 }
 
 Demangle::NodePointer
@@ -1223,8 +1218,17 @@ ASTBuilder::findDeclContext(NodePointer node) {
   case Demangle::Node::Kind::BoundGenericTypeAlias:
     return findDeclContext(node->getFirstChild());
 
-  case Demangle::Node::Kind::Module:
-    return findModule(node);
+  case Demangle::Node::Kind::Module: {
+    // A Module node is not enough information to find the decl context.
+    // The reason being that the module name in a mangled name can either be
+    // the module's ABI name, which is potentially not unique (due to the
+    // -module-abi-name flag), or the module's real name, if mangling for the
+    // debugger or USR together with the OriginallyDefinedIn attribute for
+    // example.
+    assert(false && "Looked up module as decl context directly!");
+    auto modules = findPotentialModules(node);
+    return modules.empty() ? nullptr : modules[0];
+  }
 
   case Demangle::Node::Kind::Class:
   case Demangle::Node::Kind::Enum:
@@ -1237,20 +1241,24 @@ ASTBuilder::findDeclContext(NodePointer node) {
     if (declNameNode->getKind() == Demangle::Node::Kind::LocalDeclName) {
       // Find the AST node for the defining module.
       auto moduleNode = findModuleNode(node);
-      if (!moduleNode) return nullptr;
+      if (!moduleNode)
+        return nullptr;
 
-      auto module = findModule(moduleNode);
-      if (!module) return nullptr;
+      auto potentialModules = findPotentialModules(moduleNode);
+      if (potentialModules.empty())
+        return nullptr;
 
       // Look up the local type by its mangling.
       auto mangling = Demangle::mangleNode(node);
-      if (!mangling.isSuccess()) return nullptr;
+      if (!mangling.isSuccess())
+        return nullptr;
       auto mangledName = mangling.result();
 
-      auto decl = module->lookupLocalType(mangledName);
-      if (!decl) return nullptr;
+      for (auto *module : potentialModules)
+        if (auto *decl = module->lookupLocalType(mangledName))
+          return dyn_cast<DeclContext>(decl);
 
-      return dyn_cast<DeclContext>(decl);
+      return nullptr;
     }
 
     StringRef name;
@@ -1272,7 +1280,7 @@ ASTBuilder::findDeclContext(NodePointer node) {
 
     // Ignore any other decl-name productions for now.
     } else {
-      return nullptr;
+      return nullptr;;
     }
 
     // Do some special logic for foreign type declarations.
@@ -1284,22 +1292,33 @@ ASTBuilder::findDeclContext(NodePointer node) {
       }
     }
 
-    DeclContext *dc = findDeclContext(node->getChild(0));
-    if (!dc) {
+    auto child = node->getFirstChild();
+    if (child->getKind() == Node::Kind::Module) {
+      auto potentialModules = findPotentialModules(child);
+      if (potentialModules.empty())
+        return nullptr;
+
+      for (auto *module : potentialModules)
+        if (auto typeDecl = findTypeDecl(module, Ctx.getIdentifier(name),
+                                         privateDiscriminator, node->getKind()))
+          return typeDecl;
       return nullptr;
     }
 
-    return findTypeDecl(dc, Ctx.getIdentifier(name),
-                        privateDiscriminator, node->getKind());
+    if (auto *dc = findDeclContext(child))
+      if (auto typeDecl = findTypeDecl(dc, Ctx.getIdentifier(name),
+                                       privateDiscriminator, node->getKind()))
+        return typeDecl;
+
+    return nullptr;
   }
 
   case Demangle::Node::Kind::Global:
     return findDeclContext(node->getChild(0));
 
   case Demangle::Node::Kind::Extension: {
-    auto *moduleDecl = dyn_cast_or_null<ModuleDecl>(
-        findDeclContext(node->getChild(0)));
-    if (!moduleDecl)
+    auto moduleDecls = findPotentialModules(node->getFirstChild());
+    if (moduleDecls.empty())
       return nullptr;
 
     auto *nominalDecl = dyn_cast_or_null<NominalTypeDecl>(
@@ -1321,39 +1340,43 @@ ASTBuilder::findDeclContext(NodePointer node) {
       // If the generic signature is equivalent to that of the nominal type,
       // and we're in the same module, it's due to inverse requirements.
       // Just return the nominal declaration.
-      if (genericSigMatchesNominal &&
-          nominalDecl->getParentModule() == moduleDecl) {
-        return nominalDecl;
+      for (auto *moduleDecl : moduleDecls) {
+        if (genericSigMatchesNominal &&
+            nominalDecl->getParentModule() == moduleDecl) {
+          return nominalDecl;
+        }
       }
     }
 
-    for (auto *ext : nominalDecl->getExtensions()) {
-      if (ext->getParentModule() != moduleDecl)
-        continue;
-
-      if (!ext->isConstrainedExtension()) {
-        if (!genericSig || genericSigMatchesNominal)
-          return ext;
+    for (auto *moduleDecl : moduleDecls) {
+      for (auto *ext : nominalDecl->getExtensions()) {
+        if (ext->getParentModule() != moduleDecl)
+          continue;
 
-        continue;
-      }
+        if (!ext->isConstrainedExtension()) {
+          if (!genericSig || genericSigMatchesNominal)
+            return ext;
 
-      if (!ext->isWrittenWithConstraints() && !genericSig)
-        return ext;
+          continue;
+        }
 
-      auto extSig = ext->getGenericSignature().getCanonicalSignature();
-      if (extSig == genericSig) {
-        return ext;
-      }
+        if (!ext->isWrittenWithConstraints() && !genericSig)
+          return ext;
 
-      // If the extension mangling doesn't include a generic signature, it
-      // might be because the nominal type suppresses conformance.
-      if (!genericSig) {
-        SmallVector<Requirement, 2> requirements;
-        SmallVector<InverseRequirement, 2> inverses;
-        extSig->getRequirementsWithInverses(requirements, inverses);
-        if (requirements.empty())
+        auto extSig = ext->getGenericSignature().getCanonicalSignature();
+        if (extSig == genericSig) {
           return ext;
+        }
+
+        // If the extension mangling doesn't include a generic signature, it
+        // might be because the nominal type suppresses conformance.
+        if (!genericSig) {
+          SmallVector<Requirement, 2> requirements;
+          SmallVector<InverseRequirement, 2> inverses;
+          extSig->getRequirementsWithInverses(requirements, inverses);
+          if (requirements.empty())
+            return ext;
+        }
       }
     }
 

test/TypeDecoder/clashing_abi_name.swift

@@ -0,0 +1,65 @@
+// Tests that reconstructing from mangled names whose types are defined in modules
+// with clashing ABI names works.
+
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+// RUN: cd %t
+
+// RUN: %target-build-swift -emit-library -emit-module -parse-as-library -module-name Foo -module-abi-name Bar -g %t/Foo.swift 
+// RUN: %target-build-swift -emit-executable -I %t -L %t -lFoo -g -emit-module -module-name Bar -module-abi-name Bar %t/Bar.swift
+
+
+// RUN: sed -ne '/\/\/ *DEMANGLE-TYPE: /s/\/\/ *DEMANGLE-TYPE: *//p' < %s > %t/input
+// RUN: %lldb-moduleimport-test-with-sdk %t/Bar -qualify-types=1 -type-from-mangled=%t/input | %FileCheck %s --check-prefix=CHECK-TYPE
+
+//--- Foo.swift
+public class One {
+  let i = 42
+  public init() {
+  }
+}
+
+public class Generic<T> {
+  let t: T
+  public init(t: T) {
+    self.t = t
+  }
+}
+
+//--- Bar.swift
+import Foo
+
+public class Two {
+  let j = 98
+  public init() {
+  }
+}
+
+public class Generic2<T> {
+  let t: T
+  public init(t: T) {
+    self.t = t
+  }
+}
+
+
+let one = Foo.One()
+let two = Bar.Two()
+let generic1 = Foo.Generic<Bar.Two>(t: two)
+let generic2 = Bar.Generic2<Foo.Generic<Bar.Two>>(t: generic1)
+let generic3 = Foo.Generic<Bar.Generic2<Foo.Generic<Bar.Two>>>(t: generic2)
+
+// DEMANGLE-TYPE: $s3Bar3OneCD
+// CHECK-TYPE: Foo.One
+
+// DEMANGLE-TYPE: $s3Bar3TwoCD
+// CHECK-TYPE: Bar.Two
+
+// DEMANGLE-TYPE: $s3Bar7GenericCyAA3TwoCGD
+// CHECK-TYPE: Foo.Generic<Bar.Two>
+
+// DEMANGLE-TYPE: $s3Bar8Generic2CyAA7GenericCyAA3TwoCGGD
+// CHECK-TYPE: Bar.Generic2<Foo.Generic<Bar.Two>>
+
+// DEMANGLE-TYPE: $s3Bar7GenericCyAA8Generic2CyACyAA3TwoCGGGD
+// CHECK-TYPE: Foo.Generic<Bar.Generic2<Foo.Generic<Bar.Two>>>

test/TypeDecoder/different_abi_name.swift

@@ -7,7 +7,7 @@
 // RUN: cd %t
 
 // RUN: %target-build-swift -emit-library -emit-module -parse-as-library -module-abi-name Other -g %t/TheModule.swift 
-// RUN: %target-build-swift -emit-executable -I %t -L %t -lTheModule %s -g -o %t/user -emit-module
+// RUN: %target-build-swift -emit-executable -I %t -L %t -lTheModule -g -o %t/user -emit-module %t/user.swift
 
 
 // RUN: sed -ne '/\/\/ *DEMANGLE-TYPE: /s/\/\/ *DEMANGLE-TYPE: *//p' < %s > %t/input