[C++20][Modules] Fix crash/compiler error due broken AST links (#123648)

Summary:
This PR fixes bugreport
https://github.com/llvm/llvm-project/issues/122493 The root problem is
the same as before lambda function and DeclRefExpr references a variable
that does not belong to the same module as the enclosing function body.
Therefore iteration over the function body doesn’t visit the VarDecl.
Before this change RelatedDeclsMap was created only for canonical decl
but in reality it has to be done for the definition of the function that
does not always match the canonical decl.

Test Plan: check-clang

clang/docs/ReleaseNotes.rst

@@ -975,6 +975,8 @@ Bug Fixes to C++ Support
 - Fixed a nested lambda substitution issue for constraint evaluation. (#GH123441)
 - Fixed various false diagnostics related to the use of immediate functions. (#GH123472)
 - Fix immediate escalation not propagating through inherited constructors.  (#GH112677)
+- Fixed assertions or false compiler diagnostics in the case of C++ modules for
+  lambda functions or inline friend functions defined inside templates (#GH122493).
 
 Bug Fixes to AST Handling
 ^^^^^^^^^^^^^^^^^^^^^^^^^

clang/include/clang/Serialization/ASTReader.h

@@ -546,11 +546,18 @@ private:
 
   /// Mapping from main decl ID to the related decls IDs.
   ///
-  /// These related decls have to be loaded right after the main decl.
-  /// It is required to have canonical declaration for related decls from the
-  /// same module as the enclosing main decl. Without this, due to lazy
-  /// deserialization, canonical declarations for the main decl and related can
-  /// be selected from different modules.
+  /// The key is the main decl ID, and the value is a vector of related decls
+  /// that must be loaded immediately after the main decl. This is necessary
+  /// to ensure that the definition for related decls comes from the same module
+  /// as the enclosing main decl. Without this, due to lazy deserialization,
+  /// the definition for the main decl and related decls may come from different
+  /// modules. It is used for the following cases:
+  /// - Lambda inside a template function definition: The main declaration is
+  ///   the enclosing function, and the related declarations are the lambda
+  ///   declarations.
+  /// - Friend function defined inside a template CXXRecord declaration: The
+  ///   main declaration is the enclosing record, and the related declarations
+  ///   are the friend functions.
   llvm::DenseMap<GlobalDeclID, SmallVector<GlobalDeclID, 4>> RelatedDeclsMap;
 
   struct PendingUpdateRecord {

clang/lib/Serialization/ASTWriterDecl.cpp

@@ -27,6 +27,20 @@
 using namespace clang;
 using namespace serialization;
 
+//===----------------------------------------------------------------------===//
+// Utility functions
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+// Helper function that returns true if the decl passed in the argument is
+// a defintion in dependent contxt.
+template <typename DT> bool isDefinitionInDependentContext(DT *D) {
+  return D->isDependentContext() && D->isThisDeclarationADefinition();
+}
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Declaration serialization
 //===----------------------------------------------------------------------===//
@@ -801,14 +815,14 @@ void ASTDeclWriter::VisitFunctionDecl(FunctionDecl *D) {
   }
 
   if (D->getFriendObjectKind()) {
-    // For a function defined inline within a class template, we have to force
-    // the canonical definition to be the one inside the canonical definition of
-    // the template. Remember this relation to deserialize them together.
-    if (auto *RD = dyn_cast<CXXRecordDecl>(D->getLexicalParent()))
-      if (RD->isDependentContext() && RD->isThisDeclarationADefinition()) {
-        Writer.RelatedDeclsMap[Writer.GetDeclRef(RD)].push_back(
-            Writer.GetDeclRef(D));
-      }
+    // For a friend function defined inline within a class template, we have to
+    // force the definition to be the one inside the definition of the template
+    // class. Remember this relation to deserialize them together.
+    if (auto *RD = dyn_cast<CXXRecordDecl>(D->getLexicalParent());
+        RD && isDefinitionInDependentContext(RD)) {
+      Writer.RelatedDeclsMap[Writer.GetDeclRef(RD)].push_back(
+          Writer.GetDeclRef(D));
+    }
   }
 
   Record.push_back(D->param_size());
@@ -1583,9 +1597,10 @@ void ASTDeclWriter::VisitCXXRecordDecl(CXXRecordDecl *D) {
     } else {
       Record.push_back(0);
     }
-    // For lambdas inside canonical FunctionDecl remember the mapping.
-    if (auto FD = llvm::dyn_cast_or_null<FunctionDecl>(D->getDeclContext());
-        FD && FD->isCanonicalDecl()) {
+    // For lambdas inside template functions, remember the mapping to
+    // deserialize them together.
+    if (auto *FD = llvm::dyn_cast_or_null<FunctionDecl>(D->getDeclContext());
+        FD && isDefinitionInDependentContext(FD)) {
       Writer.RelatedDeclsMap[Writer.GetDeclRef(FD)].push_back(
           Writer.GetDeclRef(D));
     }

clang/test/Headers/crash-instantiated-in-scope-cxx-modules5.cpp

@@ -0,0 +1,92 @@
+// RUN: rm -fR %t
+// RUN: split-file %s %t
+// RUN: cd %t
+// RUN: %clang_cc1 -verify -std=c++20 -Werror=uninitialized -xc++ -emit-module module.cppmap -fmodule-name=mock_resolver -o mock_resolver.pcm
+// RUN: %clang_cc1 -verify -std=c++20 -Werror=uninitialized -xc++ -emit-module module.cppmap -fmodule-name=sql_internal -o sql_internal.pcm
+// RUN: %clang_cc1 -verify -std=c++20 -Werror=uninitialized -xc++ -fmodule-file=mock_resolver.pcm -fmodule-file=sql_internal.pcm main.cc -o main.o
+
+//--- module.cppmap
+module "mock_resolver" {
+  export *
+  module "mock_resolver.h" {
+    export *
+    header "mock_resolver.h"
+  }
+}
+
+module "sql_internal" {
+  export *
+  module "sql_transform_builder.h" {
+    export *
+    header "sql_transform_builder.h"
+  }
+}
+
+//--- set_bits2.h
+// expected-no-diagnostics
+#pragma once
+
+template <typename T>
+void fwd(const T& x) {}
+
+namespace vox::bitset {
+
+template <typename TFunc>
+void ForEachSetBit2(const TFunc&) {
+  fwd([](int) {
+    const int bit_index_base = 0;
+    (void)[&](int) {
+      int v = bit_index_base;
+    };
+  });
+}
+
+}  // namespace vox::bitset
+
+//--- sql_transform_builder.h
+// expected-no-diagnostics
+#pragma once
+
+#include "set_bits2.h"
+
+class QualifyingSet3 {
+ public:
+  void GetIndexes() const {
+    vox::bitset::ForEachSetBit2([]() {});
+  }
+};
+
+template <typename T>
+void DoTransform() {
+  vox::bitset::ForEachSetBit2([]() {});
+}
+
+//--- mock_resolver.h
+// expected-no-diagnostics
+#pragma once 
+#include "set_bits2.h"
+
+class QualifyingSet2 {
+ public:
+  void GetIndexes() const {
+    vox::bitset::ForEachSetBit2([]() {});
+  }
+};
+
+//--- main.cc
+// expected-no-diagnostics
+#include "sql_transform_builder.h"
+
+template <typename Callable>
+void get(const Callable& fn) {
+  fwd<Callable>(fn);
+}
+
+namespace {
+
+void test() {
+  get([]() {});
+  DoTransform<int>();
+}
+
+} // namespace