[Clang] [C++26] Expansion Statements (Part 3)

Author: Sirraide

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -5827,6 +5827,8 @@ def note_template_nsdmi_here : Note<
   "in instantiation of default member initializer %q0 requested here">;
 def note_template_type_alias_instantiation_here : Note<
   "in instantiation of template type alias %0 requested here">;
+def note_expansion_stmt_instantiation_here : Note<
+  "in instantiation of expansion statement requested here">;
 def note_template_exception_spec_instantiation_here : Note<
   "in instantiation of exception specification for %0 requested here">;
 def note_template_requirement_instantiation_here : Note<

clang/include/clang/Sema/Sema.h

@@ -13176,6 +13176,9 @@ class Sema final : public SemaBase {
 
       /// We are performing partial ordering for template template parameters.
       PartialOrderingTTP,
+
+      /// We are instantiating an expansion statement.
+      ExpansionStmtInstantiation,
     } Kind;
 
     /// Whether we're substituting into constraints.
@@ -13371,6 +13374,12 @@ class Sema final : public SemaBase {
                           concepts::Requirement *Req,
                           SourceRange InstantiationRange = SourceRange());
 
+    /// \brief Note that we are substituting the body of an expansion statement.
+    InstantiatingTemplate(Sema &SemaRef, SourceLocation PointOfInstantiation,
+                          CXXExpansionStmtPattern *ExpansionStmt,
+                          ArrayRef<TemplateArgument> TArgs,
+                          SourceRange InstantiationRange);
+
     /// \brief Note that we are checking the satisfaction of the constraint
     /// expression inside of a nested requirement.
     InstantiatingTemplate(Sema &SemaRef, SourceLocation PointOfInstantiation,
@@ -15643,6 +15652,19 @@ class Sema final : public SemaBase {
       ArrayRef<MaterializeTemporaryExpr *> LifetimeExtendTemps);
 
   StmtResult FinishCXXExpansionStmt(Stmt *Expansion, Stmt *Body);
+
+  StmtResult BuildCXXEnumeratingExpansionStmtPattern(Decl *ESD, Stmt *Init,
+                                                     Stmt *ExpansionVar,
+                                                     SourceLocation LParenLoc,
+                                                     SourceLocation ColonLoc,
+                                                     SourceLocation RParenLoc);
+
+  ExprResult
+  BuildCXXExpansionInitListSelectExpr(CXXExpansionInitListExpr *Range,
+                                      Expr *Idx);
+
+  std::optional<uint64_t>
+  ComputeExpansionSize(CXXExpansionStmtPattern *Expansion);
   ///@}
 };
 

clang/lib/Frontend/FrontendActions.cpp

@@ -476,6 +476,8 @@ class DefaultTemplateInstCallback : public TemplateInstantiationCallback {
       return "TypeAliasTemplateInstantiation";
     case CodeSynthesisContext::PartialOrderingTTP:
       return "PartialOrderingTTP";
+    case CodeSynthesisContext::ExpansionStmtInstantiation:
+      return "ExpansionStmtInstantiation";
     }
     return "";
   }

clang/lib/Sema/SemaExpand.cpp

@@ -24,6 +24,23 @@
 using namespace clang;
 using namespace sema;
 
