Merge pull request #70172 from hamishknight/or-else

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/SIL/IR/SILProfiler.cpp

@@ -422,6 +422,13 @@ class CounterExpr {
   /// Returns true if this is a zero counter.
   bool isZero() const { return Counter.isZero(); }
 
+  friend bool operator==(const CounterExpr &LHS, const CounterExpr &RHS) {
+    return LHS.Counter == RHS.Counter;
+  }
+  friend bool operator!=(const CounterExpr &LHS, const CounterExpr &RHS) {
+    return !(LHS == RHS);
+  }
+
   llvm::coverage::Counter getLLVMCounter() const { return Counter; }
 
   void print(raw_ostream &OS,
@@ -476,8 +483,9 @@ class SourceMappingRegion {
   }
 
   SourceMappingRegion(Kind RegionKind, ASTNode Node, SourceRange Range,
+                      llvm::Optional<CounterExpr> Counter,
                       const SourceManager &SM)
-      : RegionKind(RegionKind), Node(Node) {
+      : RegionKind(RegionKind), Node(Node), Counter(Counter) {
     assert(Range.isValid());
     StartLoc = Range.Start;
     EndLoc = Lexer::getLocForEndOfToken(SM, Range.End);
@@ -492,16 +500,18 @@ class SourceMappingRegion {
 
     // Note we don't store counters for nodes, as we need to be able to fix them
     // up later.
-    return SourceMappingRegion(Kind::Node, Node, Range, SM);
+    return SourceMappingRegion(Kind::Node, Node, Range, /*Counter*/ llvm::None,
+                               SM);
   }
 
   /// Create a source region for an ASTNode that is only present for scoping of
   /// child regions, and doesn't need to be included in the resulting set of
   /// regions.
-  static SourceMappingRegion scopingOnly(ASTNode Node,
-                                         const SourceManager &SM) {
+  static SourceMappingRegion
+  scopingOnly(ASTNode Node, const SourceManager &SM,
+              llvm::Optional<CounterExpr> Counter = llvm::None) {
     return SourceMappingRegion(Kind::ScopingOnly, Node, Node.getSourceRange(),
-                               SM);
+                               Counter, SM);
   }
 
   /// Create a refined region for a given counter.
@@ -967,6 +977,14 @@ struct CoverageMapping : public ASTWalker {
           return;
         }
 
+        // Don't apply exit adjustments to if statement branches, they should
+        // be handled at the end of the statement. This avoids creating awkward
+        // overlapping exit regions for each branch, and ensures 'break'
+        // statements only have their jump counted once for the entire
+        // statement.
+        if (isa<IfStmt>(ParentStmt))
+          return;
+
         if (auto *LS = dyn_cast<LabeledStmt>(ParentStmt))
           JumpsToLabel = getCounter(LS);
       }
@@ -1225,7 +1243,7 @@ struct CoverageMapping : public ASTWalker {
       // it by break statements.
       assignCounter(IS, CounterExpr::Zero());
 
-      // FIXME: This is a redundant region.
+      // FIXME: This is a redundant region for non else-ifs.
       if (auto *Cond = getConditionNode(IS->getCond()))
         assignCounter(Cond, getCurrentCounter());
 
@@ -1244,13 +1262,68 @@ struct CoverageMapping : public ASTWalker {
         // terms of it.
         auto ThenCounter = assignKnownCounter(IS->getThenStmt());
         IS->getThenStmt()->walk(*this);
+        auto ThenDelta =
+            CounterExpr::Sub(ThenCounter, getExitCounter(), CounterBuilder);
 
+        llvm::Optional<CounterExpr> ElseDelta;
         if (auto *Else = IS->getElseStmt()) {
           auto ElseCounter = CounterExpr::Sub(ParentCounter, ThenCounter,
                                               CounterBuilder);
-          assignCounter(Else, ElseCounter);
+          // We handle `else if` and `else` slightly differently here. For
+          // `else` we have a BraceStmt, and can use the existing scoping logic
+          // to handle calculating the exit count. For `else if`, we need to
+          // set up a new scope to contain the child `if` statement, effectively
+          // we treat:
+          //
+          // if .random() {
+          // } else if .random() {
+          // } else {
+          // }
+          //
+          // the same as:
+          //
+          // if .random() {
+          // } else {
+          //   if .random() {
+          //   } else {
+          //   }
+          // }
+          //
+          // This ensures we assign a correct counter to the `else if`
+          // condition, and allows us to compute the exit count correctly. We
+          // don't need the fake `else` scope to be included in the resulting
+          // set of regions, so we mark it scoping-only.
+          if (isa<BraceStmt>(Else)) {
+            assignCounter(Else, ElseCounter);
+          } else {
+            pushRegion(SourceMappingRegion::scopingOnly(Else, SM, ElseCounter));
+          }
           Else->walk(*this);
+
+          // Once we've walked the `else`, compute the delta exit count. For
+          // a normal `else` we can use the computed exit count, for an
+          // `else if` we can take the current region count since we don't have
+          // a proper scope. This is a little hacked together, but we'll be able
+          // to do away with all of this once we re-implement as a SILOptimizer
+          // pass.
+          auto AfterElse = isa<BraceStmt>(Else) ? getExitCounter()
+                                                : getCurrentCounter();
+          if (!isa<BraceStmt>(Else))
+            popRegions(Else);
+
+          ElseDelta = CounterExpr::Sub(ElseCounter, AfterElse, CounterBuilder);
         }
+        // Compute the exit count following the `if`, taking jumps to the
+        // statement by breaks into account, and the delta of the `then` branch
+        // and `else` branch if we have one.
+        auto AfterIf = getCurrentCounter();
+        AfterIf = CounterExpr::Add(AfterIf, getCounter(IS), CounterBuilder);
+        AfterIf = CounterExpr::Sub(AfterIf, ThenDelta, CounterBuilder);
+        if (ElseDelta)
+          AfterIf = CounterExpr::Sub(AfterIf, *ElseDelta, CounterBuilder);
+
+        if (AfterIf != getCurrentCounter())
+          replaceCount(AfterIf, getEndLoc(IS));
       }
       // Already visited the children.
       // FIXME: The ASTWalker should do a post-walk for SkipChildren.

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));
       }
     }