[Clang] Reconsider the timing of instantiation of local constexpr lambdas

[zyn0217]

In the previous patch #95660, we used a strategy of eagerly instantiating local constexpr lambdas. However, that caused a regression in recursive local lambda calls.

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

Closes #98526

(The reduced examples in #97680 seem different, as they don't compile in Clang 18 either. However, #97680 per se should work again with this patch.)

[llvmbot]

@llvm/pr-subscribers-clang

Author: Younan Zhang (zyn0217)

Changes

In the previous patch #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 #98526

(The reduced examples in #97680 seem different, as they don't compile in Clang 18 either. However, #97680 per se should work again with this patch.)

---

Full diff: https://github.com/llvm/llvm-project/pull/98758.diff

2 Files Affected:

• (modified) clang/lib/Sema/SemaExpr.cpp (+45-10)
• (modified) clang/test/SemaTemplate/instantiate-local-class.cpp (+26)

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

[cor3ntin]

Thanks a lot for working on this regression

[cor3ntin]

If we are going to reuse this object, we should rename it - or maybe we need a separate class

[cor3ntin]

A separate class would avoid the pointer to SmallPtrSetImpl, which is not great

[cor3ntin]

Can you explain why we check for not having pending instantiations ?

[cor3ntin]

I think this would benefit from a comment

[zyn0217]

Curiously, if we move the call to lambda wrapper out of the function, then it would compile even on trunk -- https://gcc.godbolt.org/z/aEcajcEvE

I probably overlooked something here.

EDIT: https://gcc.godbolt.org/z/zzGEEWd3h: if we change accumulate(tree) to accumulate(nullptr), it will also compile within the function ?!

[zyn0217]

@cor3ntin I'm stuck here for hours and I think I need some help.

So let me explain what happens before #95660. Consider the example that doesn't compile before #95660:

template <typename F> constexpr int func(F f) {
  constexpr bool y = f(1UL);
  return y;
}

int main() {
  auto predicate = [](auto v) /*implicit constexpr*/ -> bool {
    return v == 1;
  };
  return func(predicate);
}

we would encounter an error saying that undefined function operator()<> cannot be used in a constant expression, which is bogus because the lambda body has been clearly defined. The issue lies in MarkFunctionReferenced(), where we deferred the instantiation of functions that belong to local classes. The logic there may originate from a time when lambdas were not supported, and then when lambda comes into play, this logic would also impact these lambdas that are defined within a function.

We won't have the corresponding lambda bodies as we have deferred the instantiation until the end of the function definition. Therefore, when we instantiate func(), which would entail the evaluation of f(1), which in turn is the evaluation of the lambda, we would fail immediately because, at this time, we couldn't see the lambda body.

If I understand correctly, our constant evaluator is designed to be oblivious of Sema and we almost always eagerly evaluate expressions for constexprs. Following this principle, my patch #95660 does the thing that lets constexpr functions always get instantiated despite the declaration location. And it just works for the case above.

However, as you can see, this also subtly caused such a regression. Again, I'll demonstrate an example:

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);
      }};

  accumulate(tree);
}

In this example, we want to deduce the return type of recursive_lambda::operator() with [Args = Tree*]. To that end, we have to instantiate fn() which is the lambda defined within sumSize. With my patch, the lambda body gets instantiated immediately, so we end up instantiating self_fn() for its return type. As that function is being instantiated, the flag willHaveBody has been set to true so that we won't instantiate it endlessly. We bail out of InstantiateFunctionDefinition early, and because we don't have a deduced return type at the time, we diagnose in DeduceReturnType().

But what will happen if we defer that lambda's instantiation? That means we don't have to instantiate self_fn for its return type, and hence, we would successfully build up a body for recursive_lambda::operator() and a return type for it - we don't have to look into the lambda body because the lambda has an explicit type.

So far, I have thought of some approaches:

1. Revert 95660. Instantiate the function as needed before we perform the evaluation. This is what this patch aims to do; however, since we don't have a common entry for constant evaluation, we probably have to add this "prerequisite function" to many places, which I feel hesitant to do.

2. Don't revert 95660. Just teach DeduceReturnType to not diagnose if the function is being instantiated. This makes things work but causes some other regressions, unfortunately.

3. Don't revert 95660. In addition, introspect the evaluation context to help decide if we should instantiate. This doesn't work because all the cases above have the same "PotentiallyEvaluated" flag set.

4. Revert 95660 and give up. This is likely a design issue, and we admit we struggle to resolve them perfectly.

This has afflicted me for half the night, and I really need some guidance!

[zyn0217]

Finally, I decide to revert #95660 and close this PR for now.

Regarding issues associated with #95660, I realized they are actually part of a bigger problem in constant evaluation. As discussed in #59966, it might be beneficial to implement on-demand template instantiation for constant evaluation. That way, we don't have to touch the logic of deferral instantiation in MarkFunctionReferenced(), (which probably can be removed then) and we could instantiate those local lambdas before evaluation.

[cor3ntin]

Regarding issues associated with #95660, I realized they are actually part of a bigger problem in constant evaluation. As discussed in #59966, it might be beneficial to implement on-demand template instantiation for constant evaluation. That way, we don't have to touch the logic of deferral instantiation in MarkFunctionReferenced(), (which probably can be removed then) and we could instantiate those local lambdas before evaluation.

I haven't had time to look in the issue in details, but this might be the best solution.
We will need to do that for reflection too
@Endilll is planning to wire Sema in constant evaluation in the next few months, that will be a good opportunity to fix all of these issues