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,16 @@ class ASTContext final {
 
   ModuleDecl *getModuleByIdentifier(Identifier ModuleID);
 
-  /// Looks up an already loaded module by its ABI name.
+  /// Looks up all modules whose real name or ABI name match ModuleName.
   ///
-  /// \returns The module if found, nullptr otherwise.
-  ModuleDecl *getLoadedModuleByABIName(StringRef ModuleName);
+  /// Modules that are being looked up by ABI name must be a dependency
+  /// of a module that has that same name as its real name, as there is no
+  /// efficient way to lazily load a module by ABI name.
+  llvm::SmallVector<ModuleDecl *, 1> *
+  getModulesByRealOrABIName(StringRef ModuleName);
+
+  /// Notifies the AST context that a loaded module's ABI name will change.
+  void moduleABINameWillChange(ModuleDecl *module, Identifier newName);
 
   /// 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 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 {

include/swift/AST/Module.h

@@ -473,9 +473,7 @@ class ModuleDecl
   Identifier getABIName() const;
 
   /// Set the ABI name of the module;
-  void setABIName(Identifier name) {
-    ModuleABIName = name;
-  }
+  void setABIName(Identifier name);
 
   /// Get the package name of this module
   /// FIXME: remove this and bump module version rdar://104723918

lib/AST/ASTContext.cpp

@@ -275,6 +275,10 @@ struct ASTContext::Implementation {
   /// DenseMap.
   llvm::MapVector<Identifier, ModuleDecl *> LoadedModules;
 
+  /// The map from a module's name to a vector of modules that share that name.
+  /// The name can be either the module's real name of the module's ABI name.
+  llvm::DenseMap<Identifier, llvm::SmallVector<ModuleDecl *, 1>> NameToModules;
+
   // FIXME: This is a StringMap rather than a StringSet because StringSet
   // doesn't allow passing in a pre-existing allocator.
   llvm::StringMap<Identifier::Aligner, llvm::BumpPtrAllocator&>
@@ -836,6 +840,7 @@ void ASTContext::Implementation::dump(llvm::raw_ostream &os) const {
                                 << llvm::capacity_in_bytes(Name) << "\n"
 
   SIZE(LoadedModules);
+  SIZE(NameToModules);
   SIZE(IdentifierTable);
   SIZE(Cleanups);
   SIZE_AND_BYTES(ModuleLoaders);
@@ -2321,12 +2326,63 @@ void ASTContext::addLoadedModule(ModuleDecl *M) {
   // and a source file has 'import Foo', a module called Bar (real name)
   // will be loaded and added to the map.
   getImpl().LoadedModules[M->getRealName()] = M;
+
+  // Add the module to the mapping from module name to list of modules that
+  // share that name.
+  getImpl().NameToModules[M->getRealName()].push_back(M);
+
+  // If the ABI name differs from the real name, also add the module to the list
+  // that share that ABI name.
+  if (M->getRealName() != M->getABIName())
+    getImpl().NameToModules[M->getABIName()].push_back(M);
 }
 
 void ASTContext::removeLoadedModule(Identifier RealName) {
+  // First remove the module from the mappings of names to modules.
+  if (ModuleDecl *M = getLoadedModule(RealName)) {
+    auto eraseModule = [&](ModuleDecl *module) {
+      return module->getRealName() == RealName;
+    };
+    auto &vector = getImpl().NameToModules[M->getRealName()];
+    llvm::erase_if(vector, eraseModule);
+    if (M->getRealName() != M->getABIName()) {
+      auto &vector = getImpl().NameToModules[M->getABIName()];
+      llvm::erase_if(vector, eraseModule);
+    }
+  }
+
   getImpl().LoadedModules.erase(RealName);
 }
 
+void ASTContext::moduleABINameWillChange(ModuleDecl *module,
+                                         Identifier newName) {
+  auto it = llvm::find_if(getLoadedModules(),
+                          [&](auto pair) { return pair.second == module; });
+
+  // If this module isn't in the loaded modules list (perhaps because there is
+  // no memory cache) theere's nothing to do.
+  if (it == getLoadedModules().end())
+    return;
+
+  // If the names are the same there's nothing to do.
+  if (module->getABIName() == newName)
+    return;
+
+  // If the real and ABI names are different, ASTContext needs to remove the
+  // module from the mapping whose key is the old ABI name.
+  if (module->getRealName() != module->getABIName()) {
+    auto &vector = getImpl().NameToModules[module->getABIName()];
+    llvm::erase_if(vector,
+                   [&](ModuleDecl *current) { return module == current; });
+  }
+
+  // Now add the module to the vector that's mapped from the new name, if it's
+  // not there already.
+  auto &vector = getImpl().NameToModules[newName];
+  if (llvm::find(vector, module) == vector.end())
+    vector.push_back(module);
+}
+
 void ASTContext::setIgnoreAdjacentModules(bool value) {
   IgnoreAdjacentModules = value;
 }
@@ -2713,11 +2769,19 @@ 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;
-  }
+llvm::SmallVector<ModuleDecl *, 1> *
+ASTContext::getModulesByRealOrABIName(StringRef ModuleName) {
+  auto Identifier = getIdentifier(ModuleName);
+  auto it = getImpl().NameToModules.find(Identifier);
+  if (it != getImpl().NameToModules.end())
+    return &it->second;
+
+  // If we didn't find the module it might have not been loaded yet, try
+  // triggering a module load and searching again.
+  getModuleByName(ModuleName);
+  it = getImpl().NameToModules.find(Identifier);
+  if (it != getImpl().NameToModules.end())
+    return &it->second;
   return nullptr;
 }
 

lib/AST/ASTDemangler.cpp

@@ -264,11 +264,15 @@ 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 || potentialParentModules->empty())
       return Type();
 
-    auto opaqueDecl = parentModule->lookupOpaqueResultType(mangledName);
+    OpaqueTypeDecl *opaqueDecl = nullptr;
+    for (auto module : *potentialParentModules)
+      if (auto decl = module->lookupOpaqueResultType(mangledName))
+        opaqueDecl = decl;
+
     if (!opaqueDecl)
       return Type();
     SmallVector<Type, 8> allArgs;
@@ -1123,17 +1127,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 +1221,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 || modules->empty()) ? nullptr : (*modules)[0];
+  }
 
   case Demangle::Node::Kind::Class:
   case Demangle::Node::Kind::Enum:
@@ -1237,20 +1244,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 || 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;
@@ -1284,22 +1295,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 || 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 || moduleDecls->empty())
       return nullptr;
 
     auto *nominalDecl = dyn_cast_or_null<NominalTypeDecl>(
@@ -1321,14 +1343,23 @@ 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)
+      bool found = false;
+      for (ModuleDecl *module : *moduleDecls) {
+        if (ext->getParentModule() == module) {
+          found = true;
+          break;
+        }
+      }
+      if (!found)
         continue;
 
       if (!ext->isConstrainedExtension()) {

lib/AST/Module.cpp

@@ -1954,6 +1954,11 @@ Identifier ModuleDecl::getABIName() const {
   return getName();
 }
 
+void ModuleDecl::setABIName(Identifier name) {
+  getASTContext().moduleABINameWillChange(this, name);
+  ModuleABIName = name;
+}
+
 StringRef ModuleDecl::getModuleFilename() const {
   // FIXME: Audit uses of this function and figure out how to migrate them to
   // per-file names. Modules can consist of more than one file.

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>>>