[AST] Tail-allocate case body variables on CaseStmt

Author: hamishknight

include/swift/AST/Stmt.h

@@ -84,8 +84,9 @@ class alignas(8) Stmt : public ASTAllocated<Stmt> {
     NumElements : 32
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(CaseStmt, Stmt, 32,
+  SWIFT_INLINE_BITFIELD_FULL(CaseStmt, Stmt, 16+32,
     : NumPadBits,
+    NumCaseBodyVars : 16,
     NumPatterns : 32
   );
 
@@ -1210,8 +1211,8 @@ enum CaseParentKind { Switch, DoCatch };
 ///
 class CaseStmt final
     : public Stmt,
-      private llvm::TrailingObjects<CaseStmt, FallthroughStmt *,
-                                    CaseLabelItem> {
+      private llvm::TrailingObjects<CaseStmt, FallthroughStmt *, CaseLabelItem,
+                                    VarDecl *> {
   friend TrailingObjects;
 
   Stmt *ParentStmt = nullptr;
@@ -1222,14 +1223,17 @@ class CaseStmt final
 
   llvm::PointerIntPair<BraceStmt *, 1, bool> BodyAndHasFallthrough;
 
-  ArrayRef<VarDecl *> CaseBodyVariables;
-
   CaseStmt(CaseParentKind ParentKind, SourceLoc ItemIntroducerLoc,
            ArrayRef<CaseLabelItem> CaseLabelItems, SourceLoc UnknownAttrLoc,
            SourceLoc ItemTerminatorLoc, BraceStmt *Body,
            ArrayRef<VarDecl *> CaseBodyVariables, std::optional<bool> Implicit,
            NullablePtr<FallthroughStmt> fallthroughStmt);
 
+  MutableArrayRef<VarDecl *> getCaseBodyVariablesBuffer() {
+    return {getTrailingObjects<VarDecl *>(),
+            static_cast<size_t>(Bits.CaseStmt.NumCaseBodyVars)};
+  }
+
 public:
   /// Create a parsed 'case'/'default' for 'switch' statement.
   static CaseStmt *
@@ -1296,7 +1300,7 @@ class CaseStmt final
   void setBody(BraceStmt *body) { BodyAndHasFallthrough.setPointer(body); }
 
   /// True if the case block declares any patterns with local variable bindings.
-  bool hasCaseBodyVariables() const { return !CaseBodyVariables.empty(); }
+  bool hasCaseBodyVariables() const { return !getCaseBodyVariables().empty(); }
 
   /// Get the source location of the 'case', 'default', or 'catch' of the first
   /// label.
@@ -1349,7 +1353,7 @@ class CaseStmt final
 
   /// Return an ArrayRef containing the case body variables of this CaseStmt.
   ArrayRef<VarDecl *> getCaseBodyVariables() const {
-    return CaseBodyVariables;
+    return const_cast<CaseStmt *>(this)->getCaseBodyVariablesBuffer();
   }
 
   /// Find the next case statement within the same 'switch' or 'do-catch',

lib/AST/Stmt.cpp

@@ -758,9 +758,11 @@ CaseStmt::CaseStmt(CaseParentKind parentKind, SourceLoc itemIntroducerLoc,
     : Stmt(StmtKind::Case, getDefaultImplicitFlag(implicit, itemIntroducerLoc)),
       UnknownAttrLoc(unknownAttrLoc), ItemIntroducerLoc(itemIntroducerLoc),
       ItemTerminatorLoc(itemTerminatorLoc), ParentKind(parentKind),
-      BodyAndHasFallthrough(body, fallthroughStmt.isNonNull()),
-      CaseBodyVariables(caseBodyVariables) {
+      BodyAndHasFallthrough(body, fallthroughStmt.isNonNull()) {
   Bits.CaseStmt.NumPatterns = caseLabelItems.size();
+  Bits.CaseStmt.NumCaseBodyVars = caseBodyVariables.size();
+  ASSERT(Bits.CaseStmt.NumCaseBodyVars == caseBodyVariables.size() &&
+         "too many case body vars");
   assert(Bits.CaseStmt.NumPatterns > 0 &&
          "case block must have at least one pattern");
   assert(
@@ -770,6 +772,9 @@ CaseStmt::CaseStmt(CaseParentKind parentKind, SourceLoc itemIntroducerLoc,
     *getTrailingObjects<FallthroughStmt *>() = fallthroughStmt.get();
   }
 
+  std::uninitialized_copy(caseBodyVariables.begin(), caseBodyVariables.end(),
+                          getCaseBodyVariablesBuffer().begin());
+
   MutableArrayRef<CaseLabelItem> items{getTrailingObjects<CaseLabelItem>(),
                                        static_cast<size_t>(Bits.CaseStmt.NumPatterns)};
 
@@ -914,10 +919,11 @@ CaseStmt *CaseStmt::create(ASTContext &ctx, CaseParentKind ParentKind,
                            BraceStmt *body, ArrayRef<VarDecl *> caseVarDecls,
                            std::optional<bool> implicit,
                            NullablePtr<FallthroughStmt> fallthroughStmt) {
-  void *mem =
-      ctx.Allocate(totalSizeToAlloc<FallthroughStmt *, CaseLabelItem>(
-                       fallthroughStmt.isNonNull(), caseLabelItems.size()),
-                   alignof(CaseStmt));
+  void *mem = ctx.Allocate(
+      totalSizeToAlloc<FallthroughStmt *, CaseLabelItem, VarDecl *>(
+          fallthroughStmt.isNonNull(), caseLabelItems.size(),
+          caseVarDecls.size()),
+      alignof(CaseStmt));
   return ::new (mem)
       CaseStmt(ParentKind, caseLoc, caseLabelItems, unknownAttrLoc, colonLoc,
                body, caseVarDecls, implicit, fallthroughStmt);