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

Close #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.

Author: ChuanqiXu9

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