[Clang] Fix Handling of Init Capture with Parameter Packs in LambdaScopeForCallOperatorInstantiationRAII

[LYP951018]

This PR addresses issues related to the handling of init capture with parameter packs in Clang's LambdaScopeForCallOperatorInstantiationRAII.

Previously, addInstantiatedCapturesToScope would add init capture containing packs to the scope using the type of the init capture to determine the expanded pack size. However, this approach resulted in a pack size of 0 because getType()->containsUnexpandedParameterPack() returns false. After extensive testing, it appears that the correct pack size can only be inferred from getInit.

But getInit may reference parameters and init capture from an outer lambda, as shown in the following example:

auto L = [](auto... z) {
    return [... w = z](auto... y) {
        // ...
    };
};

To address this, addInstantiatedCapturesToScope in LambdaScopeForCallOperatorInstantiationRAII should be called last. Additionally, addInstantiatedCapturesToScope has been modified to only add init capture to the scope. The previous implementation incorrectly called MakeInstantiatedLocalArgPack for other non-init captures containing packs, resulting in a pack size of 0.

Impact

This patch affects scenarios where LambdaScopeForCallOperatorInstantiationRAII is passed with ShouldAddDeclsFromParentScope = false, preventing the correct addition of the current lambda's init capture to the scope. There are two main scenarios for ShouldAddDeclsFromParentScope = false:

1. Constraints: Sometimes constraints are instantiated in place rather than delayed. In this case, LambdaScopeForCallOperatorInstantiationRAII does not need to add init capture to the scope.
2. noexcept Expressions: The expressions inside noexcept have already been transformed, and the packs referenced within have been expanded. Only RebuildLambdaInfo needs to add the expanded captures to the scope, without requiring addInstantiatedCapturesToScope from LambdaScopeForCallOperatorInstantiationRAII.

Considerations

An alternative approach could involve adding a data structure within the lambda to record the expanded size of the init capture pack. However, this would increase the lambda's size and require extensive modifications.

This PR is a prerequisite for implmenting #61426

[llvmbot]

@llvm/pr-subscribers-clang

Author: Yupei Liu (LYP951018)

Changes

This PR addresses issues related to the handling of init capture with parameter packs in Clang's LambdaScopeForCallOperatorInstantiationRAII.

Previously, addInstantiatedCapturesToScope would add init capture containing packs to the scope using the type of the init capture to determine the expanded pack size. However, this approach resulted in a pack size of 0 because getType()->containsUnexpandedParameterPack() returns false. After extensive testing, it appears that the correct pack size can only be inferred from getInit.

But getInit may reference parameters and init capture from an outer lambda, as shown in the following example:

auto L = [](auto... z) {
    return [... w = z](auto... y) {
        // ...
    };
};

To address this, addInstantiatedCapturesToScope in LambdaScopeForCallOperatorInstantiationRAII should be called last. Additionally, addInstantiatedCapturesToScope has been modified to only add init capture to the scope. The previous implementation incorrectly called MakeInstantiatedLocalArgPack for other non-init captures containing packs, resulting in a pack size of 0.

Impact

This patch affects scenarios where LambdaScopeForCallOperatorInstantiationRAII is passed with ShouldAddDeclsFromParentScope = false, preventing the correct addition of the current lambda's init capture to the scope. There are two main scenarios for ShouldAddDeclsFromParentScope = false:

1. Constraints: Sometimes constraints are instantiated in place rather than delayed. In this case, LambdaScopeForCallOperatorInstantiationRAII does not need to add init capture to the scope.
2. noexcept Expressions: The expressions inside noexcept have already been transformed, and the packs referenced within have been expanded. Only RebuildLambdaInfo needs to add the expanded captures to the scope, without requiring addInstantiatedCapturesToScope from LambdaScopeForCallOperatorInstantiationRAII.

Considerations

An alternative approach could involve adding a data structure within the lambda to record the expanded size of the init capture pack. However, this would increase the lambda's size and require extensive modifications.

This PR is a prerequisite for implmenting #61426

---

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

5 Files Affected:

• (modified) clang/include/clang/Sema/Sema.h (+4)
• (modified) clang/lib/Sema/SemaConcept.cpp (+13-5)
• (modified) clang/lib/Sema/SemaLambda.cpp (+7-2)
• (modified) clang/lib/Sema/SemaTemplateVariadic.cpp (+11-6)
• (modified) clang/test/SemaTemplate/concepts-lambda.cpp (+28)

