[C++20] [Modules] Merge lambdas from Sema to imported one

Close llvm#102721

Previously we tried to merge lambdas between TUs. But we may still meet
problems if we import and then include the headers containing the same
lambda.

The solution in this patch is, when we assigned a lambda numbering in a
sema, try to see if there is already a same lambda imported. If yes, try
to merge them.

clang/lib/Serialization/ASTReader.cpp

@@ -8962,17 +8962,6 @@ 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::LoadSelector(Selector Sel) {
   // It would be complicated to avoid reading the methods anyway. So don't.
   ReadMethodPool(Sel);

clang/lib/Serialization/ASTReaderDecl.cpp

@@ -4506,6 +4506,34 @@ void ASTReader::loadObjCCategories(GlobalDeclID ID, ObjCInterfaceDecl *D,
   ModuleMgr.visit(Visitor);
 }
 
+void ASTReader::AssignedLambdaNumbering(const CXXRecordDecl *Lambda) {
+  if (Lambda->getLambdaContextDecl()) {
+    auto ContextAndIndexPair = std::make_pair<const Decl *, unsigned>(
+        Lambda->getLambdaContextDecl()->getCanonicalDecl(),
+        Lambda->getLambdaIndexInContext());
+
+    auto Iter = LambdaDeclarationsForMerging.find(ContextAndIndexPair);
+    // If we import first and then include second, try to merge the second
+    // one with the first one.
+    if (Iter != LambdaDeclarationsForMerging.end() &&
+        Iter->second->isFromASTFile()) {
+      NamedDecl *Slot = Iter->second;
+      ASTDeclMerger Merger(*this);
+      // The lambda shouldn't be a key declaration as the lambda wouldn't
+      // be a canonical decl for its owning module.
+      Merger.mergeRedeclarableImpl(const_cast<CXXRecordDecl *>(Lambda),
+                                   cast<TagDecl>(Slot),
+                                   /*KeyDeclID*/ GlobalDeclID());
+      return;
+    }
+
+    // Keep track of this lambda so it can be merged with another lambda that
+    // is loaded later.
+    LambdaDeclarationsForMerging.insert(
+        {ContextAndIndexPair, const_cast<CXXRecordDecl *>(Lambda)});
+  }
+}
+
 template<typename DeclT, typename Fn>
 static void forAllLaterRedecls(DeclT *D, Fn F) {
   F(D);

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();
+}