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

Author: Sirraide

clang/include/clang/AST/ASTNodeTraverser.h

@@ -959,6 +959,12 @@ class ASTNodeTraverser
     }
   }
 
+  void VisitCXXExpansionStmtDecl(const CXXExpansionStmtDecl *Node) {
+    Visit(Node->getExpansionPattern());
+    if (Traversal != TK_IgnoreUnlessSpelledInSource)
+      Visit(Node->getInstantiations());
+  }
+
   void VisitCallExpr(const CallExpr *Node) {
     for (const auto *Child :
          make_filter_range(Node->children(), [this](const Stmt *Child) {

clang/include/clang/AST/ComputeDependence.h

@@ -94,6 +94,7 @@ class DesignatedInitExpr;
 class ParenListExpr;
 class PseudoObjectExpr;
 class AtomicExpr;
+class CXXExpansionInitListExpr;
 class ArraySectionExpr;
 class OMPArrayShapingExpr;
 class OMPIteratorExpr;
@@ -191,6 +192,8 @@ ExprDependence computeDependence(ParenListExpr *E);
 ExprDependence computeDependence(PseudoObjectExpr *E);
 ExprDependence computeDependence(AtomicExpr *E);
 
+ExprDependence computeDependence(CXXExpansionInitListExpr *E);
+
 ExprDependence computeDependence(ArraySectionExpr *E);
 ExprDependence computeDependence(OMPArrayShapingExpr *E);
 ExprDependence computeDependence(OMPIteratorExpr *E);

clang/include/clang/AST/Decl.h

@@ -1254,7 +1254,9 @@ class VarDecl : public DeclaratorDecl, public Redeclarable<VarDecl> {
     if (getKind() != Decl::Var && getKind() != Decl::Decomposition)
       return false;
     if (const DeclContext *DC = getLexicalDeclContext())
-      return DC->getRedeclContext()->isFunctionOrMethod();
+      return DC->getEnclosingNonExpansionStatementContext()
+          ->getRedeclContext()
+          ->isFunctionOrMethod();
     return false;
   }
 

clang/include/clang/AST/DeclBase.h

@@ -2195,6 +2195,10 @@ class DeclContext {
     return getDeclKind() == Decl::RequiresExprBody;
   }
 
+  bool isExpansionStmt() const {
+    return getDeclKind() == Decl::CXXExpansionStmt;
+  }
+
   bool isNamespace() const { return getDeclKind() == Decl::Namespace; }
 
   bool isStdNamespace() const;
@@ -2292,6 +2296,15 @@ class DeclContext {
     return const_cast<DeclContext *>(this)->getOuterLexicalRecordContext();
   }
 
+  /// Retrieve the innermost enclosing context that doesn't belong to an
+  /// expansion statement. Returns 'this' if this context is not an expansion
+  /// statement.
+  DeclContext *getEnclosingNonExpansionStatementContext();
+  const DeclContext *getEnclosingNonExpansionStatementContext() const {
+    return const_cast<DeclContext *>(this)
+        ->getEnclosingNonExpansionStatementContext();
+  }
+
   /// Test if this context is part of the enclosing namespace set of
   /// the context NS, as defined in C++0x [namespace.def]p9. If either context
   /// isn't a namespace, this is equivalent to Equals().

clang/include/clang/AST/DeclTemplate.h

@@ -3343,6 +3343,120 @@ class TemplateParamObjectDecl : public ValueDecl,
   static bool classofKind(Kind K) { return K == TemplateParamObject; }
 };
 
+/// Represents a C++26 expansion statement declaration.
+///
+/// This is a bit of a hack, since expansion statements shouldn't really be
+/// 'declarations' per se (they don't declare anything). Nevertheless, we *do*
+/// need them to be declaration *contexts*, because the DeclContext is used to
+/// compute the 'template depth' of entities enclosed therein. In particular,
+/// the 'template depth' is used to find instantiations of parameter variables,
+/// and a lambda enclosed within an expansion statement cannot compute its
+/// template depth without a pointer to the enclosing expansion statement.
+///
+/// For the remainder of this comment, let 'expanding' an expansion statement
+/// refer to the process of performing template substitution on its body N
+/// times, where N is the expansion size (how this size is determined depends on
+/// the kind of expansion statement); by contrast we may sometimes 'instantiate'
+/// an expansion statement (because it happens to be in a template). This is
+/// just regular template instantiation.
+///
+/// Apart from a template parameter list that contains a template parameter used
+/// as the expansion index, this node contains a 'CXXExpansionStmtPattern' as
+/// well as a 'CXXExpansionStmtInstantiation'. These two members correspond to
+/// distinct representations of the expansion statement: the former is used
+/// prior to expansion and contains all the parts needed to perform expansion;
+/// the latter holds the expanded/desugared AST nodes that result from the
+/// expansion.
+///
+/// After expansion, the 'CXXExpansionStmtPattern' is no longer updated and left
+/// as-is; this also means that, if an already-expanded expansion statement is
+/// inside a template, and that template is then instantiated, the
+/// 'CXXExpansionStmtPattern' is *not* instantiated; only the
+/// 'CXXExpansionStmtInstantiation' is. The latter is also what's used for
+/// codegen and constant evaluation.
+///
+/// For example, if the user writes the following expansion statement:
+/// \verbatim
+///   std::tuple<int, int, int> a{1, 2, 3};
+///   template for (auto x : a) {
+///     // ...
+///   }
+/// \endverbatim
+///
+/// The 'CXXExpansionStmtPattern' of this particular 'CXXExpansionStmtDecl' is a
+/// 'CXXDestructuringExpansionStmtPattern', which stores, amongst other things,
+/// the declaration of the variable 'x' as well as the expansion-initializer
+/// 'a'.
+///
+/// After expansion, we end up with a 'CXXExpansionStmtInstantiation' that
+/// contains a DecompositionDecl and 3 CompoundStmts, one for each expansion:
+///
+/// \verbatim
+/// {
+///   auto [__u0, __u1, __u2] = a;
+///   {
+///     auto x = __u0;
+///     // ...
+///   }
+///   {
+///     auto x = __u1;
+///     // ...
+///   }
+///   {
+///     auto x = __u2;
+///     // ...
+///   }
+/// }
+/// \endverbatim
+///
+/// The outer braces shown above are implicit; we don't actually create another
+/// CompoundStmt wrapping everything.
+///
+/// \see CXXExpansionStmtPattern
+/// \see CXXExpansionStmtInstantiation
+class CXXExpansionStmtDecl : public Decl, public DeclContext {
+  CXXExpansionStmtPattern *Expansion = nullptr;
+  TemplateParameterList *TParams;
+  CXXExpansionStmtInstantiation *Instantiations = nullptr;
+
+  CXXExpansionStmtDecl(DeclContext *DC, SourceLocation Loc,
+                       TemplateParameterList *TParams);
+
+public:
+  friend class ASTDeclReader;
+
+  static CXXExpansionStmtDecl *Create(ASTContext &C, DeclContext *DC,
+                                      SourceLocation Loc,
+                                      TemplateParameterList *TParams);
+  static CXXExpansionStmtDecl *CreateDeserialized(ASTContext &C,
+                                                  GlobalDeclID ID);
+
+  CXXExpansionStmtPattern *getExpansionPattern() { return Expansion; }
+  const CXXExpansionStmtPattern *getExpansionPattern() const {
+    return Expansion;
+  }
+  void setExpansionPattern(CXXExpansionStmtPattern *S) { Expansion = S; }
+
+  CXXExpansionStmtInstantiation *getInstantiations() { return Instantiations; }
+  const CXXExpansionStmtInstantiation *getInstantiations() const {
+    return Instantiations;
+  }
+
+  void setInstantiations(CXXExpansionStmtInstantiation *S) {
+    Instantiations = S;
+  }
+
+  NonTypeTemplateParmDecl *getIndexTemplateParm() const {
+    return cast<NonTypeTemplateParmDecl>(TParams->getParam(0));
+  }
+  TemplateParameterList *getTemplateParameters() const { return TParams; }
+
+  SourceRange getSourceRange() const override LLVM_READONLY;
+
+  static bool classof(const Decl *D) { return classofKind(D->getKind()); }
+  static bool classofKind(Kind K) { return K == CXXExpansionStmt; }
+};
+
 inline NamedDecl *getAsNamedDecl(TemplateParameter P) {
   if (auto *PD = P.dyn_cast<TemplateTypeParmDecl *>())
     return PD;

clang/include/clang/AST/ExprCXX.h

@@ -5499,6 +5499,165 @@ class BuiltinBitCastExpr final
   }
 };
 
+/// Represents an expansion-init-list of an enumerating expansion statement.
+///
+/// \see CXXEnumeratingExpansionStmtPattern
+class CXXExpansionInitListExpr final
+    : public Expr,
+      llvm::TrailingObjects<CXXExpansionInitListExpr, Expr *> {
+  friend class ASTStmtReader;
+  friend TrailingObjects;
+
+  const unsigned NumExprs;
+  SourceLocation LBraceLoc;
+  SourceLocation RBraceLoc;
+
+  CXXExpansionInitListExpr(EmptyShell ES, unsigned NumExprs);
+  CXXExpansionInitListExpr(ArrayRef<Expr *> Exprs, SourceLocation LBraceLoc,
+                           SourceLocation RBraceLoc);
+
+public:
+  static CXXExpansionInitListExpr *Create(const ASTContext &C,
+                                          ArrayRef<Expr *> Exprs,
+                                          SourceLocation LBraceLoc,
+                                          SourceLocation RBraceLoc);
+
+  static CXXExpansionInitListExpr *
+  CreateEmpty(const ASTContext &C, EmptyShell Empty, unsigned NumExprs);
+
+  ArrayRef<Expr *> getExprs() const { return getTrailingObjects(NumExprs); }
+  MutableArrayRef<Expr *> getExprs() { return getTrailingObjects(NumExprs); }
+  unsigned getNumExprs() const { return NumExprs; }
+
+  bool containsPackExpansion() const;
+
+  SourceLocation getBeginLoc() const { return getLBraceLoc(); }
+  SourceLocation getEndLoc() const { return getRBraceLoc(); }
+
+  SourceLocation getLBraceLoc() const { return LBraceLoc; }
+  SourceLocation getRBraceLoc() const { return RBraceLoc; }
+
+  child_range children() {
+    const_child_range CCR =
+        const_cast<const CXXExpansionInitListExpr *>(this)->children();
+    return child_range(cast_away_const(CCR.begin()),
+                       cast_away_const(CCR.end()));
+  }
+
+  const_child_range children() const {
+    Stmt **Stmts = getTrailingStmts();
+    return const_child_range(Stmts, Stmts + NumExprs);
+  }
+
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == CXXExpansionInitListExprClass;
+  }
+
+private:
+  Stmt **getTrailingStmts() const {
+    return reinterpret_cast<Stmt **>(const_cast<Expr **>(getTrailingObjects()));
+  }
+};
+
+/// Helper that selects an expression from an expansion init list depending
+/// on the current expansion index.
+///
+/// \see CXXEnumeratingExpansionStmtPattern
+class CXXExpansionInitListSelectExpr : public Expr {
+  friend class ASTStmtReader;
+
+  enum SubExpr { RANGE, INDEX, COUNT };
+  Expr *SubExprs[COUNT];
+
+public:
+  CXXExpansionInitListSelectExpr(EmptyShell Empty);
+  CXXExpansionInitListSelectExpr(const ASTContext &C,
+                                 CXXExpansionInitListExpr *Range, Expr *Idx);
+
+  CXXExpansionInitListExpr *getRangeExpr() {
+    return cast<CXXExpansionInitListExpr>(SubExprs[RANGE]);
+  }
+
+  const CXXExpansionInitListExpr *getRangeExpr() const {
+    return cast<CXXExpansionInitListExpr>(SubExprs[RANGE]);
+  }
+
+  void setRangeExpr(CXXExpansionInitListExpr *E) { SubExprs[RANGE] = E; }
+
+  Expr *getIndexExpr() { return SubExprs[INDEX]; }
+  const Expr *getIndexExpr() const { return SubExprs[INDEX]; }
+  void setIndexExpr(Expr *E) { SubExprs[INDEX] = E; }
+
+  SourceLocation getBeginLoc() const { return getRangeExpr()->getBeginLoc(); }
+  SourceLocation getEndLoc() const { return getRangeExpr()->getEndLoc(); }
+
+  child_range children() {
+    return child_range(reinterpret_cast<Stmt **>(SubExprs),
+                       reinterpret_cast<Stmt **>(SubExprs + COUNT));
+  }
+
+  const_child_range children() const {
+    return const_child_range(
+        reinterpret_cast<Stmt **>(const_cast<Expr **>(SubExprs)),
+        reinterpret_cast<Stmt **>(const_cast<Expr **>(SubExprs + COUNT)));
+  }
+
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == CXXExpansionInitListSelectExprClass;
+  }
+};
+
+/// This class serves the same purpose as CXXExpansionInitListSelectExpr, but
+/// for destructuring expansion statements; that is, instead of selecting among
+/// a list of expressions, it selects from a list of 'BindingDecl's.
+///
+/// \see CXXEnumeratingExpansionStmtPattern
+/// \see CXXDestructuringExpansionStmtPattern
+class CXXDestructuringExpansionSelectExpr : public Expr {
+  friend class ASTStmtReader;
+
+  DecompositionDecl *Decomposition;
+  Expr *Index;
+
+public:
+  CXXDestructuringExpansionSelectExpr(EmptyShell Empty);
+  CXXDestructuringExpansionSelectExpr(const ASTContext &C,
+                                      DecompositionDecl *Decomposition,
+                                      Expr *Index);
+
+  DecompositionDecl *getDecompositionDecl() {
+    return cast<DecompositionDecl>(Decomposition);
+  }
+
+  const DecompositionDecl *getDecompositionDecl() const {
+    return cast<DecompositionDecl>(Decomposition);
+  }
+
+  void setDecompositionDecl(DecompositionDecl *E) { Decomposition = E; }
+
+  Expr *getIndexExpr() { return Index; }
+  const Expr *getIndexExpr() const { return Index; }
+  void setIndexExpr(Expr *E) { Index = E; }
+
+  SourceLocation getBeginLoc() const { return Decomposition->getBeginLoc(); }
+  SourceLocation getEndLoc() const { return Decomposition->getEndLoc(); }
+
+  child_range children() {
+    return child_range(reinterpret_cast<Stmt **>(&Index),
+                       reinterpret_cast<Stmt **>(&Index + 1));
+  }
+
+  const_child_range children() const {
+    return const_child_range(
+        reinterpret_cast<Stmt **>(const_cast<Expr **>(&Index)),
+        reinterpret_cast<Stmt **>(const_cast<Expr **>(&Index + 1)));
+  }
+
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == CXXDestructuringExpansionSelectExprClass;
+  }
+};
+
 } // namespace clang
 
 #endif // LLVM_CLANG_AST_EXPRCXX_H

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -1881,6 +1881,14 @@ DEF_TRAVERSE_DECL(UsingShadowDecl, {})
 
 DEF_TRAVERSE_DECL(ConstructorUsingShadowDecl, {})
 
