From c94bb7fee15c0411d3b5ef4ed41b53ee8d385302 Mon Sep 17 00:00:00 2001
From: Dan Katz <katzdm@gmail.com>
Date: Thu, 13 Feb 2025 14:39:42 -0500
Subject: [PATCH] Experimental substitution caching.

---
 clang/include/clang/AST/ASTContext.h |  6 ++++++
 clang/lib/AST/ASTContext.cpp         | 15 +++++++++++++++
 clang/lib/AST/ExprConstantMeta.cpp   | 22 +++++++++++++++++++++-
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/clang/include/clang/AST/ASTContext.h b/clang/include/clang/AST/ASTContext.h
index 59778e430102..bbe34fcd655e 100644
--- a/clang/include/clang/AST/ASTContext.h
+++ b/clang/include/clang/AST/ASTContext.h
@@ -352,6 +352,9 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// hashes of their member-specifications.
   mutable llvm::DenseMap<QualType, unsigned> ClsMemberSpecHashes;
 
+  /// A mapping from 'substitute' inputs to specializations.
+  mutable llvm::DenseMap<unsigned, APValue> SubstitutionHashes;
+
   /// MD5 hash of CUID. It is calculated when first used and cached by this
   /// data member.
   mutable std::string CUIDHash;
@@ -3362,6 +3365,9 @@ public:
   bool checkClassMemberSpecHash(QualType QT, unsigned &Out);
   void recordClassMemberSpecHash(QualType QT, unsigned Hash);
 
+  bool checkCachedSubstitution(unsigned Hash, APValue *Out);
+  void recordCachedSubstitution(unsigned Hash, const APValue &Value);
+
   /// Return a declaration for the global GUID object representing the given
   /// GUID value.
   MSGuidDecl *getMSGuidDecl(MSGuidDeclParts Parts) const;
diff --git a/clang/lib/AST/ASTContext.cpp b/clang/lib/AST/ASTContext.cpp
index 2f4647a80121..f46d74d49cf8 100644
--- a/clang/lib/AST/ASTContext.cpp
+++ b/clang/lib/AST/ASTContext.cpp
@@ -13223,6 +13223,21 @@ void ASTContext::recordClassMemberSpecHash(QualType QT, unsigned Hash) {
   ClsMemberSpecHashes[QT] = Hash;
 }
 
+bool ASTContext::checkCachedSubstitution(unsigned Hash, APValue *Out) {
+  if (SubstitutionHashes.contains(Hash)) {
+    *Out = SubstitutionHashes[Hash];
+    return true;
+  }
+  return false;
+}
+
+void ASTContext::recordCachedSubstitution(unsigned Hash,
+                                          const APValue &Result) {
+  assert(Result.isReflection());
+
+  SubstitutionHashes[Hash] = Result;
+}
+
 MSGuidDecl *
 ASTContext::getMSGuidDecl(MSGuidDecl::Parts Parts) const {
   assert(MSGuidTagDecl && "building MS GUID without MS extensions?");
diff --git a/clang/lib/AST/ExprConstantMeta.cpp b/clang/lib/AST/ExprConstantMeta.cpp
index 225d8a7bdc79..de1a274c38b0 100644
--- a/clang/lib/AST/ExprConstantMeta.cpp
+++ b/clang/lib/AST/ExprConstantMeta.cpp
@@ -2786,6 +2786,17 @@ bool substitute(APValue &Result, ASTContext &C, MetaActions &Meta,
   SmallVector<TemplateArgument, 4> ExpandedTArgs;
   expandTemplateArgPacks(TArgs, ExpandedTArgs);
 
+  // Lookup cached specialization; if found, return it.
+  llvm::FoldingSetNodeID ID;
+  {
+    ID.AddPointer(TDecl);
+    for (const TemplateArgument &TArg : ExpandedTArgs)
+      TArg.Profile(ID, C);
+  }
+  unsigned SubstitutionHash = ID.ComputeHash();
+  if (C.checkCachedSubstitution(SubstitutionHash, &Result))
+    return false;
+
   if (!Meta.CheckTemplateArgumentList(TDecl, ExpandedTArgs, false,
                                       Args[0]->getExprLoc()))
     return true;
@@ -2810,6 +2821,7 @@ bool substitute(APValue &Result, ASTContext &C, MetaActions &Meta,
 
     APValue RV(ReflectionKind::Type,
                const_cast<Type *>(TSpecDecl->getTypeForDecl()));
+    C.recordCachedSubstitution(SubstitutionHash, RV);
     return SetAndSucceed(Result, RV);
   } else if (auto *TATD = dyn_cast<TypeAliasTemplateDecl>(TDecl)) {
     TArgs.clear();
@@ -2818,12 +2830,16 @@ bool substitute(APValue &Result, ASTContext &C, MetaActions &Meta,
     QualType QT = Meta.Substitute(TATD, TArgs, Range.getBegin());
     assert(!QT.isNull() && "substitution failed after validating arguments?");
 
+    APValue RV = makeReflection(QT);
+    C.recordCachedSubstitution(SubstitutionHash, RV);
     return SetAndSucceed(Result, makeReflection(QT));
   } else if (auto *FTD = dyn_cast<FunctionTemplateDecl>(TDecl)) {
     FunctionDecl *Spec = Meta.Substitute(FTD, ExpandedTArgs, Range.getBegin());
     assert(Spec && "substitution failed after validating arguments?");
 
-    return SetAndSucceed(Result, makeReflection(Spec));
+    APValue RV = makeReflection(Spec);
+    C.recordCachedSubstitution(SubstitutionHash, RV);
+    return SetAndSucceed(Result, RV);
   } else if (auto *VTD = dyn_cast<VarTemplateDecl>(TDecl)) {
     TArgs.clear();
     expandTemplateArgPacks(ExpandedTArgs, TArgs);
@@ -2831,6 +2847,8 @@ bool substitute(APValue &Result, ASTContext &C, MetaActions &Meta,
     VarDecl *Spec = Meta.Substitute(VTD, TArgs, Range.getBegin());
     assert(Spec && "substitution failed after validating arguments?");
 
+    APValue RV = makeReflection(Spec);
+    C.recordCachedSubstitution(SubstitutionHash, RV);
     return SetAndSucceed(Result, makeReflection(Spec));
   } else if (auto *CD = dyn_cast<ConceptDecl>(TDecl)) {
     TArgs.clear();
@@ -2843,6 +2861,8 @@ bool substitute(APValue &Result, ASTContext &C, MetaActions &Meta,
     if (!Evaluator(SatisfiesConcept, Spec, true))
       llvm_unreachable("failed to evaluate substituted concept");
 
+    APValue RV = SatisfiesConcept.Lift(C.BoolTy);
+    C.recordCachedSubstitution(SubstitutionHash, RV);
     return SetAndSucceed(Result, SatisfiesConcept.Lift(C.BoolTy));
   }
   llvm_unreachable("unimplemented for template kind");