[clang] [C++20][Modules] Fix non-determinism in serialized AST (PR #110131)

[Dmitry Polukhin via cfe-commits]

https://github.com/dmpolukhin created https://github.com/llvm/llvm-project/pull/110131

Summary:
https://github.com/llvm/llvm-project/pull/109167 serializes FunctionToLambdasMap in the order of pointers in DenseMap. It gives different order with different memory layouts. Fix this issue by using LocalDeclID instead of pointers.

Test Plan: check-clang

>From c28fd2bb947ab3325a304efed52d17181286e7f3 Mon Sep 17 00:00:00 2001
From: Dmitry Polukhin <dmitry.polukhin at gmail.com>
Date: Thu, 26 Sep 2024 07:53:55 -0700
Subject: [PATCH] [C++20][Modules] Fix non-determinism in serialized AST

Summary:
https://github.com/llvm/llvm-project/pull/109167 serialized
FunctionToLambdasMap in the order of pointers in DenseMap. It gives
different order with different memeory leayout. Fix this issue by using
LocalDeclID instead of pointers.

Test Plan: check-clang
---
 clang/include/clang/AST/DeclID.h              | 22 +++++++++++++++++++
 clang/include/clang/Serialization/ASTWriter.h |  6 ++---
 clang/lib/Serialization/ASTWriter.cpp         |  3 +--
 clang/lib/Serialization/ASTWriterDecl.cpp     |  3 ++-
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/clang/include/clang/AST/DeclID.h b/clang/include/clang/AST/DeclID.h
index f4607e42c4be38..49964b43c7d1d8 100644
--- a/clang/include/clang/AST/DeclID.h
+++ b/clang/include/clang/AST/DeclID.h
@@ -189,6 +189,7 @@ class LocalDeclID : public DeclIDBase {
   // Every Decl ID is a local decl ID to the module being writing in ASTWriter.
   friend class ASTWriter;
   friend class GlobalDeclID;
+  friend struct llvm::DenseMapInfo<clang::LocalDeclID>;
 
 public:
   LocalDeclID() : Base() {}
@@ -267,6 +268,27 @@ template <> struct DenseMapInfo<clang::GlobalDeclID> {
   }
 };
 
+template <> struct DenseMapInfo<clang::LocalDeclID> {
+  using LocalDeclID = clang::LocalDeclID;
+  using DeclID = LocalDeclID::DeclID;
+
+  static LocalDeclID getEmptyKey() {
+    return LocalDeclID(DenseMapInfo<DeclID>::getEmptyKey());
+  }
+
+  static LocalDeclID getTombstoneKey() {
+    return LocalDeclID(DenseMapInfo<DeclID>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const LocalDeclID &Key) {
+    return DenseMapInfo<DeclID>::getHashValue(Key.getRawValue());
+  }
+
+  static bool isEqual(const LocalDeclID &L, const LocalDeclID &R) {
+    return L == R;
+  }
+};
+
 } // namespace llvm
 
 #endif
diff --git a/clang/include/clang/Serialization/ASTWriter.h b/clang/include/clang/Serialization/ASTWriter.h
index 760866fd9de938..e21d41c8673143 100644
--- a/clang/include/clang/Serialization/ASTWriter.h
+++ b/clang/include/clang/Serialization/ASTWriter.h
@@ -233,13 +233,13 @@ class ASTWriter : public ASTDeserializationListener,
   /// instead of comparing the result of `getDeclID()` or `GetDeclRef()`.
   llvm::SmallPtrSet<const Decl *, 32> PredefinedDecls;
 
-  /// Mapping from FunctionDecl to the list of lambda IDs inside the function.
+  /// Mapping from FunctionDecl ID to the list of lambda IDs inside the
+  /// function.
   ///
   /// These lambdas have to be loaded right after the function they belong to.
   /// In order to have canonical declaration for lambda class from the same
   /// module as enclosing function during deserialization.
-  llvm::DenseMap<const Decl *, SmallVector<LocalDeclID, 4>>
-      FunctionToLambdasMap;
+  llvm::DenseMap<LocalDeclID, SmallVector<LocalDeclID, 4>> FunctionToLambdasMap;
 
   /// Offset of each declaration in the bitstream, indexed by
   /// the declaration's ID.
diff --git a/clang/lib/Serialization/ASTWriter.cpp b/clang/lib/Serialization/ASTWriter.cpp
index 223727366f61bd..7a40c5c65d39d9 100644
--- a/clang/lib/Serialization/ASTWriter.cpp
+++ b/clang/lib/Serialization/ASTWriter.cpp
@@ -5713,8 +5713,7 @@ void ASTWriter::WriteDeclAndTypes(ASTContext &Context) {
     // efficent becuase it allows lazy deserialization.
     RecordData FunctionToLambdasMapRecord;
     for (const auto &Pair : FunctionToLambdasMap) {
-      FunctionToLambdasMapRecord.push_back(
-          GetDeclRef(Pair.first).getRawValue());
+      FunctionToLambdasMapRecord.push_back(Pair.first.getRawValue());
       FunctionToLambdasMapRecord.push_back(Pair.second.size());
       for (const auto &Lambda : Pair.second)
         FunctionToLambdasMapRecord.push_back(Lambda.getRawValue());
diff --git a/clang/lib/Serialization/ASTWriterDecl.cpp b/clang/lib/Serialization/ASTWriterDecl.cpp
index 50c090b195d619..b9222a1b33fd74 100644
--- a/clang/lib/Serialization/ASTWriterDecl.cpp
+++ b/clang/lib/Serialization/ASTWriterDecl.cpp
@@ -1524,7 +1524,8 @@ void ASTDeclWriter::VisitCXXRecordDecl(CXXRecordDecl *D) {
     // For lambdas inside canonical FunctionDecl remember the mapping.
     if (auto FD = llvm::dyn_cast_or_null<FunctionDecl>(D->getDeclContext());
         FD && FD->isCanonicalDecl()) {
-      Writer.FunctionToLambdasMap[FD].push_back(Writer.GetDeclRef(D));
+      Writer.FunctionToLambdasMap[Writer.GetDeclRef(FD)].push_back(
+          Writer.GetDeclRef(D));
     }
   } else {
     Record.push_back(CXXRecNotTemplate);