[clang] [Clang] Reconsider the timing of instantiation of local constexpr lambdas (PR #98758)

[Younan Zhang via cfe-commits]

https://github.com/zyn0217 created https://github.com/llvm/llvm-project/pull/98758

In the previous patch https://github.com/llvm/llvm-project/pull/95660, we used a strategy of eagerly instantiating local constexpr lambdas. However, that caused a regression in recursive local lambda calls.

This patch addressed that by adding an instantiation requirement to the expression constant evaluation, like what we did when deducing the function return type.

Closes https://github.com/llvm/llvm-project/issues/98526

(The reduced examples in https://github.com/llvm/llvm-project/issues/97680 seem different, as they don't compile in Clang 18 either. However, #97680 per se should work again with this patch.)

>From d615943c09abecd777485279645a4c80a63ba199 Mon Sep 17 00:00:00 2001
From: Younan Zhang <zyn7109 at gmail.com>
Date: Sun, 14 Jul 2024 02:06:59 +0800
Subject: [PATCH] [Clang] Reconsider the timing of instantiation of local
 constexpr lambdas

---
 clang/lib/Sema/SemaExpr.cpp                   | 55 +++++++++++++++----
 .../SemaTemplate/instantiate-local-class.cpp  | 26 +++++++++
 2 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
index 852344d895ffd..8414f5c7e46c4 100644
--- a/clang/lib/Sema/SemaExpr.cpp
+++ b/clang/lib/Sema/SemaExpr.cpp
@@ -59,6 +59,7 @@
 #include "clang/Sema/Template.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLForwardCompat.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ConvertUTF.h"
@@ -5340,20 +5341,31 @@ bool Sema::CheckCXXDefaultArgExpr(SourceLocation CallLoc, FunctionDecl *FD,
 
 struct ImmediateCallVisitor : public RecursiveASTVisitor<ImmediateCallVisitor> {
   const ASTContext &Context;
-  ImmediateCallVisitor(const ASTContext &Ctx) : Context(Ctx) {}
+  llvm::SmallPtrSetImpl<const FunctionDecl *> *ReferencedFunctions;
+
+  ImmediateCallVisitor(const ASTContext &Ctx,
+                       llvm::SmallPtrSetImpl<const FunctionDecl *>
+                           *ReferencedFunctions = nullptr)
+      : Context(Ctx), ReferencedFunctions(ReferencedFunctions) {}
 
   bool HasImmediateCalls = false;
   bool shouldVisitImplicitCode() const { return true; }
 
   bool VisitCallExpr(CallExpr *E) {
-    if (const FunctionDecl *FD = E->getDirectCallee())
+    if (const FunctionDecl *FD = E->getDirectCallee()) {
       HasImmediateCalls |= FD->isImmediateFunction();
+      if (ReferencedFunctions)
+        ReferencedFunctions->insert(FD);
+    }
     return RecursiveASTVisitor<ImmediateCallVisitor>::VisitStmt(E);
   }
 
   bool VisitCXXConstructExpr(CXXConstructExpr *E) {
-    if (const FunctionDecl *FD = E->getConstructor())
+    if (const FunctionDecl *FD = E->getConstructor()) {
       HasImmediateCalls |= FD->isImmediateFunction();
+      if (ReferencedFunctions)
+        ReferencedFunctions->insert(FD);
+    }
     return RecursiveASTVisitor<ImmediateCallVisitor>::VisitStmt(E);
   }
 
@@ -16983,6 +16995,30 @@ Sema::VerifyIntegerConstantExpression(Expr *E, llvm::APSInt *Result,
   SmallVector<PartialDiagnosticAt, 8> Notes;
   EvalResult.Diag = &Notes;
 
+  // Check if the expression refers to local functions yet to be instantiated.
+  // If so, instantiate them now, as the constant evaluation requires the
+  // function definition.
+  if (!PendingLocalImplicitInstantiations.empty()) {
+    llvm::SmallPtrSet<const FunctionDecl *, 4> ReferencedFunctions;
+    ImmediateCallVisitor V(getASTContext(), &ReferencedFunctions);
+    V.TraverseStmt(E);
+
+    auto Pred = [&](PendingImplicitInstantiation Pair) {
+      ValueDecl *VD = Pair.first;
+      return isa<FunctionDecl>(VD) &&
+             ReferencedFunctions.contains(cast<FunctionDecl>(VD));
+    };
+    // Workaround: A lambda with captures cannot be copy-assigned, which is
+    // required by llvm::make_filter_range().
+    llvm::function_ref<bool(PendingImplicitInstantiation)> PredRef = Pred;
+
+    auto R =
+        llvm::make_filter_range(PendingLocalImplicitInstantiations, PredRef);
+    LocalEagerInstantiationScope InstantiateReferencedLocalFunctions(*this);
+    PendingLocalImplicitInstantiations = {R.begin(), R.end()};
+    InstantiateReferencedLocalFunctions.perform();
+  }
+
   // Try to evaluate the expression, and produce diagnostics explaining why it's
   // not a constant expression as a side-effect.
   bool Folded =
@@ -17938,17 +17974,16 @@ void Sema::MarkFunctionReferenced(SourceLocation Loc, FunctionDecl *Func,
 
         if (FirstInstantiation || TSK != TSK_ImplicitInstantiation ||
             Func->isConstexpr()) {
-          if (Func->isConstexpr())
+          if (isa<CXXRecordDecl>(Func->getDeclContext()) &&
+              cast<CXXRecordDecl>(Func->getDeclContext())->isLocalClass() &&
+              CodeSynthesisContexts.size())
+            PendingLocalImplicitInstantiations.push_back(
+                std::make_pair(Func, PointOfInstantiation));
+          else if (Func->isConstexpr())
             // Do not defer instantiations of constexpr functions, to avoid the
             // expression evaluator needing to call back into Sema if it sees a
             // call to such a function.
             InstantiateFunctionDefinition(PointOfInstantiation, Func);
-          else if (isa<CXXRecordDecl>(Func->getDeclContext()) &&
-                   cast<CXXRecordDecl>(Func->getDeclContext())
-                       ->isLocalClass() &&
-                   CodeSynthesisContexts.size())
-            PendingLocalImplicitInstantiations.push_back(
-                std::make_pair(Func, PointOfInstantiation));
           else {
             Func->setInstantiationIsPending(true);
             PendingInstantiations.push_back(
diff --git a/clang/test/SemaTemplate/instantiate-local-class.cpp b/clang/test/SemaTemplate/instantiate-local-class.cpp
index 7eee131e28d60..bee4cbaf33dd7 100644
--- a/clang/test/SemaTemplate/instantiate-local-class.cpp
+++ b/clang/test/SemaTemplate/instantiate-local-class.cpp
@@ -532,4 +532,30 @@ int main() {
 
 } // namespace GH35052
 
+namespace GH98526 {
+
+template <typename F> struct recursive_lambda {
+  template <typename... Args> auto operator()(Args &&...args) const {
+    return fn(*this, args...);
+  }
+  F fn;
+};
+
+template <typename F> recursive_lambda(F) -> recursive_lambda<F>;
+
+struct Tree {
+  Tree *left, *right;
+};
+
+int sumSize(Tree *tree) {
+  auto accumulate =
+      recursive_lambda{[&](auto &self_fn, Tree *element_node) -> int {
+        return 1 + self_fn(tree->left) + self_fn(tree->right);
+      }};
+
+  return accumulate(tree);
+}
+
+} // namespace GH98526
+
 #endif