diff --git a/clang/include/clang/Sema/Sema.h b/clang/include/clang/Sema/Sema.h
index 2ec6367eccea0..d8f903bc05eb6 100644
--- a/clang/include/clang/Sema/Sema.h
+++ b/clang/include/clang/Sema/Sema.h
@@ -14186,6 +14186,10 @@ class Sema final : public SemaBase {
   std::optional<unsigned> getNumArgumentsInExpansion(
       QualType T, const MultiLevelTemplateArgumentList &TemplateArgs);
 
+  std::optional<unsigned> getNumArgumentsInExpansionFromUnexpanded(
+      llvm::ArrayRef<UnexpandedParameterPack> Unexpanded,
+      const MultiLevelTemplateArgumentList &TemplateArgs);
+
   /// Determine whether the given declarator contains any unexpanded
   /// parameter packs.
   ///
diff --git a/clang/lib/Sema/SemaConcept.cpp b/clang/lib/Sema/SemaConcept.cpp
index 9e16b67284be4..49065af10b452 100644
--- a/clang/lib/Sema/SemaConcept.cpp
+++ b/clang/lib/Sema/SemaConcept.cpp
@@ -712,8 +712,8 @@ bool Sema::addInstantiatedCapturesToScope(
   auto AddSingleCapture = [&](const ValueDecl *CapturedPattern,
                               unsigned Index) {
     ValueDecl *CapturedVar = LambdaClass->getCapture(Index)->getCapturedVar();
-    if (CapturedVar->isInitCapture())
-      Scope.InstantiatedLocal(CapturedPattern, CapturedVar);
+    assert(CapturedVar->isInitCapture());
+    Scope.InstantiatedLocal(CapturedPattern, CapturedVar);
   };
 
   for (const LambdaCapture &CapturePattern : LambdaPattern->captures()) {
@@ -721,13 +721,21 @@ bool Sema::addInstantiatedCapturesToScope(
       Instantiated++;
       continue;
     }
-    const ValueDecl *CapturedPattern = CapturePattern.getCapturedVar();
+    ValueDecl *CapturedPattern = CapturePattern.getCapturedVar();
+
+    if (!CapturedPattern->isInitCapture()) {
+      continue;
+    }
+
     if (!CapturedPattern->isParameterPack()) {
       AddSingleCapture(CapturedPattern, Instantiated++);
     } else {
       Scope.MakeInstantiatedLocalArgPack(CapturedPattern);
-      std::optional<unsigned> NumArgumentsInExpansion =
-          getNumArgumentsInExpansion(CapturedPattern->getType(), TemplateArgs);
+      SmallVector<UnexpandedParameterPack, 2> Unexpanded;
+      SemaRef.collectUnexpandedParameterPacks(
+          dyn_cast<VarDecl>(CapturedPattern)->getInit(), Unexpanded);
+      auto NumArgumentsInExpansion =
+          getNumArgumentsInExpansionFromUnexpanded(Unexpanded, TemplateArgs);
       if (!NumArgumentsInExpansion)
         continue;
       for (unsigned Arg = 0; Arg < *NumArgumentsInExpansion; ++Arg)
diff --git a/clang/lib/Sema/SemaLambda.cpp b/clang/lib/Sema/SemaLambda.cpp
index 601077e9f3334..f39d35df56ed6 100644
--- a/clang/lib/Sema/SemaLambda.cpp
+++ b/clang/lib/Sema/SemaLambda.cpp
@@ -2389,11 +2389,10 @@ Sema::LambdaScopeForCallOperatorInstantiationRAII::
   if (!FDPattern)
     return;
 
-  SemaRef.addInstantiatedCapturesToScope(FD, FDPattern, Scope, MLTAL);
-
   if (!ShouldAddDeclsFromParentScope)
     return;
 
+  FunctionDecl *OutermostFD = FD, *OutermostFDPattern = FDPattern;
   llvm::SmallVector<std::pair<FunctionDecl *, FunctionDecl *>, 4>
       ParentInstantiations;
   while (true) {
@@ -2417,5 +2416,11 @@ Sema::LambdaScopeForCallOperatorInstantiationRAII::
   for (const auto &[FDPattern, FD] : llvm::reverse(ParentInstantiations)) {
     SemaRef.addInstantiatedParametersToScope(FD, FDPattern, Scope, MLTAL);
     SemaRef.addInstantiatedLocalVarsToScope(FD, FDPattern, Scope);
+
+    if (isLambdaCallOperator(FD))
+      SemaRef.addInstantiatedCapturesToScope(FD, FDPattern, Scope, MLTAL);
   }
+
+  SemaRef.addInstantiatedCapturesToScope(OutermostFD, OutermostFDPattern, Scope,
+                                         MLTAL);
 }
diff --git a/clang/lib/Sema/SemaTemplateVariadic.cpp b/clang/lib/Sema/SemaTemplateVariadic.cpp
index 3d4ccaf68c700..d9886bd06fee4 100644
--- a/clang/lib/Sema/SemaTemplateVariadic.cpp
+++ b/clang/lib/Sema/SemaTemplateVariadic.cpp
@@ -825,12 +825,9 @@ bool Sema::CheckParameterPacksForExpansion(
   return false;
 }
 
-std::optional<unsigned> Sema::getNumArgumentsInExpansion(
-    QualType T, const MultiLevelTemplateArgumentList &TemplateArgs) {
-  QualType Pattern = cast<PackExpansionType>(T)->getPattern();
-  SmallVector<UnexpandedParameterPack, 2> Unexpanded;
-  CollectUnexpandedParameterPacksVisitor(Unexpanded).TraverseType(Pattern);
-
+std::optional<unsigned> Sema::getNumArgumentsInExpansionFromUnexpanded(
+    llvm::ArrayRef<UnexpandedParameterPack> Unexpanded,
+    const MultiLevelTemplateArgumentList &TemplateArgs) {
   std::optional<unsigned> Result;
   for (unsigned I = 0, N = Unexpanded.size(); I != N; ++I) {
     // Compute the depth and index for this parameter pack.
@@ -878,6 +875,14 @@ std::optional<unsigned> Sema::getNumArgumentsInExpansion(
   return Result;
 }
 
+std::optional<unsigned> Sema::getNumArgumentsInExpansion(
+    QualType T, const MultiLevelTemplateArgumentList &TemplateArgs) {
+  QualType Pattern = cast<PackExpansionType>(T)->getPattern();
+  SmallVector<UnexpandedParameterPack, 2> Unexpanded;
+  CollectUnexpandedParameterPacksVisitor(Unexpanded).TraverseType(Pattern);
+  return getNumArgumentsInExpansionFromUnexpanded(Unexpanded, TemplateArgs);
+}
+
 bool Sema::containsUnexpandedParameterPacks(Declarator &D) {
   const DeclSpec &DS = D.getDeclSpec();
   switch (DS.getTypeSpecType()) {
diff --git a/clang/test/SemaTemplate/concepts-lambda.cpp b/clang/test/SemaTemplate/concepts-lambda.cpp
index 252ef08549a48..9c5807bbabdcb 100644
--- a/clang/test/SemaTemplate/concepts-lambda.cpp
+++ b/clang/test/SemaTemplate/concepts-lambda.cpp
@@ -251,3 +251,31 @@ void dependent_param() {
   L(0, 1)(1, 2)(1);
 }
 } // namespace dependent_param_concept
+
+namespace init_captures {
+template <int N> struct V {};
+
+void sink(V<0>, V<1>, V<2>, V<3>, V<4>) {}
+
+void init_capture_pack() {
+  auto L = [](auto... z) {
+    return [=](auto... y) {
+      return [... w = z, y...](auto)
+        requires requires { sink(w..., y...); }
+      {};
+    };
+  };
+  L(V<0>{}, V<1>{}, V<2>{})(V<3>{}, V<4>{})(1);
+}
+
+void dependent_capture_packs() {
+  auto L = [](auto... z) {
+    return [... w = z](auto... y) {
+      return [... c = w](auto)
+        requires requires { sink(c..., y...); }
+      {};
+    };
+  };
+  L(V<0>{}, V<1>{}, V<2>{})(V<3>{}, V<4>{})(1);
+}
+} // namespace init_captures

[LYP951018]

ping~

[cor3ntin]

When is CapturedPattern not an initCapture?

[LYP951018]

captures like

[a]() {}

[cor3ntin]

Hum, nvm, I misread that code , I failed to realize the assert L715 was in a different scope, so I was confused.

[cor3ntin]

LGTM

[cor3ntin]

Hum, nvm, I misread that code , I failed to realize the assert L715 was in a different scope, so I was confused.

[LYP951018]

Thanks for the review ;)

[zwuis]

Why not use cast?