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

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/Decl.h

@@ -1247,14 +1247,16 @@ class VarDecl : public DeclaratorDecl, public Redeclarable<VarDecl> {
 
   /// Returns true for local variable declarations other than parameters.
   /// Note that this includes static variables inside of functions. It also
-  /// includes variables inside blocks.
+  /// includes variables inside blocks and expansion statements.
   ///
   ///   void foo() { int x; static int y; extern int z; }
   bool isLocalVarDecl() const {
     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,127 @@ 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.
+///
+/// There are different kinds of expansion statements; see the comment on
+/// 'CXXExpansionStmtPattern' for more information.
+///
+/// As an 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'
+/// 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
+/// is *equivalent* to the AST shown below. Note that only the inner '{}' (i.e.
+/// those marked as 'Actual "CompoundStmt"' below) are actually present as
+/// 'CompoundStmt's in the AST; the outer braces that wrap everything do *not*
+/// correspond to an actual 'CompoundStmt' and are implicit in the sense that we
+/// simply push a scope when evaluating or emitting IR for a
+/// 'CXXExpansionStmtInstantiation'.
+///
+/// \verbatim
+/// { // Not actually present in the AST.
+///   auto [__u0, __u1, __u2] = a;
+///   { // Actual 'CompoundStmt'.
+///     auto x = __u0;
+///     // ...
+///   }
+///   { // Actual 'CompoundStmt'.
+///     auto x = __u1;
+///     // ...
+///   }
+///   { // Actual 'CompoundStmt'.
+///     auto x = __u2;
+///     // ...
+///   }
+/// }
+/// \endverbatim
+///
+/// See the documentation around 'CXXExpansionStmtInstantiation' for more notes
+/// as to why this node exist and how it is used.
+///
+/// \see CXXExpansionStmtPattern
+/// \see CXXExpansionStmtInstantiation
+class CXXExpansionStmtDecl : public Decl, public DeclContext {
+  CXXExpansionStmtPattern *Expansion = nullptr;
+  NonTypeTemplateParmDecl *IndexNTTP = nullptr;
+  CXXExpansionStmtInstantiation *Instantiations = nullptr;
+
+  CXXExpansionStmtDecl(DeclContext *DC, SourceLocation Loc,
+                       NonTypeTemplateParmDecl *NTTP);
+
+public:
+  friend class ASTDeclReader;
+
+  static CXXExpansionStmtDecl *Create(ASTContext &C, DeclContext *DC,
+                                      SourceLocation Loc,
+                                      NonTypeTemplateParmDecl *NTTP);
+  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() { return IndexNTTP; }
+  const NonTypeTemplateParmDecl *getIndexTemplateParm() const {
+    return IndexNTTP;
+  }
+
+  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,52 @@ class BuiltinBitCastExpr final
   }
 };
 
+/// Helper that selects an expression from an InitListExpr depending
+/// on the current expansion index.
+///
+/// \see CXXExpansionStmtPattern
+class CXXExpansionSelectExpr : public Expr {
+  friend class ASTStmtReader;
+
+  enum SubExpr { RANGE, INDEX, COUNT };
+  Expr *SubExprs[COUNT];
+
+public:
+  CXXExpansionSelectExpr(EmptyShell Empty);
+  CXXExpansionSelectExpr(const ASTContext &C, InitListExpr *Range, Expr *Idx);
+
+  InitListExpr *getRangeExpr() {
+    return cast<InitListExpr>(SubExprs[RANGE]);
+  }
+
+  const InitListExpr *getRangeExpr() const {
+    return cast<InitListExpr>(SubExprs[RANGE]);
+  }
+
+  void setRangeExpr(InitListExpr *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() == CXXExpansionSelectExprClass;
+  }
+};
 } // 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,10 @@ DEF_TRAVERSE_STMT(RequiresExpr, {
     TRY_TO(TraverseConceptRequirement(Req));
 })
 
+DEF_TRAVERSE_STMT(CXXExpansionStmtPattern, {})
+DEF_TRAVERSE_STMT(CXXExpansionStmtInstantiation, {})
+DEF_TRAVERSE_STMT(CXXExpansionSelectExpr, {})
+
 // These literals (all of them) do not need any action.
 DEF_TRAVERSE_STMT(IntegerLiteral, {})
 DEF_TRAVERSE_STMT(FixedPointLiteral, {})