[AST] Only store the needed data in WhileStmt.

[riccibruno]

Summary

Don't store the data for the condition variable if not needed.
This cuts the size of WhileStmt by up to a pointer.
The order of the children is kept the same.

[riccibruno]

Add a flag to the output of -ast-dump indicating which sub-statement
a WhileStmt is storing.

[rjmccall]

LGTM.

cfe/trunk/include/clang/AST/Stmt.h

class SwitchStmtBitfields {
     /// value is meant to be a hint for possible clients.
     unsigned AllEnumCasesCovered : 1;
 
     /// The location of the "switch".
     SourceLocation SwitchLoc;
   };
 
   class WhileStmtBitfields {
+    friend class ASTStmtReader;
     friend class WhileStmt;
 
     unsigned : NumStmtBits;
 
+    /// True if the WhileStmt has storage for a condition variable.
+    unsigned HasVar : 1;
+
     /// The location of the "while".
     SourceLocation WhileLoc;
   };
 
   class DoStmtBitfields {
     friend class DoStmt;
 
     unsigned : NumStmtBits;
public:
   }
 
   static bool classof(const Stmt *T) {
     return T->getStmtClass() == SwitchStmtClass;
   }
 };
 
 /// WhileStmt - This represents a 'while' stmt.
-class WhileStmt : public Stmt {
-  enum { VAR, COND, BODY, END_EXPR };
-  Stmt* SubExprs[END_EXPR];
+class WhileStmt final : public Stmt,
+                        private llvm::TrailingObjects<WhileStmt, Stmt *> {
+  friend TrailingObjects;
 
-public:
-  WhileStmt(const ASTContext &C, VarDecl *Var, Expr *cond, Stmt *body,
+  // WhileStmt is followed by several trailing objects,
+  // some of which optional. Note that it would be more
+  // convenient to put the optional trailing object at the end
+  // but this would affect children().
+  // The trailing objects are in order:
+  //
+  // * A "Stmt *" for the condition variable.
+  //    Present if and only if hasVarStorage(). This is in fact a "DeclStmt *".
+  //
+  // * A "Stmt *" for the condition.
+  //    Always present. This is in fact an "Expr *".
+  //
+  // * A "Stmt *" for the body.
+  //    Always present.
+  //
+  enum { VarOffset = 0, BodyOffsetFromCond = 1 };
+  enum { NumMandatoryStmtPtr = 2 };
+
+  unsigned varOffset() const { return VarOffset; }
+  unsigned condOffset() const { return VarOffset + hasVarStorage(); }
+  unsigned bodyOffset() const { return condOffset() + BodyOffsetFromCond; }
+
+  unsigned numTrailingObjects(OverloadToken<Stmt *>) const {
+    return NumMandatoryStmtPtr + hasVarStorage();
+  }
+
+  /// Build a while statement.
+  WhileStmt(const ASTContext &Ctx, VarDecl *Var, Expr *Cond, Stmt *Body,
             SourceLocation WL);
 
   /// Build an empty while statement.
-  explicit WhileStmt(EmptyShell Empty) : Stmt(WhileStmtClass, Empty) {}
+  explicit WhileStmt(EmptyShell Empty, bool HasVar);
+
+public:
+  /// Create a while statement.
+  static WhileStmt *Create(const ASTContext &Ctx, VarDecl *Var, Expr *Cond,
+                           Stmt *Body, SourceLocation WL);
+
+  /// Create an empty while statement optionally with storage for
+  /// a condition variable.
+  static WhileStmt *CreateEmpty(const ASTContext &Ctx, bool HasVar);
+
+  /// True if this WhileStmt has storage for a condition variable.
+  bool hasVarStorage() const { return WhileStmtBits.HasVar; }
+
+  Expr *getCond() {
+    return reinterpret_cast<Expr *>(getTrailingObjects<Stmt *>()[condOffset()]);
+  }
+
+  const Expr *getCond() const {
+    return reinterpret_cast<Expr *>(getTrailingObjects<Stmt *>()[condOffset()]);
+  }
+
+  void setCond(Expr *Cond) {
+    getTrailingObjects<Stmt *>()[condOffset()] = reinterpret_cast<Stmt *>(Cond);
+  }
+
+  Stmt *getBody() { return getTrailingObjects<Stmt *>()[bodyOffset()]; }
+  const Stmt *getBody() const {
+    return getTrailingObjects<Stmt *>()[bodyOffset()];
+  }
+
+  void setBody(Stmt *Body) {
+    getTrailingObjects<Stmt *>()[bodyOffset()] = Body;
+  }
 
   /// Retrieve the variable declared in this "while" statement, if any.
   ///
   /// In the following example, "x" is the condition variable.
   /// \code
   /// while (int x = random()) {
   ///   // ...
   /// }
   /// \endcode
-  VarDecl *getConditionVariable() const;
-  void setConditionVariable(const ASTContext &C, VarDecl *V);
+  VarDecl *getConditionVariable();
+  const VarDecl *getConditionVariable() const {
+    return const_cast<WhileStmt *>(this)->getConditionVariable();
+  }
+
+  /// Set the condition variable of this while statement.
+  /// The while statement must have storage for it.
+  void setConditionVariable(const ASTContext &Ctx, VarDecl *V);
 
   /// If this WhileStmt has a condition variable, return the faux DeclStmt
   /// associated with the creation of that condition variable.
-  const DeclStmt *getConditionVariableDeclStmt() const {
-    return reinterpret_cast<DeclStmt*>(SubExprs[VAR]);
+  DeclStmt *getConditionVariableDeclStmt() {
+    return hasVarStorage() ? static_cast<DeclStmt *>(
+                                 getTrailingObjects<Stmt *>()[varOffset()])
+                           : nullptr;
   }
 
-  Expr *getCond() { return reinterpret_cast<Expr*>(SubExprs[COND]); }
-  const Expr *getCond() const { return reinterpret_cast<Expr*>(SubExprs[COND]);}
-  void setCond(Expr *E) { SubExprs[COND] = reinterpret_cast<Stmt*>(E); }
-  Stmt *getBody() { return SubExprs[BODY]; }
-  const Stmt *getBody() const { return SubExprs[BODY]; }
-  void setBody(Stmt *S) { SubExprs[BODY] = S; }
+  const DeclStmt *getConditionVariableDeclStmt() const {
+    return hasVarStorage() ? static_cast<DeclStmt *>(
+                                 getTrailingObjects<Stmt *>()[varOffset()])
+                           : nullptr;
+  }
 
   SourceLocation getWhileLoc() const { return WhileStmtBits.WhileLoc; }
   void setWhileLoc(SourceLocation L) { WhileStmtBits.WhileLoc = L; }
 
   SourceLocation getBeginLoc() const { return getWhileLoc(); }
-
-  SourceLocation getEndLoc() const { return getBody()->getEndLoc(); }
+  SourceLocation getEndLoc() const LLVM_READONLY {
+    return getBody()->getEndLoc();
+  }
 
   static bool classof(const Stmt *T) {
     return T->getStmtClass() == WhileStmtClass;
   }
 
   // Iterators
   child_range children() {
-    return child_range(&SubExprs[0], &SubExprs[0]+END_EXPR);
+    return child_range(getTrailingObjects<Stmt *>(),
+                       getTrailingObjects<Stmt *>() +
+                           numTrailingObjects(OverloadToken<Stmt *>()));
   }
 };
 
 /// DoStmt - This represents a 'do/while' stmt.
 class DoStmt : public Stmt {
   enum { BODY, COND, END_EXPR };
   Stmt *SubExprs[END_EXPR];
   SourceLocation WhileLoc;

cfe/trunk/lib/AST/ASTDumper.cpp

public:
     void VisitBlockDecl(const BlockDecl *D);
 
     // Stmts.
     void VisitStmt(const Stmt *Node);
     void VisitDeclStmt(const DeclStmt *Node);
     void VisitAttributedStmt(const AttributedStmt *Node);
     void VisitIfStmt(const IfStmt *Node);
     void VisitSwitchStmt(const SwitchStmt *Node);
+    void VisitWhileStmt(const WhileStmt *Node);
     void VisitLabelStmt(const LabelStmt *Node);
     void VisitGotoStmt(const GotoStmt *Node);
     void VisitCXXCatchStmt(const CXXCatchStmt *Node);
     void VisitCaseStmt(const CaseStmt *Node);
     void VisitCapturedStmt(const CapturedStmt *Node);
 
     // OpenMP
     void VisitOMPExecutableDirective(const OMPExecutableDirective *Node);

 void ASTDumper::VisitSwitchStmt(const SwitchStmt *Node) {
   VisitStmt(Node);
   if (Node->hasInitStorage())
     OS << " has_init";
   if (Node->hasVarStorage())
     OS << " has_var";
 }
 
+void ASTDumper::VisitWhileStmt(const WhileStmt *Node) {
+  VisitStmt(Node);
+  if (Node->hasVarStorage())
+    OS << " has_var";
+}
+
 void ASTDumper::VisitLabelStmt(const LabelStmt *Node) {
   VisitStmt(Node);
   OS << " '" << Node->getName() << "'";
 }
 
 void ASTDumper::VisitGotoStmt(const GotoStmt *Node) {
   VisitStmt(Node);
   OS << " '" << Node->getLabel()->getName() << "'";

cfe/trunk/lib/AST/ASTImporter.cpp

if (!Imp)
     return Imp.takeError();
 
   VarDecl *ToConditionVariable;
   Expr *ToCond;
   Stmt *ToBody;
   SourceLocation ToWhileLoc;
   std::tie(ToConditionVariable, ToCond, ToBody, ToWhileLoc) = *Imp;
 
-  return new (Importer.getToContext()) WhileStmt(
-      Importer.getToContext(),
-      ToConditionVariable, ToCond, ToBody, ToWhileLoc);
+  return WhileStmt::Create(Importer.getToContext(), ToConditionVariable, ToCond,
+                           ToBody, ToWhileLoc);
 }
 
 ExpectedStmt ASTNodeImporter::VisitDoStmt(DoStmt *S) {
   auto Imp = importSeq(
       S->getBody(), S->getCond(), S->getDoLoc(), S->getWhileLoc(),
       S->getRParenLoc());
   if (!Imp)
     return Imp.takeError();

cfe/trunk/lib/AST/Stmt.cpp

if (!V) {
     return;
   }
 
   SourceRange VarRange = V->getSourceRange();
   getTrailingObjects<Stmt *>()[varOffset()] = new (Ctx)
       DeclStmt(DeclGroupRef(V), VarRange.getBegin(), VarRange.getEnd());
 }
 
-WhileStmt::WhileStmt(const ASTContext &C, VarDecl *Var, Expr *cond, Stmt *body,
-                     SourceLocation WL)
+WhileStmt::WhileStmt(const ASTContext &Ctx, VarDecl *Var, Expr *Cond,
+                     Stmt *Body, SourceLocation WL)
     : Stmt(WhileStmtClass) {
-  setConditionVariable(C, Var);
-  SubExprs[COND] = cond;
-  SubExprs[BODY] = body;
-  WhileStmtBits.WhileLoc = WL;
+  bool HasVar = Var != nullptr;
+  WhileStmtBits.HasVar = HasVar;
+
+  setCond(Cond);
+  setBody(Body);
+  if (HasVar)
+    setConditionVariable(Ctx, Var);
+
+  setWhileLoc(WL);
 }
 
-VarDecl *WhileStmt::getConditionVariable() const {
-  if (!SubExprs[VAR])
-    return nullptr;
+WhileStmt::WhileStmt(EmptyShell Empty, bool HasVar)
+    : Stmt(WhileStmtClass, Empty) {
+  WhileStmtBits.HasVar = HasVar;
+}
+
+WhileStmt *WhileStmt::Create(const ASTContext &Ctx, VarDecl *Var, Expr *Cond,
+                             Stmt *Body, SourceLocation WL) {
+  bool HasVar = Var != nullptr;
+  void *Mem =
+      Ctx.Allocate(totalSizeToAlloc<Stmt *>(NumMandatoryStmtPtr + HasVar),
+                   alignof(WhileStmt));
+  return new (Mem) WhileStmt(Ctx, Var, Cond, Body, WL);
+}
 
-  auto *DS = cast<DeclStmt>(SubExprs[VAR]);
+WhileStmt *WhileStmt::CreateEmpty(const ASTContext &Ctx, bool HasVar) {
+  void *Mem =
+      Ctx.Allocate(totalSizeToAlloc<Stmt *>(NumMandatoryStmtPtr + HasVar),
+                   alignof(WhileStmt));
+  return new (Mem) WhileStmt(EmptyShell(), HasVar);
+}
+
+VarDecl *WhileStmt::getConditionVariable() {
+  auto *DS = getConditionVariableDeclStmt();
+  if (!DS)
+    return nullptr;
   return cast<VarDecl>(DS->getSingleDecl());
 }
 
-void WhileStmt::setConditionVariable(const ASTContext &C, VarDecl *V) {
+void WhileStmt::setConditionVariable(const ASTContext &Ctx, VarDecl *V) {
+  assert(hasVarStorage() &&
+         "This while statement has no storage for a condition variable!");
+
   if (!V) {
-    SubExprs[VAR] = nullptr;
+    getTrailingObjects<Stmt *>()[varOffset()] = nullptr;
     return;
   }
 
   SourceRange VarRange = V->getSourceRange();
-  SubExprs[VAR] = new (C) DeclStmt(DeclGroupRef(V), VarRange.getBegin(),
-                                   VarRange.getEnd());
+  getTrailingObjects<Stmt *>()[varOffset()] = new (Ctx)
+      DeclStmt(DeclGroupRef(V), VarRange.getBegin(), VarRange.getEnd());
 }
 
 // IndirectGotoStmt
 LabelDecl *IndirectGotoStmt::getConstantTarget() {
   if (auto *E = dyn_cast<AddrLabelExpr>(getTarget()->IgnoreParenImpCasts()))
     return E->getLabel();
   return nullptr;
 }

cfe/trunk/lib/Sema/SemaStmt.cpp

if (CondVal.second &&
       !Diags.isIgnored(diag::warn_comma_operator, CondVal.second->getExprLoc()))
     CommaVisitor(*this).Visit(CondVal.second);
 
   DiagnoseUnusedExprResult(Body);
 
   if (isa<NullStmt>(Body))
     getCurCompoundScope().setHasEmptyLoopBodies();
 
-  return new (Context)
-      WhileStmt(Context, CondVal.first, CondVal.second, Body, WhileLoc);
+  return WhileStmt::Create(Context, CondVal.first, CondVal.second, Body,
+                           WhileLoc);
 }
 
 StmtResult
 Sema::ActOnDoStmt(SourceLocation DoLoc, Stmt *Body,
                   SourceLocation WhileLoc, SourceLocation CondLParen,
                   Expr *Cond, SourceLocation CondRParen) {
   assert(Cond && "ActOnDoStmt(): missing expression");
 

cfe/trunk/lib/Serialization/ASTReaderStmt.cpp

else
       S->setSwitchCaseList(SC);
 
     PrevSC = SC;
   }
 }
 
 void ASTStmtReader::VisitWhileStmt(WhileStmt *S) {
   VisitStmt(S);
-  S->setConditionVariable(Record.getContext(), ReadDeclAs<VarDecl>());
+
+  bool HasVar = Record.readInt();
 
   S->setCond(Record.readSubExpr());
   S->setBody(Record.readSubStmt());
+  if (HasVar)
+    S->setConditionVariable(Record.getContext(), ReadDeclAs<VarDecl>());
+
   S->setWhileLoc(ReadSourceLocation());
 }
 
 void ASTStmtReader::VisitDoStmt(DoStmt *S) {
   VisitStmt(S);
   S->setCond(Record.readSubExpr());
   S->setBody(Record.readSubStmt());
   S->setDoLoc(ReadSourceLocation());
while (true) {
     case STMT_SWITCH:
       S = SwitchStmt::CreateEmpty(
           Context,
           /* HasInit=*/Record[ASTStmtReader::NumStmtFields + 0],
           /* HasVar=*/Record[ASTStmtReader::NumStmtFields + 1]);
       break;
 
     case STMT_WHILE:
-      S = new (Context) WhileStmt(Empty);
+      S = WhileStmt::CreateEmpty(
+          Context,
+          /* HasVar=*/Record[ASTStmtReader::NumStmtFields + 0]);
       break;
 
     case STMT_DO:
       S = new (Context) DoStmt(Empty);
       break;
 
     case STMT_FOR:
       S = new (Context) ForStmt(Empty);

cfe/trunk/lib/Serialization/ASTWriterStmt.cpp

void ASTStmtWriter::VisitSwitchStmt(SwitchStmt *S) {
   for (SwitchCase *SC = S->getSwitchCaseList(); SC;
        SC = SC->getNextSwitchCase())
     Record.push_back(Writer.RecordSwitchCaseID(SC));
   Code = serialization::STMT_SWITCH;
 }
 
 void ASTStmtWriter::VisitWhileStmt(WhileStmt *S) {
   VisitStmt(S);
-  Record.AddDeclRef(S->getConditionVariable());
+
+  bool HasVar = S->getConditionVariableDeclStmt() != nullptr;
+  Record.push_back(HasVar);
+
   Record.AddStmt(S->getCond());
   Record.AddStmt(S->getBody());
+  if (HasVar)
+    Record.AddDeclRef(S->getConditionVariable());
+
   Record.AddSourceLocation(S->getWhileLoc());
   Code = serialization::STMT_WHILE;
 }
 
 void ASTStmtWriter::VisitDoStmt(DoStmt *S) {
   VisitStmt(S);
   Record.AddStmt(S->getCond());
   Record.AddStmt(S->getBody());

cfe/trunk/test/Import/while-stmt/test.cpp

 // RUN: clang-import-test -dump-ast -import %S/Inputs/F.cpp -expression %s | FileCheck %s
 
 // CHECK: WhileStmt
-// CHECK-NEXT: <<NULL>>
 // CHECK-NEXT: CXXBoolLiteralExpr
 // CHECK-NEXT: NullStmt
 
 // CHECK: WhileStmt
-// CHECK-NEXT: <<NULL>>
 // CHECK-NEXT: CXXBoolLiteralExpr
 // CHECK-NEXT: CompoundStmt
 
 // CHECK: WhileStmt
 // CHECK-NEXT: DeclStmt
 // CHECK-NEXT: VarDecl
 // CHECK-NEXT: CXXBoolLiteralExpr
 // CHECK-NEXT: ImplicitCastExpr
 // CHECK-NEXT: DeclRefExpr
 // CHECK-NEXT: NullStmt
 
 void expr() {
   f();
 }