+DEF_TRAVERSE_DECL(CXXExpansionStmtDecl, {
+  if (D->getInstantiations() &&
+      getDerived().shouldVisitTemplateInstantiations())
+    TRY_TO(TraverseStmt(D->getInstantiations()));
+
+  TRY_TO(TraverseStmt(D->getExpansionPattern()));
+})
+
 DEF_TRAVERSE_DECL(OMPThreadPrivateDecl, {
   for (auto *I : D->varlist()) {
     TRY_TO(TraverseStmt(I));
@@ -3117,6 +3125,15 @@ DEF_TRAVERSE_STMT(RequiresExpr, {
     TRY_TO(TraverseConceptRequirement(Req));
 })
 
+DEF_TRAVERSE_STMT(CXXEnumeratingExpansionStmtPattern, {})
+DEF_TRAVERSE_STMT(CXXIteratingExpansionStmtPattern, {})
+DEF_TRAVERSE_STMT(CXXDestructuringExpansionStmtPattern, {})
+DEF_TRAVERSE_STMT(CXXDependentExpansionStmtPattern, {})
+DEF_TRAVERSE_STMT(CXXExpansionStmtInstantiation, {})
+DEF_TRAVERSE_STMT(CXXExpansionInitListExpr, {})
+DEF_TRAVERSE_STMT(CXXExpansionInitListSelectExpr, {})
+DEF_TRAVERSE_STMT(CXXDestructuringExpansionSelectExpr, {})
+
 // These literals (all of them) do not need any action.
 DEF_TRAVERSE_STMT(IntegerLiteral, {})
 DEF_TRAVERSE_STMT(FixedPointLiteral, {})