[NFC] [Modules] Extract the logic to decide whether the module units belongs to the same module

This patch extracts the logci to decide how we decide the module units
belongs to the same module into a member function of ASTContext. This is
helpful to refactor the implementation in the future.

Author: ChuanqiXu9

clang/include/clang/AST/ASTContext.h

@@ -1073,6 +1073,12 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// Get module under construction, nullptr if this is not a C++20 module.
   Module *getCurrentNamedModule() const { return CurrentCXXNamedModule; }
 
+  /// If the two module \p M1 and \p M2 are in the same module.
+  ///
+  /// FIXME: The signature may be confusing since `clang::Module` means to
+  /// a module fragment or a module unit but not a C++20 module.
+  bool isInSameModule(const Module *M1, const Module *M2);
+
   TranslationUnitDecl *getTranslationUnitDecl() const {
     return TUDecl->getMostRecentDecl();
   }

clang/lib/AST/ASTContext.cpp

@@ -1110,6 +1110,15 @@ void ASTContext::setCurrentNamedModule(Module *M) {
   CurrentCXXNamedModule = M;
 }
 
+bool ASTContext::isInSameModule(const Module *M1, const Module *M2) {
+  if (!M1 != !M2)
+    return false;
+
+  assert(M1 && "Shouldn't call `isInSameModule` if both M1 and M2 are none.");
+  return M1->getPrimaryModuleInterfaceName() ==
+         M2->getPrimaryModuleInterfaceName();
+}
+
 ExternCContextDecl *ASTContext::getExternCContextDecl() const {
   if (!ExternCContext)
     ExternCContext = ExternCContextDecl::Create(*this, getTranslationUnitDecl());

clang/lib/Sema/SemaDecl.cpp

@@ -1663,8 +1663,7 @@ bool Sema::CheckRedeclarationModuleOwnership(NamedDecl *New, NamedDecl *Old) {
     // Partitions are part of the module, but a partition could import another
     // module, so verify that the PMIs agree.
     if ((NewM->isModulePartition() || OldM->isModulePartition()) &&
-        NewM->getPrimaryModuleInterfaceName() ==
-            OldM->getPrimaryModuleInterfaceName())
+        getASTContext().isInSameModule(NewM, OldM))
       return false;
   }
 

clang/lib/Sema/SemaLookup.cpp

