[C++20] [Modules] Handle transitive import in the module properly

Close https://github.com/llvm/llvm-project/issues/84002

Per [module.import]p7:

> Additionally, when a module-import-declaration in a module unit of
> some module M imports another module unit U of M, it also imports all
> translation units imported by non-exported module-import-declarations
> in the module unit purview of U.

However, we only tried to implement it during the implicit import of
primary module interface for module implementation unit.

Also we didn't implement the last sentence from [module.import]p7
completely:

> These rules can in turn lead to the importation of yet more
> translation units.

This patch tries to care the both issues.

5 files changed, 210 insertions, 17 deletions

diff --git a/clang/docs/ReleaseNotes.rst b/clang/docs/ReleaseNotes.rst
index 5e0352a..b074055 100644
--- a/clang/docs/ReleaseNotes.rst
+++ b/clang/docs/ReleaseNotes.rst
@@ -86,6 +86,11 @@ C++20 Feature Support
 - Implemented the `__is_layout_compatible` intrinsic to support
   `P0466R5: Layout-compatibility and Pointer-interconvertibility Traits <https://wg21.link/P0466R5>`_.
 
+- Clang now implements [module.import]p7 fully. Clang now will import module
+  units transitively for the module units coming from the same module of the
+  current module units.
+  Fixes `#84002 <https://github.com/llvm/llvm-project/issues/84002>`_.
+
 C++23 Feature Support
 ^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/clang/include/clang/Basic/Module.h b/clang/include/clang/Basic/Module.h
index 30ec9c9..9f62c05 100644
--- a/clang/include/clang/Basic/Module.h
+++ b/clang/include/clang/Basic/Module.h
@@ -598,6 +598,11 @@ public:
            Kind == ModulePartitionImplementation;
   }
 
+  /// Is this a module partition implementation unit.
+  bool isModulePartitionImplementation() const {
+    return Kind == ModulePartitionImplementation;
+  }
+
   /// Is this a module implementation.
   bool isModuleImplementation() const {
     return Kind == ModuleImplementationUnit;
@@ -853,12 +858,6 @@ public:
                   VisibleCallback Vis = [](Module *) {},
                   ConflictCallback Cb = [](ArrayRef<Module *>, Module *,
                                            StringRef) {});
-
-  /// Make transitive imports visible for [module.import]/7.
-  void makeTransitiveImportsVisible(
-      Module *M, SourceLocation Loc, VisibleCallback Vis = [](Module *) {},
-      ConflictCallback Cb = [](ArrayRef<Module *>, Module *, StringRef) {});
-
 private:
   /// Import locations for each visible module. Indexed by the module's
   /// VisibilityID.
diff --git a/clang/lib/Basic/Module.cpp b/clang/lib/Basic/Module.cpp
index 1c5043a..9f597dc 100644
--- a/clang/lib/Basic/Module.cpp
+++ b/clang/lib/Basic/Module.cpp
@@ -722,14 +722,6 @@ void VisibleModuleSet::setVisible(Module *M, SourceLocation Loc,
   VisitModule({M, nullptr});
 }
 
