[C++20] [Modules] Merge lambdas in source to imported lambdas

clang/include/clang/AST/Decl.h

@@ -3366,6 +3366,7 @@ class IndirectFieldDecl : public ValueDecl,
 /// 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

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).

clang/include/clang/Sema/ExternalSemaSource.h

@@ -233,7 +233,7 @@ class ExternalSemaSource : public ExternalASTSource {
   /// 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.
   /// \{

clang/include/clang/Sema/MultiplexExternalSemaSource.h

@@ -361,7 +361,7 @@ class MultiplexExternalSemaSource : public ExternalSemaSource {
                                         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.
   /// \{

clang/include/clang/Serialization/ASTReader.h

@@ -2149,7 +2149,7 @@ class ASTReader
       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);

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

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) {

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