+// Build a 'DeclRefExpr' designating the template parameter '__N'.
+static DeclRefExpr *BuildIndexDRE(Sema &S, CXXExpansionStmtDecl *ESD) {
+  return S.BuildDeclRefExpr(ESD->getIndexTemplateParm(),
+                            S.Context.getPointerDiffType(), VK_PRValue,
+                            ESD->getBeginLoc());
+}
+
+static bool FinaliseExpansionVar(Sema &S, VarDecl *ExpansionVar,
+                                 ExprResult Initializer) {
+  if (Initializer.isInvalid()) {
+    S.ActOnInitializerError(ExpansionVar);
+    return true;
+  }
+
+  S.AddInitializerToDecl(ExpansionVar, Initializer.get(), /*DirectInit=*/false);
+  return ExpansionVar->isInvalidDecl();
+}
 
 CXXExpansionStmtDecl *
 Sema::ActOnCXXExpansionStmtDecl(unsigned TemplateDepth,
@@ -66,13 +83,147 @@ StmtResult Sema::ActOnCXXExpansionStmtPattern(
     Expr *ExpansionInitializer, SourceLocation LParenLoc,
     SourceLocation ColonLoc, SourceLocation RParenLoc,
     ArrayRef<MaterializeTemporaryExpr *> LifetimeExtendTemps) {
+  if (!ExpansionInitializer || !ExpansionVarStmt)
+    return StmtError();
+
+  assert(CurContext->isExpansionStmt());
+  auto *DS = cast<DeclStmt>(ExpansionVarStmt);
+  if (!DS->isSingleDecl()) {
+    Diag(DS->getBeginLoc(), diag::err_type_defined_in_for_range);
+    return StmtError();
+  }
+
+  VarDecl *ExpansionVar = dyn_cast<VarDecl>(DS->getSingleDecl());
+  if (!ExpansionVar || ExpansionVar->isInvalidDecl() ||
+      ExpansionInitializer->containsErrors())
+    return StmtError();
+
+  // This is an enumerating expansion statement.
+  if (auto *ILE = dyn_cast<CXXExpansionInitListExpr>(ExpansionInitializer)) {
+    ExprResult Initializer =
+        BuildCXXExpansionInitListSelectExpr(ILE, BuildIndexDRE(*this, ESD));
+    if (FinaliseExpansionVar(*this, ExpansionVar, Initializer))
+      return StmtError();
+
+    // Note that lifetime extension only applies to destructuring expansion
+    // statements, so we just ignore 'LifetimeExtendedTemps' entirely for other
+    // types of expansion statements (this is CWG 3043).
+    return BuildCXXEnumeratingExpansionStmtPattern(ESD, Init, DS, LParenLoc,
+                                                   ColonLoc, RParenLoc);
+  }
+
   Diag(ESD->getLocation(), diag::err_expansion_statements_todo);
   return StmtError();
 }
 
+StmtResult Sema::BuildCXXEnumeratingExpansionStmtPattern(
+    Decl *ESD, Stmt *Init, Stmt *ExpansionVar, SourceLocation LParenLoc,
+    SourceLocation ColonLoc, SourceLocation RParenLoc) {
+  return new (Context) CXXEnumeratingExpansionStmtPattern(
+      cast<CXXExpansionStmtDecl>(ESD), Init, cast<DeclStmt>(ExpansionVar),
+      LParenLoc, ColonLoc, RParenLoc);
+}
+
 StmtResult Sema::FinishCXXExpansionStmt(Stmt *Exp, Stmt *Body) {
   if (!Exp || !Body)
     return StmtError();
 
+  auto *Expansion = cast<CXXExpansionStmtPattern>(Exp);
+  assert(!Expansion->getDecl()->getInstantiations() &&
+         "should not rebuild expansion statement after instantiation");
+
+  Expansion->setBody(Body);
+  if (Expansion->hasDependentSize())
+    return Expansion;
+
+  // This can fail if this is an iterating expansion statement.
+  std::optional<uint64_t> NumInstantiations = ComputeExpansionSize(Expansion);
+  if (!NumInstantiations)
+    return StmtError();
+
+  // Collect shared statements.
+  SmallVector<Stmt *, 1> Shared;
+  if (Expansion->getInit())
+    Shared.push_back(Expansion->getInit());
+
+  assert(isa<CXXEnumeratingExpansionStmtPattern>(Expansion) && "TODO");
+
+  // Return an empty statement if the range is empty.
+  if (*NumInstantiations == 0) {
+    Expansion->getDecl()->setInstantiations(
+        CXXExpansionStmtInstantiation::Create(
+            Context, Expansion->getBeginLoc(), Expansion->getEndLoc(),
+            /*Instantiations=*/{}, Shared,
+            isa<CXXDestructuringExpansionStmtPattern>(Expansion)));
+    return Expansion;
+  }
+
+  // Create a compound statement binding the expansion variable and body.
+  Stmt *VarAndBody[] = {Expansion->getExpansionVarStmt(), Body};
+  Stmt *CombinedBody =
+      CompoundStmt::Create(Context, VarAndBody, FPOptionsOverride(),
+                           Body->getBeginLoc(), Body->getEndLoc());
+
+  // Expand the body for each instantiation.
+  SmallVector<Stmt *, 4> Instantiations;
+  CXXExpansionStmtDecl *ESD = Expansion->getDecl();
+  for (uint64_t I = 0; I < *NumInstantiations; ++I) {
+    // Now that we're expanding this, exit the context of the expansion stmt
+    // so that we no longer treat this as dependent.
+    ContextRAII CtxGuard(*this, CurContext->getParent(),
+                         /*NewThis=*/false);
+
+    TemplateArgument Arg{Context, llvm::APSInt::get(I),
+                         Context.getPointerDiffType()};
+    MultiLevelTemplateArgumentList MTArgList(ESD, Arg, true);
+    MTArgList.addOuterRetainedLevels(
+        Expansion->getDecl()->getIndexTemplateParm()->getDepth());
+
+    LocalInstantiationScope LIScope(*this, /*CombineWithOuterScope=*/true);
+    NonSFINAEContext _(*this);
+    InstantiatingTemplate Inst(*this, Body->getBeginLoc(), Expansion, Arg,
+                               Body->getSourceRange());
+
+    StmtResult Instantiation = SubstStmt(CombinedBody, MTArgList);
+    if (Instantiation.isInvalid())
+      return StmtError();
+    Instantiations.push_back(Instantiation.get());
+  }
+
+  auto *InstantiationsStmt = CXXExpansionStmtInstantiation::Create(
+      Context, Expansion->getBeginLoc(), Expansion->getEndLoc(), Instantiations,
+      Shared, isa<CXXDestructuringExpansionStmtPattern>(Expansion));
+
+  Expansion->getDecl()->setInstantiations(InstantiationsStmt);
+  return Expansion;
+}
+
+ExprResult
+Sema::BuildCXXExpansionInitListSelectExpr(CXXExpansionInitListExpr *Range,
+                                          Expr *Idx) {
+  if (Range->containsPackExpansion() || Idx->isValueDependent())
+    return new (Context) CXXExpansionInitListSelectExpr(Context, Range, Idx);
+
+  // The index is a DRE to a template parameter; we should never
+  // fail to evaluate it.
+  Expr::EvalResult ER;
+  if (!Idx->EvaluateAsInt(ER, Context))
+    llvm_unreachable("Failed to evaluate expansion index");
+
+  uint64_t I = ER.Val.getInt().getZExtValue();
+  return Range->getExprs()[I];
+}
+
+std::optional<uint64_t>
+Sema::ComputeExpansionSize(CXXExpansionStmtPattern *Expansion) {
+  assert(!Expansion->hasDependentSize());
+
+  if (isa<CXXEnumeratingExpansionStmtPattern>(Expansion))
+    return cast<CXXExpansionInitListSelectExpr>(
+               Expansion->getExpansionVariable()->getInit())
+        ->getRangeExpr()
+        ->getExprs()
+        .size();
+
   llvm_unreachable("TODO");
 }

clang/lib/Sema/SemaTemplateInstantiate.cpp

@@ -573,6 +573,7 @@ bool Sema::CodeSynthesisContext::isInstantiationRecord() const {
   case PriorTemplateArgumentSubstitution:
   case ConstraintsCheck:
   case NestedRequirementConstraintsCheck:
+  case ExpansionStmtInstantiation:
     return true;
 
   case RequirementInstantiation:
@@ -759,6 +760,15 @@ Sema::InstantiatingTemplate::InstantiatingTemplate(
           PointOfInstantiation, InstantiationRange, /*Entity=*/nullptr,
           /*Template=*/nullptr, /*TemplateArgs=*/{}) {}
 
+Sema::InstantiatingTemplate::InstantiatingTemplate(
+    Sema &SemaRef, SourceLocation PointOfInstantiation,
+    CXXExpansionStmtPattern *ExpansionStmt, ArrayRef<TemplateArgument> TArgs,
+    SourceRange InstantiationRange)
+    : InstantiatingTemplate(
+          SemaRef, CodeSynthesisContext::ExpansionStmtInstantiation,
+          PointOfInstantiation, InstantiationRange, /*Entity=*/nullptr,
+          /*Template=*/nullptr, /*TemplateArgs=*/TArgs) {}
+
 Sema::InstantiatingTemplate::InstantiatingTemplate(
     Sema &SemaRef, SourceLocation PointOfInstantiation,
     concepts::NestedRequirement *Req, ConstraintsCheck,
@@ -1260,6 +1270,9 @@ void Sema::PrintInstantiationStack(InstantiationContextDiagFuncRef DiagFunc) {
                                << /*isTemplateTemplateParam=*/true
                                << Active->InstantiationRange);
       break;
+    case CodeSynthesisContext::ExpansionStmtInstantiation:
+      Diags.Report(Active->PointOfInstantiation,
+                   diag::note_expansion_stmt_instantiation_here);
     }
   }
 }
