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

Author: Sirraide

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -3702,6 +3702,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

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

clang/lib/Sema/SemaExpand.cpp

@@ -173,6 +173,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) {
@@ -317,8 +363,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) {
@@ -347,8 +416,9 @@ StmtResult Sema::FinishCXXExpansionStmt(Stmt *Exp, Stmt *Body) {
     Shared.push_back(Iter->getRangeVarStmt());
     Shared.push_back(Iter->getBeginVarStmt());
     Shared.push_back(Iter->getEndVarStmt());
-  } else {
-    assert(isa<CXXEnumeratingExpansionStmtPattern>(Expansion) && "TODO");
+  } else if (auto *Destructuring =
+                 dyn_cast<CXXDestructuringExpansionStmtPattern>(Expansion)) {
+    Shared.push_back(Destructuring->getDecompositionDeclStmt());
   }
 
   // Return an empty statement if the range is empty.
@@ -417,6 +487,23 @@ Sema::BuildCXXExpansionInitListSelectExpr(CXXExpansionInitListExpr *Range,
   return Range->getExprs()[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(!Expansion->hasDependentSize());
@@ -476,5 +563,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

@@ -9383,20 +9383,35 @@ StmtResult TreeTransform<Derived>::TransformCXXDependentExpansionStmtPattern(
     CXXDependentExpansionStmtPattern *S) {
   TransformCXXExpansionStmtPatternResult Common;
   ExprResult ExpansionInitializer;
+  SmallVector<MaterializeTemporaryExpr *, 8> LifetimeExtendTemps;
 
-  Common = TransformCXXExpansionStmtPatternCommonParts(S);
-  if (!Common.isValid())
-    return StmtError();
+  // 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;
 
-  ExpansionInitializer =
-      getDerived().TransformExpr(S->getExpansionInitializer());
-  if (ExpansionInitializer.isInvalid())
-    return StmtError();
+    Common = TransformCXXExpansionStmtPatternCommonParts(S);
+    if (!Common.isValid())
+      return StmtError();
+
+    ExpansionInitializer =
+        getDerived().TransformExpr(S->getExpansionInitializer());
+    if (ExpansionInitializer.isInvalid())
+      return StmtError();
+
+    LifetimeExtendTemps =
+        SemaRef.currentEvaluationContext().ForRangeLifetimeExtendTemps;
+  }
 
   StmtResult Expansion = SemaRef.BuildNonEnumeratingCXXExpansionStmtPattern(
       Common.NewESD, Common.NewInit, Common.NewExpansionVarDecl,
       ExpansionInitializer.get(), S->getLParenLoc(), S->getColonLoc(),
-      S->getRParenLoc(), /*LifetimeExtendTemps=*/{});
+      S->getRParenLoc(), LifetimeExtendTemps);
   if (Expansion.isInvalid())
     return StmtError();
 
@@ -9455,8 +9470,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();
@@ -9488,7 +9518,18 @@ ExprResult TreeTransform<Derived>::TransformCXXExpansionInitListSelectExpr(
 template <typename Derived>
 ExprResult TreeTransform<Derived>::TransformCXXDestructuringExpansionSelectExpr(
     CXXDestructuringExpansionSelectExpr *E) {
-  llvm_unreachable("TOOD");
+  Decl *DD = getDerived().TransformDecl(
+      E->getDecompositionDecl()->getLocation(), E->getDecompositionDecl());
+  ExprResult Idx = getDerived().TransformExpr(E->getIndexExpr());
+  if (!DD || Idx.isInvalid())
+    return ExprError();
+
+  if (!getDerived().AlwaysRebuild() && DD == E->getDecompositionDecl() &&
+      Idx.get() == E->getIndexExpr())
+    return E;
+
+  return SemaRef.BuildCXXDestructuringExpansionSelectExpr(
+      cast<DecompositionDecl>(DD), Idx.get());
 }
 
 template<typename Derived>