Add module alias map and lookup func to ASTContext

Set module underlying name in ModuleDecl if a module alias exists
rdar://83682112

Author: elsh

include/swift/AST/ASTContext.h

@@ -347,6 +347,9 @@ class ASTContext final {
   /// Cache of module names that fail the 'canImport' test in this context.
   mutable llvm::SmallPtrSet<Identifier, 8> FailedModuleImportNames;
 
+  /// Mapping between aliases and underlying names of modules imported or referenced modules.
+  mutable llvm::DenseMap<Identifier, Identifier> ModuleAliasMap;
+
   /// Retrieve the allocator for the given arena.
   llvm::BumpPtrAllocator &
   getAllocator(AllocationArena arena = AllocationArena::Permanent) const;
@@ -471,6 +474,12 @@ class ASTContext final {
   /// specified string.
   Identifier getIdentifier(StringRef Str) const;
 
+  /// Convert a given alias map to a map of Identifiers between module aliases and underlying names.
+  void setModuleAliases(const llvm::StringMap<StringRef> &aliasMap);
+
+  /// Retrieve the underlying name given an alias name key.
+  Identifier lookupModuleAlias(Identifier key) const;
+
   /// Decide how to interpret two precedence groups.
   Associativity associateInfixOperators(PrecedenceGroupDecl *left,
                                         PrecedenceGroupDecl *right) const;

include/swift/AST/Module.h

@@ -171,6 +171,9 @@ class ModuleDecl
   /// The ABI name of the module, if it differs from the module name.
   mutable Identifier ModuleABIName;
 
+  /// The underlying name for an alias used for this module (if any).
+  mutable Identifier ModuleUnderlyingName;
+
 public:
   /// Produces the components of a given module's full name in reverse order.
   ///
@@ -357,6 +360,10 @@ class ModuleDecl
     ModuleABIName = name;
   }
 
+  /// Retrieve the underlying name of the alias (if any) used for this module.
+  /// If no module alias is set, it returns getName().
+  Identifier getUnderlyingName() const;
+
   /// User-defined module version number.
   llvm::VersionTuple UserModuleVersion;
   void setUserModuleVersion(llvm::VersionTuple UserVer) {
@@ -365,6 +372,7 @@ class ModuleDecl
   llvm::VersionTuple getUserModuleVersion() const {
     return UserModuleVersion;
   }
+
 private:
   /// A cache of this module's underlying module and required bystander if it's
   /// an underscored cross-import overlay.
@@ -380,6 +388,12 @@ class ModuleDecl
   ///  module if one exists.
   ModuleDecl *getUnderlyingModuleIfOverlay() const;
 
+  /// If a module alias is used, set the corresponding underlying name,
+  /// which will be used for contents including metadata and mangling.
+  void setUnderlyingName(Identifier name) {
+    ModuleUnderlyingName = name;
+  }
+
 public:
 
   /// Returns true if this module is an underscored cross import overlay

lib/AST/ASTContext.cpp

@@ -1642,6 +1642,21 @@ void ASTContext::addModuleInterfaceChecker(
   getImpl().InterfaceChecker = std::move(checker);
 }
 
+void ASTContext::setModuleAliases(const llvm::StringMap<StringRef> &aliasMap) {
+  for (auto k: aliasMap.keys()) {
+    auto val = aliasMap.lookup(k);
+    ModuleAliasMap[getIdentifier(k)] = getIdentifier(val);
+  }
+}
+
+Identifier ASTContext::lookupModuleAlias(Identifier key) const {
+  auto found = ModuleAliasMap.find(key);
+  if (found != ModuleAliasMap.end()) {
+    return found->second;
+  }
+  return Identifier();
+}
+
 Optional<ModuleDependencies> ASTContext::getModuleDependencies(
     StringRef moduleName, bool isUnderlyingClangModule,
     ModuleDependenciesCache &cache, InterfaceSubContextDelegate &delegate,

lib/AST/Module.cpp

@@ -475,8 +475,8 @@ ModuleDecl::ModuleDecl(Identifier name, ASTContext &ctx,
   ctx.addDestructorCleanup(*this);
   setImplicit();
   setInterfaceType(ModuleType::get(this));
-
   setAccess(AccessLevel::Public);
+  setUnderlyingName(ctx.lookupModuleAlias(name));
 
   Bits.ModuleDecl.StaticLibrary = 0;
   Bits.ModuleDecl.TestingEnabled = 0;
@@ -1564,6 +1564,12 @@ ImportedModule::removeDuplicates(SmallVectorImpl<ImportedModule> &imports) {
   imports.erase(last, imports.end());
 }
 
+Identifier ModuleDecl::getUnderlyingName() const {
+    if (!ModuleUnderlyingName.empty())
+        return ModuleUnderlyingName;
+    return getName();
+}
+
 Identifier ModuleDecl::getABIName() const {
   if (!ModuleABIName.empty())
     return ModuleABIName;

lib/Frontend/Frontend.cpp

@@ -219,6 +219,9 @@ bool CompilerInstance::setUpASTContextIfNeeded() {
       Invocation.getClangImporterOptions(),
       Invocation.getSymbolGraphOptions(),
       SourceMgr, Diagnostics));
+  if (!Invocation.getFrontendOptions().ModuleAliasMap.empty())
+    Context->setModuleAliases(Invocation.getFrontendOptions().ModuleAliasMap);
+
   registerParseRequestFunctions(Context->evaluator);
   registerTypeCheckerRequestFunctions(Context->evaluator);
   registerSILGenRequestFunctions(Context->evaluator);