@@ -1894,6 +1907,12 @@ Decl *TemplateInstantiator::TransformDecl(SourceLocation Loc, Decl *D) {
       maybeInstantiateFunctionParameterToScope(PVD))
     return nullptr;
 
+  if (isa<CXXExpansionStmtDecl>(D)) {
+    assert(SemaRef.CurrentInstantiationScope);
+    return cast<Decl *>(
+        *SemaRef.CurrentInstantiationScope->findInstantiationOf(D));
+  }
+
   return SemaRef.FindInstantiatedDecl(Loc, cast<NamedDecl>(D), TemplateArgs);
 }
 
@@ -2288,9 +2307,13 @@ ExprResult
 TemplateInstantiator::TransformFunctionParmPackRefExpr(DeclRefExpr *E,
                                                        ValueDecl *PD) {
   typedef LocalInstantiationScope::DeclArgumentPack DeclArgumentPack;
-  llvm::PointerUnion<Decl *, DeclArgumentPack *> *Found
-    = getSema().CurrentInstantiationScope->findInstantiationOf(PD);
-  assert(Found && "no instantiation for parameter pack");
+  llvm::PointerUnion<Decl *, DeclArgumentPack *> *Found =
+      getSema().CurrentInstantiationScope->getInstantiationOfIfExists(PD);
+
+  // This can happen when instantiating an expansion statement that contains
+  // a pack (e.g. `template for (auto x : {{ts...}})`).
+  if (!Found)
+    return E;
 
   Decl *TransformedDecl;
   if (DeclArgumentPack *Pack = dyn_cast<DeclArgumentPack *>(*Found)) {

clang/lib/Sema/SemaTemplateInstantiateDecl.cpp

@@ -2089,7 +2089,37 @@ Decl *TemplateDeclInstantiator::VisitStaticAssertDecl(StaticAssertDecl *D) {
 
 Decl *TemplateDeclInstantiator::VisitCXXExpansionStmtDecl(
     CXXExpansionStmtDecl *OldESD) {
-  llvm_unreachable("TODO");
+  Decl *Index = VisitNonTypeTemplateParmDecl(OldESD->getIndexTemplateParm());
+  CXXExpansionStmtDecl *NewESD = SemaRef.BuildCXXExpansionStmtDecl(
+      Owner, OldESD->getBeginLoc(), cast<NonTypeTemplateParmDecl>(Index));
+  SemaRef.CurrentInstantiationScope->InstantiatedLocal(OldESD, NewESD);
+
+  // If this was already expanded, only instantiate the expansion and
+  // don't touch the unexpanded expansion statement.
+  if (CXXExpansionStmtInstantiation *OldInst = OldESD->getInstantiations()) {
+    StmtResult NewInst = SemaRef.SubstStmt(OldInst, TemplateArgs);
+    if (NewInst.isInvalid())
+      return nullptr;
+
+    NewESD->setInstantiations(NewInst.getAs<CXXExpansionStmtInstantiation>());
+    NewESD->setExpansionPattern(OldESD->getExpansionPattern());
+    return NewESD;
+  }
+
+  // Enter the scope of this expansion statement; don't do this if we've
+  // already expanded it, as in that case we no longer want to treat its
+  // content as dependent.
+  Sema::ContextRAII Context(SemaRef, NewESD, /*NewThis=*/false);
+
+  StmtResult Expansion =
+      SemaRef.SubstStmt(OldESD->getExpansionPattern(), TemplateArgs);
+  if (Expansion.isInvalid())
+    return nullptr;
+
+  // The code that handles CXXExpansionStmtPattern takes care of calling
+  // setInstantiation() on the ESD if there was an expansion.
+  NewESD->setExpansionPattern(cast<CXXExpansionStmtPattern>(Expansion.get()));
+  return NewESD;
 }
 
 Decl *TemplateDeclInstantiator::VisitEnumDecl(EnumDecl *D) {
@@ -7093,6 +7123,12 @@ NamedDecl *Sema::FindInstantiatedDecl(SourceLocation Loc, NamedDecl *D,
     // anonymous unions in class templates).
   }
 
+  if (CurrentInstantiationScope) {
+    if (auto Found = CurrentInstantiationScope->getInstantiationOfIfExists(D))
+      if (auto *FD = dyn_cast<NamedDecl>(cast<Decl *>(*Found)))
+        return FD;
+  }
+
   if (!ParentDependsOnArgs)
     return D;
 

clang/lib/Sema/SemaTemplateVariadic.cpp

@@ -912,10 +912,14 @@ bool Sema::CheckParameterPacksForExpansion(
     unsigned NewPackSize, PendingPackExpansionSize = 0;
     if (IsVarDeclPack) {
       // Figure out whether we're instantiating to an argument pack or not.
+      //
+      // The instantiation may not exist; this can happen when instantiating an
+      // expansion statement that contains a pack (e.g.
+      // `template for (auto x : {{ts...}})`).
       llvm::PointerUnion<Decl *, DeclArgumentPack *> *Instantiation =
-          CurrentInstantiationScope->findInstantiationOf(
+          CurrentInstantiationScope->getInstantiationOfIfExists(
               cast<NamedDecl *>(ParmPack.first));
-      if (isa<DeclArgumentPack *>(*Instantiation)) {
+      if (Instantiation && isa<DeclArgumentPack *>(*Instantiation)) {
         // We could expand this function parameter pack.
         NewPackSize = cast<DeclArgumentPack *>(*Instantiation)->size();
       } else {