Merge pull request #39496 from apple/es-malias2

[NFC] Add a module alias map and a lookup func to ASTContext.
Add a getter for the actual module name in ModuleDecl if a module alias is used.
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 real (physical) names of 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,16 @@ class ASTContext final {
   /// specified string.
   Identifier getIdentifier(StringRef Str) const;
 
+  /// Convert a given alias map to a map of Identifiers between module aliases and their actual names.
+  /// For example, if '-module-alias Foo=X -module-alias Bar=Y' input is passed in, the aliases Foo and Bar are
+  /// the names of the imported or referenced modules in source files in the main module, and X and Y
+  /// are the real (physical) module names on disk.
+  void setModuleAliases(const llvm::StringMap<StringRef> &aliasMap);
+
+  /// Retrieve the actual module name if a module alias is used via '-module-alias Foo=X', where Foo is
+  /// a module alias and X is the real (physical) name. Returns \p key if no aliasing is used.
+  Identifier getRealModuleName(Identifier key) const;
+
   /// Decide how to interpret two precedence groups.
   Associativity associateInfixOperators(PrecedenceGroupDecl *left,
                                         PrecedenceGroupDecl *right) const;

include/swift/AST/Module.h

@@ -357,6 +357,15 @@ class ModuleDecl
     ModuleABIName = name;
   }
 
+  /// Retrieve the actual module name of an alias used for this module (if any).
+  ///
+  /// For example, if '-module-alias Foo=Bar' is passed in when building the main module,
+  /// and this module is (a) not the main module and (b) is named Foo, then it returns
+  /// the real (physically on-disk) module name Bar.
+  ///
+  /// If no module aliasing is set, it will return getName(), i.e. Foo.
+  Identifier getRealName() const;
+
   /// User-defined module version number.
   llvm::VersionTuple UserModuleVersion;
   void setUserModuleVersion(llvm::VersionTuple UserVer) {
@@ -365,6 +374,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.

lib/AST/ASTContext.cpp

@@ -1642,6 +1642,23 @@ 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);
+    if (!val.empty()) {
+      ModuleAliasMap[getIdentifier(k)] = getIdentifier(val);
+    }
+  }
+}
+
+Identifier ASTContext::getRealModuleName(Identifier key) const {
+  auto found = ModuleAliasMap.find(key);
+  if (found != ModuleAliasMap.end()) {
+    return found->second;
+  }
+  return key;
+}
+
 Optional<ModuleDependencies> ASTContext::getModuleDependencies(
     StringRef moduleName, bool isUnderlyingClangModule,
     ModuleDependenciesCache &cache, InterfaceSubContextDelegate &delegate,

lib/AST/Module.cpp

@@ -475,7 +475,6 @@ ModuleDecl::ModuleDecl(Identifier name, ASTContext &ctx,
   ctx.addDestructorCleanup(*this);
   setImplicit();
   setInterfaceType(ModuleType::get(this));
-
   setAccess(AccessLevel::Public);
 
   Bits.ModuleDecl.StaticLibrary = 0;
@@ -1564,6 +1563,11 @@ ImportedModule::removeDuplicates(SmallVectorImpl<ImportedModule> &imports) {
   imports.erase(last, imports.end());
 }
 
+Identifier ModuleDecl::getRealName() const {
+  // This will return the real name for an alias (if used) or getName()
+  return getASTContext().getRealModuleName(getName());
+}
+
 Identifier ModuleDecl::getABIName() const {
   if (!ModuleABIName.empty())
     return ModuleABIName;

lib/Frontend/Frontend.cpp

@@ -233,6 +233,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);
   registerClangImporterRequestFunctions(Context->evaluator);