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

Author: Sirraide

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -3712,6 +3712,8 @@ def err_conflicting_codeseg_attribute : Error<
 def warn_duplicate_codeseg_attribute : Warning<
   "duplicate code segment specifiers">, InGroup<Section>;
 
+def err_expansion_stmt_invalid_init : Error<
+  "cannot expand expression of type %0">;
 def err_expansion_stmt_vla : Error<
   "cannot expand variable length array type %0">;
 def err_expansion_stmt_incomplete : Error<

clang/include/clang/Sema/Sema.h

@@ -15713,6 +15713,9 @@ class Sema final : public SemaBase {
 
   ExprResult BuildCXXExpansionSelectExpr(InitListExpr *Range, Expr *Idx);
 
+  ExprResult BuildCXXDestructuringExpansionSelectExpr(DecompositionDecl *DD,
+                                                      Expr *Idx);
+
   std::optional<uint64_t>
   ComputeExpansionSize(CXXExpansionStmtPattern *Expansion);
   ///@}

clang/lib/Sema/SemaExpand.cpp

@@ -220,6 +220,52 @@ TryBuildIterableExpansionStmtInitializer(Sema &S, Expr *ExpansionInitializer,
   return Data;
 }
 
+static StmtResult BuildDestructuringCXXExpansionStmt(
+    Sema &S, Expr *ExpansionInitializer, SourceLocation ColonLoc,
+    bool VarIsConstexpr,
+    ArrayRef<MaterializeTemporaryExpr *> LifetimeExtendTemps) {
+  auto Ctx = Sema::ExpressionEvaluationContext::PotentiallyEvaluated;
+  if (VarIsConstexpr)
+    Ctx = Sema::ExpressionEvaluationContext::ImmediateFunctionContext;
+  EnterExpressionEvaluationContext ExprEvalCtx(S, Ctx);
+
+  // The declarations should be attached to the parent decl context.
+  Sema::ContextRAII CtxGuard(
+      S, S.CurContext->getEnclosingNonExpansionStatementContext(),
+      /*NewThis=*/false);
+
+  UnsignedOrNone Arity =
+      S.GetDecompositionElementCount(ExpansionInitializer->getType(), ColonLoc);
+
+  if (!Arity) {
+    S.Diag(ExpansionInitializer->getBeginLoc(),
+           diag::err_expansion_stmt_invalid_init)
+        << ExpansionInitializer->getType()
+        << ExpansionInitializer->getSourceRange();
+    return StmtError();
+  }
+
+  QualType AutoRRef = S.Context.getAutoRRefDeductType();
+  SmallVector<BindingDecl *> Bindings;
+  for (unsigned I = 0; I < *Arity; ++I)
+    Bindings.push_back(BindingDecl::Create(
+        S.Context, S.CurContext, ColonLoc,
+        S.getPreprocessor().getIdentifierInfo("__u" + std::to_string(I)),
+        AutoRRef));
+
+  TypeSourceInfo *TSI = S.Context.getTrivialTypeSourceInfo(AutoRRef);
+  auto *DD =
+      DecompositionDecl::Create(S.Context, S.CurContext, ColonLoc, ColonLoc,
+                                AutoRRef, TSI, SC_Auto, Bindings);
+
+  if (VarIsConstexpr)
+    DD->setConstexpr(true);
+
+  S.ApplyForRangeOrExpansionStatementLifetimeExtension(DD, LifetimeExtendTemps);
+  S.AddInitializerToDecl(DD, ExpansionInitializer, false);
+  return S.ActOnDeclStmt(S.ConvertDeclToDeclGroup(DD), ColonLoc, ColonLoc);
+}
+
 CXXExpansionStmtDecl *
 Sema::ActOnCXXExpansionStmtDecl(unsigned TemplateDepth,
                                 SourceLocation TemplateKWLoc) {
@@ -368,8 +414,31 @@ StmtResult Sema::BuildNonEnumeratingCXXExpansionStmtPattern(
         Data.EndDecl, LParenLoc, ColonLoc, RParenLoc);
   }
 
-  Diag(ESD->getLocation(), diag::err_expansion_statements_todo);
-  return StmtError();
+  // If not, try destructuring.
+  StmtResult DecompDeclStmt = BuildDestructuringCXXExpansionStmt(
+      *this, ExpansionInitializer, ColonLoc, ExpansionVar->isConstexpr(),
+      LifetimeExtendTemps);
+  if (DecompDeclStmt.isInvalid()) {
+    ActOnInitializerError(ExpansionVar);
+    return StmtError();
+  }
+
+  auto *DS = DecompDeclStmt.getAs<DeclStmt>();
+  auto *DD = cast<DecompositionDecl>(DS->getSingleDecl());
+  if (DD->isInvalidDecl())
+    return StmtError();
+
+  ExprResult Select = BuildCXXDestructuringExpansionSelectExpr(DD, Index);
+  if (Select.isInvalid()) {
+    ActOnInitializerError(ExpansionVar);
+    return StmtError();
+  }
+
+  if (FinaliseExpansionVar(*this, ExpansionVar, Select))
+    return StmtError();
+
+  return new (Context) CXXDestructuringExpansionStmtPattern(
+      ESD, Init, ExpansionVarStmt, DS, LParenLoc, ColonLoc, RParenLoc);
 }
 
 StmtResult Sema::FinishCXXExpansionStmt(Stmt *Exp, Stmt *Body) {
@@ -409,8 +478,8 @@ StmtResult Sema::FinishCXXExpansionStmt(Stmt *Exp, Stmt *Body) {
     Shared.push_back(Expansion->getRangeVarStmt());
     Shared.push_back(Expansion->getBeginVarStmt());
     Shared.push_back(Expansion->getEndVarStmt());
-  } else {
-    assert(Expansion->isEnumerating() && "TODO");
+  } else if (Expansion->isDestructuring()) {
+    Shared.push_back(Expansion->getDecompositionDeclStmt());
   }
 
   // Return an empty statement if the range is empty.
@@ -471,6 +540,23 @@ ExprResult Sema::BuildCXXExpansionSelectExpr(InitListExpr *Range, Expr *Idx) {
   return Range->getInit(I);
 }
 
+ExprResult Sema::BuildCXXDestructuringExpansionSelectExpr(DecompositionDecl *DD,
+                                                          Expr *Idx) {
+  if (Idx->isValueDependent())
+    return new (Context) CXXDestructuringExpansionSelectExpr(Context, DD, Idx);
+
+  Expr::EvalResult ER;
+  if (!Idx->EvaluateAsInt(ER, Context))
+    llvm_unreachable("Failed to evaluate expansion index");
+
+  uint64_t I = ER.Val.getInt().getZExtValue();
+  MarkAnyDeclReferenced(Idx->getBeginLoc(), DD, true);
+  if (auto *BD = DD->bindings()[I]; auto *HVD = BD->getHoldingVar())
+    return HVD->getInit();
+  else
+    return BD->getBinding();
+}
+
 std::optional<uint64_t>
 Sema::ComputeExpansionSize(CXXExpansionStmtPattern *Expansion) {
   assert(!HasDependentSize(Expansion));
@@ -663,5 +749,9 @@ Sema::ComputeExpansionSize(CXXExpansionStmtPattern *Expansion) {
     return ER.Val.getInt().getZExtValue();
   }
 
-  llvm_unreachable("TODO");
+  if (auto *Destructuring =
+          dyn_cast<CXXDestructuringExpansionStmtPattern>(Expansion))
+    return Destructuring->getDecompositionDecl()->bindings().size();
+
+  llvm_unreachable("Invalid expansion statement class");
 }

clang/lib/Sema/TreeTransform.h

@@ -9297,28 +9297,51 @@ TreeTransform<Derived>::TransformCXXForRangeStmt(CXXForRangeStmt *S) {
 template <typename Derived>
 StmtResult TreeTransform<Derived>::TransformCXXExpansionStmtPattern(
     CXXExpansionStmtPattern *S) {
+  ExprResult ExpansionInitializer;
+  SmallVector<MaterializeTemporaryExpr *, 8> LifetimeExtendTemps;
   CXXExpansionStmtDecl *NewESD = nullptr;
   Stmt *Init = nullptr;
   DeclStmt *ExpansionVarStmt = nullptr;
-  Decl *ESD =
-      getDerived().TransformDecl(S->getDecl()->getLocation(), S->getDecl());
-  if (!ESD || ESD->isInvalidDecl())
-    return StmtError();
-  NewESD = cast<CXXExpansionStmtDecl>(ESD);
 
-  Init = S->getInit();
-  if (Init) {
-    StmtResult SR = getDerived().TransformStmt(Init);
-    if (SR.isInvalid())
+  // Collect lifetime-extended temporaries in case this ends up being a
+  // destructuring expansion statement (for other kinds of expansion statements,
+  // this should make no difference since we ignore 'LifetimeExtendTemps' for
+  // those).
+  {
+    EnterExpressionEvaluationContext ExprEvalCtx(
+        SemaRef, SemaRef.currentEvaluationContext().Context);
+    SemaRef.currentEvaluationContext().InLifetimeExtendingContext = true;
+    SemaRef.currentEvaluationContext().RebuildDefaultArgOrDefaultInit = true;
+    Decl *ESD =
+        getDerived().TransformDecl(S->getDecl()->getLocation(), S->getDecl());
+    if (!ESD || ESD->isInvalidDecl())
       return StmtError();
-    Init = SR.get();
-  }
+    NewESD = cast<CXXExpansionStmtDecl>(ESD);
 
-  StmtResult ExpansionVar =
-      getDerived().TransformStmt(S->getExpansionVarStmt());
-  if (ExpansionVar.isInvalid())
-    return StmtError();
-  ExpansionVarStmt = cast<DeclStmt>(ExpansionVar.get());
+    Init = S->getInit();
+    if (Init) {
+      StmtResult SR = getDerived().TransformStmt(Init);
+      if (SR.isInvalid())
+        return StmtError();
+      Init = SR.get();
+    }
+
+    StmtResult ExpansionVar =
+        getDerived().TransformStmt(S->getExpansionVarStmt());
+    if (ExpansionVar.isInvalid())
+      return StmtError();
+    ExpansionVarStmt = cast<DeclStmt>(ExpansionVar.get());
+
+    if (S->isDependent()) {
+      ExpansionInitializer =
+          getDerived().TransformExpr(S->getExpansionInitializer());
+      if (ExpansionInitializer.isInvalid())
+        return StmtError();
+
+      LifetimeExtendTemps =
+          SemaRef.currentEvaluationContext().ForRangeLifetimeExtendTemps;
+    }
+  }
 
   CXXExpansionStmtPattern *NewPattern = nullptr;
   if (S->isEnumerating()) {
@@ -9355,7 +9378,12 @@ StmtResult TreeTransform<Derived>::TransformCXXExpansionStmtPattern(
 
     NewPattern = cast<CXXExpansionStmtPattern>(Res.get());
   } else {
-    llvm_unreachable("TODO");
+    // The only time we instantiate an expansion statement is if its expansion
+    // size is dependent (otherwise, we only instantiate the expansions and
+    // leave the underlying CXXExpansionStmtPattern as-is). Since destructuring
+    // expansion statements never have a dependent size, we should never get
+    // here.
+    llvm_unreachable("destructuring pattern should never be instantiated");
   }
 
   StmtResult Body = getDerived().TransformStmt(S->getBody());
@@ -9386,8 +9414,23 @@ StmtResult TreeTransform<Derived>::TransformCXXExpansionStmtInstantiation(
   SmallVector<Stmt *> SharedStmts;
   SmallVector<Stmt *> Instantiations;
 
-  if (TransformStmts(SharedStmts, S->getSharedStmts()))
-    return StmtError();
+  // Apply lifetime extension to the shared statements if this was a
+  // destructuring expansion statement.
+  {
+    EnterExpressionEvaluationContext ExprEvalCtx(
+        SemaRef, SemaRef.currentEvaluationContext().Context);
+    SemaRef.currentEvaluationContext().InLifetimeExtendingContext = true;
+    SemaRef.currentEvaluationContext().RebuildDefaultArgOrDefaultInit = true;
+    if (TransformStmts(SharedStmts, S->getSharedStmts()))
+      return StmtError();
+
+    if (S->shouldApplyLifetimeExtensionToSharedStmts()) {
+      auto *VD =
+          cast<VarDecl>(cast<DeclStmt>(SharedStmts.front())->getSingleDecl());
+      SemaRef.ApplyForRangeOrExpansionStatementLifetimeExtension(
+          VD, SemaRef.currentEvaluationContext().ForRangeLifetimeExtendTemps);
+    }
+  }
 
   if (TransformStmts(Instantiations, S->getInstantiations()))
     return StmtError();