[AST] Tighten up invariants around IfStmt

The 'then' statement must be a BraceStmt, and
the 'else' must either be a BraceStmt or an IfStmt.

Author: hamishknight

include/swift/AST/ASTBridging.h

@@ -734,7 +734,8 @@ SWIFT_NAME("BridgedIfStmt.createParsed(_:ifKeywordLoc:condition:thenStmt:"
            "elseLoc:elseStmt:)")
 BridgedIfStmt BridgedIfStmt_createParsed(BridgedASTContext cContext,
                                          BridgedSourceLoc cIfLoc,
-                                         BridgedExpr cond, BridgedStmt then,
+                                         BridgedExpr cond,
+                                         BridgedBraceStmt then,
                                          BridgedSourceLoc cElseLoc,
                                          BridgedNullableStmt elseStmt);
 

include/swift/AST/Stmt.h

@@ -757,20 +757,25 @@ class LabeledConditionalStmt : public LabeledStmt {
 class IfStmt : public LabeledConditionalStmt {
   SourceLoc IfLoc;
   SourceLoc ElseLoc;
-  Stmt *Then;
+  BraceStmt *Then;
   Stmt *Else;
 
 public:
   IfStmt(LabeledStmtInfo LabelInfo, SourceLoc IfLoc, StmtCondition Cond,
-         Stmt *Then, SourceLoc ElseLoc, Stmt *Else,
+         BraceStmt *Then, SourceLoc ElseLoc, Stmt *Else,
          llvm::Optional<bool> implicit = llvm::None)
       : LabeledConditionalStmt(StmtKind::If,
                                getDefaultImplicitFlag(implicit, IfLoc),
                                LabelInfo, Cond),
-        IfLoc(IfLoc), ElseLoc(ElseLoc), Then(Then), Else(Else) {}
+        IfLoc(IfLoc), ElseLoc(ElseLoc), Then(Then), Else(Else) {
+    assert(Then && "Must have non-null 'then' statement");
+    assert(!Else || isa<BraceStmt>(Else) ||
+           isa<IfStmt>(Else) &&
+               "Else statement must either be BraceStmt or IfStmt");
+  }
 
-  IfStmt(SourceLoc IfLoc, Expr *Cond, Stmt *Then, SourceLoc ElseLoc, Stmt *Else,
-         llvm::Optional<bool> implicit, ASTContext &Ctx);
+  IfStmt(SourceLoc IfLoc, Expr *Cond, BraceStmt *Then, SourceLoc ElseLoc,
+         Stmt *Else, llvm::Optional<bool> implicit, ASTContext &Ctx);
 
   SourceLoc getIfLoc() const { return IfLoc; }
   SourceLoc getElseLoc() const { return ElseLoc; }
@@ -782,8 +787,8 @@ class IfStmt : public LabeledConditionalStmt {
     return (Else ? Else->getEndLoc() : Then->getEndLoc());
   }
 
-  Stmt *getThenStmt() const { return Then; }
-  void setThenStmt(Stmt *s) { Then = s; }
+  BraceStmt *getThenStmt() const { return Then; }
+  void setThenStmt(BraceStmt *s) { Then = s; }
 
   Stmt *getElseStmt() const { return Else; }
   void setElseStmt(Stmt *s) { Else = s; }

lib/AST/ASTBridging.cpp

@@ -1051,7 +1051,8 @@ BridgedBraceStmt BridgedBraceStmt_createParsed(BridgedASTContext cContext,
 
 BridgedIfStmt BridgedIfStmt_createParsed(BridgedASTContext cContext,
                                          BridgedSourceLoc cIfLoc,
-                                         BridgedExpr cond, BridgedStmt then,
+                                         BridgedExpr cond,
+                                         BridgedBraceStmt then,
                                          BridgedSourceLoc cElseLoc,
                                          BridgedNullableStmt elseStmt) {
   ASTContext &context = cContext.unbridged();

lib/AST/ASTWalker.cpp

@@ -1813,8 +1813,8 @@ Stmt *Traversal::visitIfStmt(IfStmt *IS) {
   if (doIt(IS->getCond()))
     return nullptr;
 
-  if (Stmt *S2 = doIt(IS->getThenStmt()))
-    IS->setThenStmt(S2);
+  if (auto *S2 = doIt(IS->getThenStmt()))
+    IS->setThenStmt(cast<BraceStmt>(S2));
   else
     return nullptr;
 

lib/AST/Stmt.cpp

@@ -661,7 +661,7 @@ static StmtCondition exprToCond(Expr *C, ASTContext &Ctx) {
   return Ctx.AllocateCopy(Arr);
 }
 
-IfStmt::IfStmt(SourceLoc IfLoc, Expr *Cond, Stmt *Then, SourceLoc ElseLoc,
+IfStmt::IfStmt(SourceLoc IfLoc, Expr *Cond, BraceStmt *Then, SourceLoc ElseLoc,
                Stmt *Else, llvm::Optional<bool> implicit, ASTContext &Ctx)
     : IfStmt(LabeledStmtInfo(), IfLoc, exprToCond(Cond, Ctx), Then, ElseLoc,
              Else, implicit) {}

lib/ASTGen/Sources/ASTGen/Stmts.swift

@@ -73,7 +73,7 @@ extension ASTGenVisitor {
       self.ctx,
       ifKeywordLoc: node.ifKeyword.bridgedSourceLoc(in: self),
       condition: conditions.first!.castToExpr,
-      thenStmt: self.generate(codeBlock: node.body).asStmt,
+      thenStmt: self.generate(codeBlock: node.body),
       elseLoc: node.elseKeyword.bridgedSourceLoc(in: self),
       elseStmt: node.elseBody.map {
         switch $0 {

lib/Sema/BuilderTransform.cpp

@@ -494,7 +494,7 @@ class ResultBuilderTransform
         availabilityCond && builder.supports(ctx.Id_buildLimitedAvailability);
 
     NullablePtr<Expr> thenVarRef;
-    NullablePtr<Stmt> thenBranch;
+    NullablePtr<BraceStmt> thenBranch;
     {
       SmallVector<ASTNode, 4> thenBody;
 

lib/Sema/CSSyntacticElement.cpp

@@ -1803,7 +1803,7 @@ class SyntacticElementSolutionApplication
     else
       hadError = true;
 
-    ifStmt->setThenStmt(visit(ifStmt->getThenStmt()).get<Stmt *>());
+    ifStmt->setThenStmt(castToStmt<BraceStmt>(visit(ifStmt->getThenStmt())));
 
     if (auto elseStmt = ifStmt->getElseStmt()) {
       ifStmt->setElseStmt(visit(elseStmt).get<Stmt *>());

lib/Sema/PCMacro.cpp

@@ -142,10 +142,10 @@ class Instrumenter : InstrumenterBase {
     transformStmtCondition(SC, IS->getStartLoc());
     IS->setCond(SC); // FIXME: is setting required?..
 
-    if (Stmt *TS = IS->getThenStmt()) {
-      Stmt *NTS = transformStmt(TS);
+    if (auto *TS = IS->getThenStmt()) {
+      auto *NTS = transformStmt(TS);
       if (NTS != TS) {
-        IS->setThenStmt(NTS);
+        IS->setThenStmt(cast<BraceStmt>(NTS));
       }
     }
 

lib/Sema/PlaygroundTransform.cpp

@@ -205,10 +205,10 @@ class Instrumenter : InstrumenterBase {
   // transform*() return their input if it's unmodified,
   // or a modified copy of their input otherwise.
   IfStmt *transformIfStmt(IfStmt *IS) {
-    if (Stmt *TS = IS->getThenStmt()) {
-      Stmt *NTS = transformStmt(TS);
+    if (auto *TS = IS->getThenStmt()) {
+      auto *NTS = transformStmt(TS);
       if (NTS != TS) {
-        IS->setThenStmt(NTS);
+        IS->setThenStmt(cast<BraceStmt>(NTS));
       }
     }