-void VisibleModuleSet::makeTransitiveImportsVisible(Module *M,
-                                                    SourceLocation Loc,
-                                                    VisibleCallback Vis,
-                                                    ConflictCallback Cb) {
-  for (auto *I : M->Imports)
-    setVisible(I, Loc, Vis, Cb);
-}
-
 ASTSourceDescriptor::ASTSourceDescriptor(Module &M)
     : Signature(M.Signature), ClangModule(&M) {
   if (M.Directory)
diff --git a/clang/lib/Sema/SemaModule.cpp b/clang/lib/Sema/SemaModule.cpp
index ed7f626..f08c1cb3 100644
--- a/clang/lib/Sema/SemaModule.cpp
+++ b/clang/lib/Sema/SemaModule.cpp
@@ -73,6 +73,90 @@ static std::string stringFromPath(ModuleIdPath Path) {
   return Name;
 }
 
+/// Helper function for makeTransitiveImportsVisible to decide whether
+/// the \param Imported module unit is in the same module with the \param
+/// CurrentModule.
+/// \param FoundPrimaryModuleInterface is a helper parameter to record the
+/// primary module interface unit corresponding to the module \param
+/// 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,
+                                    Module *&FoundPrimaryModuleInterface) {
+  if (!Imported->isNamedModule())
+    return false;
+
+  // The a partition unit we're importing must be in the same module of the
+  // current module.
+  if (Imported->isModulePartition())
+    return true;
+
+  // If we found the primary module interface during the search process, we can
+  // return quickly to avoid expensive string comparison.
+  if (FoundPrimaryModuleInterface)
+    return Imported == FoundPrimaryModuleInterface;
+
+  if (!CurrentModule)
+    return false;
+
+  // Then the imported module must be a primary module interface unit.  It
+  // is only allowed to import the primary module interface unit from the same
+  // module in the implementation unit and the implementation partition unit.
+
+  // Since we'll handle implementation unit above. We can only care
+  // about the implementation partition unit here.
+  if (!CurrentModule->isModulePartitionImplementation())
+    return false;
+
+  if (Imported->getPrimaryModuleInterfaceName() ==
+      CurrentModule->getPrimaryModuleInterfaceName()) {
+    assert(!FoundPrimaryModuleInterface ||
+           FoundPrimaryModuleInterface == Imported);
+    FoundPrimaryModuleInterface = Imported;
+    return true;
+  }
+
+  return false;
+}
+
+/// [module.import]p7:
+///   Additionally, when a module-import-declaration in a module unit of some
+///   module M imports another module unit U of M, it also imports all
+///   translation units imported by non-exported module-import-declarations in
+///   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,
+                             bool IsImportingPrimaryModuleInterface = false) {
+  assert(Imported->isNamedModule() &&
+         "'makeTransitiveImportsVisible()' is intended for standard C++ named "
+         "modules only.");
+
+  llvm::SmallVector<Module *, 4> Worklist;
+  Worklist.push_back(Imported);
+
+  Module *FoundPrimaryModuleInterface =
+      IsImportingPrimaryModuleInterface ? Imported : nullptr;
+
+  while (!Worklist.empty()) {
+    Module *Importing = Worklist.pop_back_val();
+
+    if (VisibleModules.isVisible(Importing))
+      continue;
+
+    // FIXME: The ImportLoc here is not meaningful. It may be problematic if we
+    // use the sourcelocation loaded from the visible modules.
+    VisibleModules.setVisible(Importing, ImportLoc);
+
+    if (isImportingModuleUnitFromSameModule(Importing, CurrentModule,
+                                            FoundPrimaryModuleInterface))
+      for (Module *TransImported : Importing->Imports)
+        if (!VisibleModules.isVisible(TransImported))
+          Worklist.push_back(TransImported);
+  }
+}
+
 Sema::DeclGroupPtrTy
 Sema::ActOnGlobalModuleFragmentDecl(SourceLocation ModuleLoc) {
   // We start in the global module;
@@ -396,8 +480,8 @@ Sema::ActOnModuleDecl(SourceLocation StartLoc, SourceLocation ModuleLoc,
   // and return the import decl to be added to the current TU.
   if (Interface) {
 
-    VisibleModules.setVisible(Interface, ModuleLoc);
-    VisibleModules.makeTransitiveImportsVisible(Interface, ModuleLoc);
+    makeTransitiveImportsVisible(VisibleModules, Interface, Mod, ModuleLoc,
+                                 /*IsImportingPrimaryModuleInterface=*/true);
 
     // Make the import decl for the interface in the impl module.
     ImportDecl *Import = ImportDecl::Create(Context, CurContext, ModuleLoc,
@@ -554,7 +638,11 @@ DeclResult Sema::ActOnModuleImport(SourceLocation StartLoc,
   if (Mod->isHeaderUnit())
     Diag(ImportLoc, diag::warn_experimental_header_unit);
 
-  VisibleModules.setVisible(Mod, ImportLoc);
+  if (Mod->isNamedModule())
+    makeTransitiveImportsVisible(VisibleModules, Mod, getCurrentModule(),
+                                 ImportLoc);
+  else
+    VisibleModules.setVisible(Mod, ImportLoc);
 
   checkModuleImportContext(*this, Mod, ImportLoc, CurContext);
 
diff --git a/clang/test/Modules/transitive-import.cppm b/clang/test/Modules/transitive-import.cppm
new file mode 100644
index 0000000..0eed8cf
--- /dev/null
+++ b/clang/test/Modules/transitive-import.cppm
@@ -0,0 +1,109 @@
+// RUN: rm -rf %t
+// RUN: mkdir -p %t
+// RUN: split-file %s %t
+//
+// RUN: %clang_cc1 -std=c++20 %t/Invisible.cppm -emit-module-interface -o %t/Invisible.pcm
+// RUN: %clang_cc1 -std=c++20 %t/Other.cppm -emit-module-interface -fprebuilt-module-path=%t \
+// RUN:     -o %t/Other.pcm
+// RUN: %clang_cc1 -std=c++20 %t/Another.cppm -emit-module-interface -o %t/Another.pcm
+// RUN: %clang_cc1 -std=c++20 %t/A-interface.cppm -emit-module-interface \
+// RUN:     -fprebuilt-module-path=%t -o %t/A-interface.pcm
+// RUN: %clang_cc1 -std=c++20 %t/A-interface2.cppm -emit-module-interface \
+// RUN:     -fprebuilt-module-path=%t -o %t/A-interface2.pcm
+// RUN: %clang_cc1 -std=c++20 %t/A-interface3.cppm -emit-module-interface \
+// RUN:     -fprebuilt-module-path=%t -o %t/A-interface3.pcm
+// RUN: %clang_cc1 -std=c++20 %t/A.cppm -emit-module-interface \
+// RUN:     -fprebuilt-module-path=%t -o %t/A.pcm
+
+// RUN: %clang_cc1 -std=c++20 %t/A.cpp -fprebuilt-module-path=%t -fsyntax-only -verify
+// RUN: %clang_cc1 -std=c++20 %t/A-impl.cppm -fprebuilt-module-path=%t -fsyntax-only -verify
+
+// RUN: %clang_cc1 -std=c++20 %t/A-impl2.cppm -fprebuilt-module-path=%t -fsyntax-only -verify
+
+//--- Invisible.cppm
+export module Invisible;
+export void invisible() {}
+
+//--- Other.cppm
+export module Other;
+import Invisible;
+export void other() {}
+
+//--- Another.cppm
+export module Another;
+export void another() {}
+
+//--- A-interface.cppm
+export module A:interface;
+import Other;
+export void a_interface() {}
+
+//--- A-interface2.cppm
+export module A:interface2;
+import Another;
+export void a_interface2() {}
+
+//--- A-interface3.cppm
+export module A:interface3;
+import :interface;
+import :interface2;
+export void a_interface3() {}
+
+//--- A.cppm
+export module A;
+import Another;
+import :interface;
+import :interface2;
+import :interface3;
+
+export void a() {}
+export void impl();
+
+//--- A.cpp
+module A;
+void impl() {
+    a_interface();
+    a_interface2();
+    a_interface3();
+
+    other();
+    another();
+
+    invisible(); // expected-error {{declaration of 'invisible' must be imported from module 'Invisible' before it is required}}
+                 // expected-note@* {{declaration here is not visible}}
+}
+
+//--- A-impl.cppm
+module A:impl;
+import :interface3;
+
+void impl_part() {
+    a_interface();
+    a_interface2();
+    a_interface3();
+
+    other();
+    another();
+
+    invisible(); // expected-error {{declaration of 'invisible' must be imported from module 'Invisible' before it is required}}
+                 // expected-note@* {{declaration here is not visible}}
+}
+
+//--- A-impl2.cppm
+module A:impl2;
+import A;
+
+void impl_part2() {
+    a();
+    impl();
+
+    a_interface();
+    a_interface2();
+    a_interface3();
+
+    other();
+    another();
+
+    invisible(); // expected-error {{declaration of 'invisible' must be imported from module 'Invisible' before it is required}}
+                 // expected-note@* {{declaration here is not visible}}
+}