[C++20] [Modules] Merge lambdas in source to imported lambdas (#106483)

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

Generally, the type of merged decls will be reused in ASTContext. But
for lambda, in the import and then include case, we can't decide its
previous decl in the imported modules so that we can't assign the
previous decl before creating the type for it. Since we can't decide its
numbering before creating it. So we have to assign the previous decl and
the canonical type for it after creating it, which is unusual and
slightly hack.

8 files changed, 66 insertions, 12 deletions

diff --git a/clang/include/clang/AST/Decl.h b/clang/include/clang/AST/Decl.h
index 6d84bd0..0600ecc 100644
--- a/clang/include/clang/AST/Decl.h
+++ b/clang/include/clang/AST/Decl.h
@@ -3366,6 +3366,7 @@ public:
 /// Represents a declaration of a type.
 class TypeDecl : public NamedDecl {
   friend class ASTContext;
+  friend class ASTReader;
 
   /// This indicates the Type object that represents
   /// this TypeDecl.  It is a cache maintained by
diff --git a/clang/include/clang/AST/Type.h b/clang/include/clang/AST/Type.h
index 575f3c1..5004191 100644
--- a/clang/include/clang/AST/Type.h
+++ b/clang/include/clang/AST/Type.h
@@ -1684,6 +1684,7 @@ class ExtQualsTypeCommonBase {
   friend class ExtQuals;
   friend class QualType;
   friend class Type;
+  friend class ASTReader;
 
   /// The "base" type of an extended qualifiers type (\c ExtQuals) or
   /// a self-referential pointer (for \c Type).
diff --git a/clang/include/clang/Sema/ExternalSemaSource.h b/clang/include/clang/Sema/ExternalSemaSource.h
index 22d1ee2d..11cd69d 100644
--- a/clang/include/clang/Sema/ExternalSemaSource.h
+++ b/clang/include/clang/Sema/ExternalSemaSource.h
@@ -233,7 +233,7 @@ public:
   /// Notify the external source that a lambda was assigned a mangling number.
   /// This enables the external source to track the correspondence between
   /// lambdas and mangling numbers if necessary.
-  virtual void AssignedLambdaNumbering(const CXXRecordDecl *Lambda) {}
+  virtual void AssignedLambdaNumbering(CXXRecordDecl *Lambda) {}
 
   /// LLVM-style RTTI.
   /// \{
diff --git a/clang/include/clang/Sema/MultiplexExternalSemaSource.h b/clang/include/clang/Sema/MultiplexExternalSemaSource.h
index 238fb39..3d1906d 100644
--- a/clang/include/clang/Sema/MultiplexExternalSemaSource.h
+++ b/clang/include/clang/Sema/MultiplexExternalSemaSource.h
@@ -361,7 +361,7 @@ public:
                                         QualType T) override;
 
   // Inform all attached sources that a mangling number was assigned.
-  void AssignedLambdaNumbering(const CXXRecordDecl *Lambda) override;
+  void AssignedLambdaNumbering(CXXRecordDecl *Lambda) override;
 
   /// LLVM-style RTTI.
   /// \{
diff --git a/clang/include/clang/Serialization/ASTReader.h b/clang/include/clang/Serialization/ASTReader.h
index 2d8952d..898f439 100644
--- a/clang/include/clang/Serialization/ASTReader.h
+++ b/clang/include/clang/Serialization/ASTReader.h
@@ -2149,7 +2149,7 @@ public:
       llvm::MapVector<const FunctionDecl *, std::unique_ptr<LateParsedTemplate>>
           &LPTMap) override;
 
-  void AssignedLambdaNumbering(const CXXRecordDecl *Lambda) override;
+  void AssignedLambdaNumbering(CXXRecordDecl *Lambda) override;
 
   /// Load a selector from disk, registering its ID if it exists.
   void LoadSelector(Selector Sel);
diff --git a/clang/lib/Sema/MultiplexExternalSemaSource.cpp b/clang/lib/Sema/MultiplexExternalSemaSource.cpp
index 79e656e..cd44483 100644
--- a/clang/lib/Sema/MultiplexExternalSemaSource.cpp
+++ b/clang/lib/Sema/MultiplexExternalSemaSource.cpp
@@ -343,7 +343,7 @@ bool MultiplexExternalSemaSource::MaybeDiagnoseMissingCompleteType(
 }
 
 void MultiplexExternalSemaSource::AssignedLambdaNumbering(
-    const CXXRecordDecl *Lambda) {
+    CXXRecordDecl *Lambda) {
   for (auto *Source : Sources)
     Source->AssignedLambdaNumbering(Lambda);
 }
diff --git a/clang/lib/Serialization/ASTReader.cpp b/clang/lib/Serialization/ASTReader.cpp
index ffdaec4..a0abd4a 100644
--- a/clang/lib/Serialization/ASTReader.cpp
+++ b/clang/lib/Serialization/ASTReader.cpp
@@ -8963,15 +8963,34 @@ void ASTReader::ReadLateParsedTemplates(
   LateParsedTemplates.clear();
 }
 
-void ASTReader::AssignedLambdaNumbering(const CXXRecordDecl *Lambda) {
-  if (Lambda->getLambdaContextDecl()) {
-    // Keep track of this lambda so it can be merged with another lambda that
-    // is loaded later.
-    LambdaDeclarationsForMerging.insert(
-        {{Lambda->getLambdaContextDecl()->getCanonicalDecl(),
-          Lambda->getLambdaIndexInContext()},
-         const_cast<CXXRecordDecl *>(Lambda)});
+void ASTReader::AssignedLambdaNumbering(CXXRecordDecl *Lambda) {
+  if (!Lambda->getLambdaContextDecl())
+    return;
+
+  auto LambdaInfo =
+      std::make_pair(Lambda->getLambdaContextDecl()->getCanonicalDecl(),
+                     Lambda->getLambdaIndexInContext());
+
+  // Handle the import and then include case for lambdas.
+  if (auto Iter = LambdaDeclarationsForMerging.find(LambdaInfo);
+      Iter != LambdaDeclarationsForMerging.end() &&
+      Iter->second->isFromASTFile() && Lambda->getFirstDecl() == Lambda) {
+    CXXRecordDecl *Previous =
+        cast<CXXRecordDecl>(Iter->second)->getMostRecentDecl();
+    Lambda->setPreviousDecl(Previous);
+    // FIXME: It will be best to use the Previous type when we creating the
+    // lambda directly. But that requires us to get the lambda context decl and
+    // lambda index before creating the lambda, which needs a drastic change in
+    // the parser.
+    const_cast<QualType &>(Lambda->TypeForDecl->CanonicalType) =
+        Previous->TypeForDecl->CanonicalType;
+    return;
   }
+
+  // Keep track of this lambda so it can be merged with another lambda that
+  // is loaded later.
+  LambdaDeclarationsForMerging.insert(
+      {LambdaInfo, const_cast<CXXRecordDecl *>(Lambda)});
 }
 
 void ASTReader::LoadSelector(Selector Sel) {
diff --git a/clang/test/Modules/pr102721.cppm b/clang/test/Modules/pr102721.cppm
new file mode 100644
index 0000000..6a84393
--- /dev/null
+++ b/clang/test/Modules/pr102721.cppm
@@ -0,0 +1,33 @@
+// RUN: rm -rf %t
+// RUN: mkdir -p %t
+// RUN: split-file %s %t
+//
+// RUN: %clang_cc1 -std=c++20 %t/a.cppm -emit-module-interface -o %t/a.pcm
+// RUN: %clang_cc1 -std=c++20 %t/b.cppm -emit-module-interface -o %t/b.pcm \
+// RUN:   -fprebuilt-module-path=%t
+// RUN: %clang_cc1 -std=c++20 %t/test.cc -fsyntax-only -verify \
+// RUN:   -fprebuilt-module-path=%t
+
+//--- foo.h
+inline auto x = []{};
+
+//--- a.cppm
+module;
+#include "foo.h"
+export module a;
+export using ::x;
+
+//--- b.cppm
+module;
+import a;
+#include "foo.h"
+export module b;
+export using ::x;
+
+//--- test.cc
+// expected-no-diagnostics
+import a;
+import b;
+void test() {
+  x();
+}