Visit lambdas right after FunctionDecl

Author: dmpolukhin

clang/lib/AST/DeclTemplate.cpp

@@ -349,12 +349,8 @@ void RedeclarableTemplateDecl::loadLazySpecializationsImpl() const {
     GlobalDeclID *Specs = CommonBasePtr->LazySpecializations;
     CommonBasePtr->LazySpecializations = nullptr;
     unsigned SpecSize = (*Specs++).getRawValue();
-    // Load the specializations in reverse order so that the most recent
-    // specialization are visited first so they become canonical declarations.
-    // This order matches the order in which namelookup discovers declarations
-    // coming from modules.
-    for (unsigned I = SpecSize; I != 0; --I)
-      (void)Context.getExternalSource()->GetExternalDecl(Specs[I - 1]);
+    for (unsigned I = 0; I != SpecSize; ++I)
+      (void)Context.getExternalSource()->GetExternalDecl(Specs[I]);
   }
 }
 

clang/lib/Serialization/ASTReaderDecl.cpp

@@ -1155,6 +1155,15 @@ void ASTDeclReader::VisitFunctionDecl(FunctionDecl *FD) {
   for (unsigned I = 0; I != NumParams; ++I)
     Params.push_back(readDeclAs<ParmVarDecl>());
   FD->setParams(Reader.getContext(), Params);
+
+  // For the first decl read all lambdas inside, otherwise skip them.
+  unsigned NumLambdas = Record.readInt();
+  if (FD->isFirstDecl()) {
+    for (unsigned I = 0; I != NumLambdas; ++I)
+      readDecl();
+  } else {
+    (void)Record.readIntArray(NumLambdas);
+  }
 }
 
 void ASTDeclReader::VisitObjCMethodDecl(ObjCMethodDecl *MD) {

clang/lib/Serialization/ASTWriterDecl.cpp

@@ -18,6 +18,7 @@
 #include "clang/AST/Expr.h"
 #include "clang/AST/OpenMPClause.h"
 #include "clang/AST/PrettyDeclStackTrace.h"
+#include "clang/AST/StmtVisitor.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Serialization/ASTReader.h"
 #include "clang/Serialization/ASTRecordWriter.h"
@@ -625,6 +626,33 @@ void ASTDeclWriter::VisitDeclaratorDecl(DeclaratorDecl *D) {
                                            : QualType());
 }
 
+static llvm::SmallVector<const Decl *, 2> collectLambdas(FunctionDecl *D) {
+  struct LambdaCollector : public ConstStmtVisitor<LambdaCollector> {
+    llvm::SmallVectorImpl<const Decl *> &Lambdas;
+
+    LambdaCollector(llvm::SmallVectorImpl<const Decl *> &Lambdas)
+        : Lambdas(Lambdas) {}
+
+    void VisitLambdaExpr(const LambdaExpr *E) {
+      VisitStmt(E);
+      Lambdas.push_back(E->getLambdaClass());
+    }
+
+    void VisitStmt(const Stmt *S) {
+      if (!S)
+        return;
+      for (const Stmt *Child : S->children())
+        if (Child)
+          Visit(Child);
+    }
+  };
+
+  llvm::SmallVector<const Decl *, 2> Lambdas;
+  if (D->hasBody())
+    LambdaCollector(Lambdas).VisitStmt(D->getBody());
+  return Lambdas;
+}
+
 void ASTDeclWriter::VisitFunctionDecl(FunctionDecl *D) {
   static_assert(DeclContext::NumFunctionDeclBits == 44,
                 "You need to update the serializer after you change the "
@@ -764,6 +792,15 @@ void ASTDeclWriter::VisitFunctionDecl(FunctionDecl *D) {
   Record.push_back(D->param_size());
   for (auto *P : D->parameters())
     Record.AddDeclRef(P);
+
+  // Store references to all lambda decls inside function to load them
+  // immediately after loading the function to make sure that canonical
+  // decls for lambdas will be from the same module.
+  llvm::SmallVector<const Decl *, 2> Lambdas = collectLambdas(D);
+  Record.push_back(Lambdas.size());
+  for (const auto *L : Lambdas)
+    Record.AddDeclRef(L);
+
   Code = serialization::DECL_FUNCTION;
 }
 

clang/test/Modules/odr_hash.cpp

@@ -3084,8 +3084,8 @@ struct S5 {
 };
 #else
 S5 s5;
-// expected-error@first.h:* {{'PointersAndReferences::S5::x' from module 'FirstModule' is not present in definition of 'PointersAndReferences::S5' in module 'SecondModule'}}
-// expected-note@second.h:* {{declaration of 'x' does not match}}
+// expected-error@second.h:* {{'PointersAndReferences::S5::x' from module 'SecondModule' is not present in definition of 'PointersAndReferences::S5' in module 'FirstModule'}}
+// expected-note@first.h:* {{declaration of 'x' does not match}}
 #endif
 
 #if defined(FIRST)