Add module alias lookup option enum

Pass the enum param to ASTContext::getRealModuleName

Author: elsh

include/swift/AST/ASTContext.h

@@ -483,9 +483,27 @@ class ASTContext final {
   /// 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;
+  /// Look up option used in \c getRealModuleName when module aliasing is applied.
+  enum class ModuleAliasLookupOption {
+    alwaysRealName,
+    realNameFromAlias,
+    aliasFromRealName
+  };
+
+  /// Look up the module alias map by the given \p key and a lookup \p option.
+  ///
+  /// \param key A module alias or real name to look up the map by.
+  /// \param option A look up option \c ModuleAliasLookupOption. Defaults to alwaysRealName.
+  ///
+  /// \return The real name or alias mapped to the key.
+  ///         If no aliasing is used, return \p key regardless of \p option.
+  ///         If \p option is alwaysRealName, return the real module name whether the \p key is an alias
+  ///         or a real name.
+  ///         If \p option is realNameFromAlias, only return a real name if \p key is an alias.
+  ///         If \p option is aliasFromRealName, only return an alias if \p key is a real name.
+  ///         Else return a real name or an alias mapped to the \p key.
+  Identifier getRealModuleName(Identifier key,
+                               ModuleAliasLookupOption option = ModuleAliasLookupOption::alwaysRealName) const;
 
   /// Decide how to interpret two precedence groups.
   Associativity associateInfixOperators(PrecedenceGroupDecl *left,

lib/AST/ASTContext.cpp

@@ -1645,12 +1645,31 @@ void ASTContext::setModuleAliases(const llvm::StringMap<StringRef> &aliasMap) {
   }
 }
 
-Identifier ASTContext::getRealModuleName(Identifier key) const {
+Identifier ASTContext::getRealModuleName(Identifier key, ModuleAliasLookupOption option) const {
   auto found = ModuleAliasMap.find(key);
-  if (found != ModuleAliasMap.end()) {
-    return found->second;
+  if (found == ModuleAliasMap.end())
+    return key; // No module aliasing was used, so just return the given key
+
+  // Found an entry
+  auto value = found->second;
+
+  // With the alwaysRealName option, look up the real name by treating
+  // the given key as an alias; if the key's not an alias, return the key
+  // itself since that's the real name.
+  if (option == ModuleAliasLookupOption::alwaysRealName) {
+     return value.second ? value.first : key;
   }
-  return key;
+
+  // With realNameFromAlias or aliasFromRealName option, only return the value
+  // if the given key matches the description (whether it's an alias or real name)
+  // by looking up the value.second (true if keyed by an alias). If not matched,
+  // return an empty Identifier.
+  if ((option == ModuleAliasLookupOption::realNameFromAlias && !value.second) ||
+      (option == ModuleAliasLookupOption::aliasFromRealName && value.second))
+      return Identifier();
+
+  // Otherwise return the value found (whether the key is an alias or real name)
+  return value.first;
 }
 
 Optional<ModuleDependencies> ASTContext::getModuleDependencies(

lib/AST/UnqualifiedLookup.cpp

@@ -431,7 +431,21 @@ void UnqualifiedLookupFactory::lookForAModuleWithTheGivenName(
 #endif
     return;
   }
-  ModuleDecl *desiredModule = Ctx.getLoadedModule(Name.getBaseIdentifier());
+
+  ModuleDecl *desiredModule = nullptr;
+  auto givenName = Name.getBaseIdentifier();
+  // Check if the given name appearing in the source file is a module
+  // real name or alias; for example, if `-module-alias Foo=Bar` was
+  // passed, the alias 'Foo' should appear in source files, not 'Bar'.
+  // If the real name 'Bar' was used, looking up getRealModuleName with
+  // the real name 'Bar' and realNameFromAlias option should return
+  // an empty Identifier.
+  if (!Ctx.getRealModuleName(givenName, ASTContext::ModuleAliasLookupOption::realNameFromAlias).empty()) {
+    // Only load the module if the lookup value is not empty, i.e. given
+    // name is a module alias, not a real module name.
+    desiredModule = Ctx.getLoadedModule(givenName);
+  }
+
   if (!desiredModule && Name.getFullName() == Ctx.TheBuiltinModule->getName())
     desiredModule = Ctx.TheBuiltinModule;
   if (desiredModule) {

lib/Parse/ParseDecl.cpp

@@ -4819,6 +4819,20 @@ ParserResult<ImportDecl> Parser::parseDeclImport(ParseDeclOptions Flags,
     return nullptr;
   }
 
+  // Look up if the imported module is being aliased via -module-alias,
+  // and check that the module alias appeared in source files instead of
+  // its corresponding real name
+  auto parsedModuleID = importPath.get().front().Item;
+  if (Context.getRealModuleName(parsedModuleID, ASTContext::ModuleAliasLookupOption::realNameFromAlias).empty()) {
+    // If reached here, it means the parsed module name is a real module name
+    // which appeared in the source file; only a module alias should be allowed
+    auto aliasName = Context.getRealModuleName(parsedModuleID, ASTContext::ModuleAliasLookupOption::aliasFromRealName);
+    diagnose(importPath.front().Loc, diag::expected_module_alias,
+                     parsedModuleID, aliasName)
+      .fixItReplace(importPath.front().Loc, aliasName.str());
+    return nullptr;
+  }
+
   auto *ID = ImportDecl::create(Context, CurDeclContext, ImportLoc, Kind,
                                 KindLoc, importPath.get());
   ID->getAttrs() = Attributes;