@@ -1606,22 +1606,32 @@ bool Sema::isUsableModule(const Module *M) {
   // [module.global.frag]p1:
   //   The global module fragment can be used to provide declarations that are
   //   attached to the global module and usable within the module unit.
-  if (M == TheGlobalModuleFragment || M == TheImplicitGlobalModuleFragment ||
-      // If M is the module we're parsing, it should be usable. This covers the
-      // private module fragment. The private module fragment is usable only if
-      // it is within the current module unit. And it must be the current
-      // parsing module unit if it is within the current module unit according
-      // to the grammar of the private module fragment. NOTE: This is covered by
-      // the following condition. The intention of the check is to avoid string
-      // comparison as much as possible.
-      M == getCurrentModule() ||
-      // The module unit which is in the same module with the current module
-      // unit is usable.
-      //
-      // FIXME: Here we judge if they are in the same module by comparing the
-      // string. Is there any better solution?
-      M->getPrimaryModuleInterfaceName() ==
-          llvm::StringRef(getLangOpts().CurrentModule).split(':').first) {
+  if (M == TheGlobalModuleFragment || M == TheImplicitGlobalModuleFragment) {
+    UsableModuleUnitsCache.insert(M);
+    return true;
+  }
+
+  // Otherwise, the global module fragment from other translation unit is not
+  // directly usable.
+  if (M->isGlobalModule())
+    return false;
+
+  Module *Current = getCurrentModule();
+
+  // If we're not parsing a module, we can't use all the declarations from
+  // another module easily.
+  if (!Current)
+    return false;
+
+  // If M is the module we're parsing or M and the current module unit lives in
+  // the same module, M should be usable.
+  //
+  // Note: It should be fine to search the vector `ModuleScopes` linearly since
+  // it should be generally small enough. There should be rare module fragments
+  // in a named module unit.
+  if (llvm::count_if(ModuleScopes,
+                     [&M](const ModuleScope &MS) { return MS.Module == M; }) ||
+      getASTContext().isInSameModule(M, Current)) {
     UsableModuleUnitsCache.insert(M);
     return true;
   }
@@ -5816,19 +5826,13 @@ void Sema::diagnoseMissingImport(SourceLocation UseLoc, const NamedDecl *Decl,
     if (M->isModuleMapModule())
       return M->getFullModuleName();
 
-    Module *CurrentModule = getCurrentModule();
-
     if (M->isImplicitGlobalModule())
       M = M->getTopLevelModule();
 
-    bool IsInTheSameModule =
-        CurrentModule && CurrentModule->getPrimaryModuleInterfaceName() ==
-                             M->getPrimaryModuleInterfaceName();
-
     // If the current module unit is in the same module with M, it is OK to show
     // the partition name. Otherwise, it'll be sufficient to show the primary
     // module name.
-    if (IsInTheSameModule)
+    if (getASTContext().isInSameModule(M, getCurrentModule()))
       return M->getTopLevelModuleName().str();
     else
       return M->getPrimaryModuleInterfaceName().str();

clang/lib/Sema/SemaModule.cpp

@@ -82,7 +82,8 @@ static std::string stringFromPath(ModuleIdPath Path) {
 /// CurrentModule. Since currently it is expensive to decide whether two module
 /// units come from the same module by comparing the module name.
 static bool
-isImportingModuleUnitFromSameModule(Module *Imported, Module *CurrentModule,
+isImportingModuleUnitFromSameModule(ASTContext &Ctx, Module *Imported,
+                                    Module *CurrentModule,
                                     Module *&FoundPrimaryModuleInterface) {
   if (!Imported->isNamedModule())
     return false;
@@ -109,8 +110,7 @@ isImportingModuleUnitFromSameModule(Module *Imported, Module *CurrentModule,
   if (!CurrentModule->isModulePartitionImplementation())
     return false;
 
-  if (Imported->getPrimaryModuleInterfaceName() ==
-      CurrentModule->getPrimaryModuleInterfaceName()) {
+  if (Ctx.isInSameModule(Imported, CurrentModule)) {
     assert(!FoundPrimaryModuleInterface ||
            FoundPrimaryModuleInterface == Imported);
     FoundPrimaryModuleInterface = Imported;
@@ -127,8 +127,9 @@ isImportingModuleUnitFromSameModule(Module *Imported, Module *CurrentModule,
 ///   the module unit purview of U. These rules can in turn lead to the
 ///   importation of yet more translation units.
 static void
-makeTransitiveImportsVisible(VisibleModuleSet &VisibleModules, Module *Imported,
-                             Module *CurrentModule, SourceLocation ImportLoc,
+makeTransitiveImportsVisible(ASTContext &Ctx, VisibleModuleSet &VisibleModules,
+                             Module *Imported, Module *CurrentModule,
+                             SourceLocation ImportLoc,
                              bool IsImportingPrimaryModuleInterface = false) {
   assert(Imported->isNamedModule() &&
          "'makeTransitiveImportsVisible()' is intended for standard C++ named "
@@ -150,7 +151,7 @@ makeTransitiveImportsVisible(VisibleModuleSet &VisibleModules, Module *Imported,
     // use the sourcelocation loaded from the visible modules.
     VisibleModules.setVisible(Importing, ImportLoc);
 
-    if (isImportingModuleUnitFromSameModule(Importing, CurrentModule,
+    if (isImportingModuleUnitFromSameModule(Ctx, Importing, CurrentModule,
                                             FoundPrimaryModuleInterface))
       for (Module *TransImported : Importing->Imports)
         if (!VisibleModules.isVisible(TransImported))
@@ -484,7 +485,8 @@ Sema::ActOnModuleDecl(SourceLocation StartLoc, SourceLocation ModuleLoc,
   // and return the import decl to be added to the current TU.
   if (Interface) {
 
-    makeTransitiveImportsVisible(VisibleModules, Interface, Mod, ModuleLoc,
+    makeTransitiveImportsVisible(getASTContext(), VisibleModules, Interface,
+                                 Mod, ModuleLoc,
                                  /*IsImportingPrimaryModuleInterface=*/true);
 
     // Make the import decl for the interface in the impl module.
@@ -643,8 +645,8 @@ DeclResult Sema::ActOnModuleImport(SourceLocation StartLoc,
     Diag(ImportLoc, diag::warn_experimental_header_unit);
 
   if (Mod->isNamedModule())
-    makeTransitiveImportsVisible(VisibleModules, Mod, getCurrentModule(),
-                                 ImportLoc);
+    makeTransitiveImportsVisible(getASTContext(), VisibleModules, Mod,
+                                 getCurrentModule(), ImportLoc);
   else
     VisibleModules.setVisible(Mod, ImportLoc);