[SIL Diagnostics] Improve diagnostics for yield-once coroutines

when the coroutines yield in some paths but not in all paths.

<rdar://48184430>

Author: ravikandhadai

include/swift/AST/DiagnosticsSIL.def

@@ -369,7 +369,20 @@ NOTE(previous_yield, none, "previous yield was here", ())
 ERROR(possible_return_before_yield, none,
       "accessor must yield on all paths before returning", ())
 
-NOTE(one_yield, none, "yield along one path is here", ())
+NOTE(branch_doesnt_yield, none,
+     "missing yield when the condition is %select{false|true}0", (bool))
+
+NOTE(named_case_doesnt_yield, none, "missing yield in the %0 case",
+    (Identifier))
+
+NOTE(case_doesnt_yield, none, "missing yield in "
+     "%select{this|the nil|the non-nil}0 case", (unsigned))
+
+NOTE(switch_value_case_doesnt_yield, none, "missing yield in the %0 case",
+    (StringRef))
+
+NOTE(try_branch_doesnt_yield, none, "missing yield when error is "
+     "%select{not |}0thrown", (bool))
 
 #ifndef DIAG_NO_UNDEF
 # if defined(DIAG)

include/swift/SIL/SILInstruction.h

@@ -7189,6 +7189,15 @@ class SwitchValueInst final
     assert(hasDefault() && "doesn't have a default");
     return getSuccessorBuf()[getNumCases()];
   }
+
+  Optional<unsigned> getUniqueCaseForDestination(SILBasicBlock *bb) const {
+    for (unsigned i = 0; i < getNumCases(); ++i) {
+      if (getCase(i).second == bb) {
+        return i + 1;
+      }
+    }
+    return None;
+  }
 };
 
 /// Common implementation for the switch_enum and

include/swift/SIL/SILSuccessor.h

@@ -113,6 +113,12 @@ class SILSuccessor {
       return *this;
     }
 
+    pred_iterator operator++(int) {
+      auto old = *this;
+      ++*this;
+      return old;
+    }
+
     pred_iterator operator+(unsigned distance) const {
       auto copy = *this;
       if (distance == 0)

lib/SILOptimizer/Mandatory/YieldOnceCheck.cpp

@@ -22,8 +22,10 @@
 #include "swift/AST/Expr.h"
 #include "swift/AST/Stmt.h"
 #include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SIL/CFG.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "llvm/ADT/BreadthFirstIterator.h"
 #include "llvm/ADT/DenseSet.h"
 
 using namespace swift;
@@ -48,32 +50,54 @@ class YieldOnceCheck : public SILFunctionTransform {
     enum YieldState { BeforeYield, AfterYield, Conflict } yieldState;
 
   private:
-    // The following state is maintained for emitting diagnostics.
+    // The following states are maintained for emitting diagnostics.
 
     /// For AfterYield and Conflict states, this field records the yield
     /// instruction that was seen while propagating the state.
     SILInstruction *yieldInst;
 
-    BBState(YieldState yState, SILInstruction *yieldI)
-        : yieldState(yState), yieldInst(yieldI) {}
+    /// For Conflict state, these fields record the basic blocks that
+    /// propagated the 'AfterYield' and 'BeforeYield' states which resulted
+    /// in the Conflict.
+    SILBasicBlock *yieldingPred;
+    SILBasicBlock *nonYieldingPred;
+
+    BBState(YieldState yState, SILInstruction *yieldI, SILBasicBlock *yieldPred,
+            SILBasicBlock *noYieldPred)
+        : yieldState(yState), yieldInst(yieldI), yieldingPred(yieldPred),
+          nonYieldingPred(noYieldPred) {}
 
   public:
-    static BBState getInitialState() { return BBState(BeforeYield, nullptr); }
+    static BBState getInitialState() {
+      return BBState(BeforeYield, nullptr, nullptr, nullptr);
+    }
 
     static BBState getAfterYieldState(SILInstruction *yieldI) {
       assert(yieldI);
-      return BBState(AfterYield, yieldI);
+      return BBState(AfterYield, yieldI, nullptr, nullptr);
     }
 
-    static BBState getConflictState(SILInstruction *yieldI) {
-      assert(yieldI);
-      return BBState(Conflict, yieldI);
+    static BBState getConflictState(SILInstruction *yieldI,
+                                    SILBasicBlock *yieldPred,
+                                    SILBasicBlock *noYieldPred) {
+      assert(yieldI && yieldPred && noYieldPred);
+      return BBState(Conflict, yieldI, yieldPred, noYieldPred);
     }
 
     SILInstruction *getYieldInstruction() const {
       assert(yieldState == AfterYield || yieldState == Conflict);
       return yieldInst;
     }
+
+    SILBasicBlock *getYieldingPred() {
+      assert(yieldState == Conflict);
+      return yieldingPred;
+    }
+
+    SILBasicBlock *getNonYieldingPred() {
+      assert(yieldState == Conflict);
+      return nonYieldingPred;
+    }
   };
 
   /// A structure that records an error found during the analysis along with
@@ -157,9 +181,16 @@ class YieldOnceCheck : public SILFunctionTransform {
   /// \param mergeBlock the basic block that is reached with two states
   /// \param oldState the previous state at the entry of the basic block
   /// \param newState the current state at the entry of the basic block
+  /// \param newStatePred the predecessor of 'mergeBlock' that has propagated
+  ///        the 'newState'.
+  /// \param bbToStateMap a map from the basic blocks visited by the analysis
+  ///        to the BBStates in which they were seen. This is used to identify
+  ///        blocks that propagate conflicting states when the merge results
+  ///        in a conflict.
   /// \return the new state obtained by merging the oldState with the newState
   static BBState merge(SILBasicBlock *mergeBlock, BBState oldState,
-                       BBState newState) {
+                       BBState newState, SILBasicBlock *newStatePred,
+                       llvm::DenseMap<SILBasicBlock *, BBState> &bbToStateMap) {
     auto oldYieldState = oldState.yieldState;
     auto newYieldState = newState.yieldState;
 
@@ -182,11 +213,27 @@ class YieldOnceCheck : public SILFunctionTransform {
            (newYieldState == BBState::BeforeYield &&
             oldYieldState == BBState::AfterYield));
 
-    SILInstruction *yieldInst = (newYieldState == BBState::AfterYield)
-                                    ? newState.getYieldInstruction()
-                                    : oldState.getYieldInstruction();
-
-    return BBState::getConflictState(yieldInst);
+    // For diagnostics, find another predecessor of 'mergeBlock' that was
+    // previously seen by the analysis. This predecessor would have
+    // propagated the 'oldState'.
+    SILBasicBlock *oldStatePred = nullptr;
+    for (auto predBB : mergeBlock->getPredecessorBlocks()) {
+      if (predBB != newStatePred && bbToStateMap.count(predBB)) {
+        oldStatePred = predBB;
+        break;
+      }
+    }
+    assert(oldStatePred);
+
+    if (oldState.yieldState == BBState::BeforeYield) {
+      return BBState::getConflictState(newState.getYieldInstruction(),
+                                       /* yieldPred */ newStatePred,
+                                       /* noYieldPred */ oldStatePred);
+    } else {
+      return BBState::getConflictState(oldState.getYieldInstruction(),
+                                       /* yieldPred */ oldStatePred,
+                                       /* noYieldPred */ newStatePred);
+    }
   }
 
   /// Perform a data-flow analysis to check whether there is exactly one
@@ -195,6 +242,7 @@ class YieldOnceCheck : public SILFunctionTransform {
   /// the control-flow graph.
   void diagnoseYieldOnceUsage(SILFunction &fun) {
     llvm::DenseMap<SILBasicBlock *, BBState> bbToStateMap;
+
     SmallVector<SILBasicBlock *, 16> worklist;
 
     auto *entryBB = fun.getEntryBlock();
@@ -234,7 +282,7 @@ class YieldOnceCheck : public SILFunctionTransform {
           continue;
         }
 
-        emitDiagnostics(error, fun);
+        emitDiagnostics(error, fun, bbToStateMap);
         return;
       }
 
@@ -255,7 +303,7 @@ class YieldOnceCheck : public SILFunctionTransform {
         // previous states and propagate it if it is different from the
         // old state.
         const auto &oldState = insertResult.first->second;
-        auto mergedState = merge(succBB, oldState, nextState);
+        auto mergedState = merge(succBB, oldState, nextState, bb, bbToStateMap);
 
         if (mergedState.yieldState == oldState.yieldState)
           continue;
@@ -271,11 +319,12 @@ class YieldOnceCheck : public SILFunctionTransform {
     }
 
     if (returnBeforeYieldError.hasValue()) {
-      emitDiagnostics(returnBeforeYieldError.getValue(), fun);
+      emitDiagnostics(returnBeforeYieldError.getValue(), fun, bbToStateMap);
     }
   }
 
-  void emitDiagnostics(YieldError &error, SILFunction &fun) {
+  void emitDiagnostics(YieldError &error, SILFunction &fun,
+                       llvm::DenseMap<SILBasicBlock *, BBState> &visitedBBs) {
     ASTContext &astCtx = fun.getModule().getASTContext();
 
     switch (error.errorKind) {
@@ -299,9 +348,162 @@ class YieldOnceCheck : public SILFunctionTransform {
       diagnose(astCtx, error.termInst->getLoc().getSourceLoc(),
                diag::possible_return_before_yield);
 
-      // Add a note that points to the yield instruction found.
-      auto *yieldInst = error.inState.getYieldInstruction();
-      diagnose(astCtx, yieldInst->getLoc().getSourceLoc(), diag::one_yield);
+      // Here, the yield state of 'error' is Conflict.
+      auto &conflictState = error.inState;
+
+      // Emit a note that pin points the branch construct that resulted in
+      // this conflict. Note that a conflict state is created at a merge block
+      // when along one incoming edge the analysis sees a BeforeYield state
+      // and along another it sees an AfterYield state.
+      // Also note that, by the definition of the merge operation,
+      // 'error.yieldingPred()' is the immediate predecessor of the merge block
+      // that propagates AfterYield state, and 'error.nonYieldingPred()' is
+      // the immediate predecessor of the merge block that propagates a
+      // BeforeYield state.
+      auto yieldPred = conflictState.getYieldingPred();
+      auto noYieldPred = conflictState.getNonYieldingPred();
+
+      // Step 1: find a branching SIL instruction where one branch has
+      // 'yieldPred' and another branch has 'noYieldPred'. For instance,
+      // in the following example, cond_br is the instruction to find.
+      //              cond_br bb1, bb2
+      //              bb1:
+      //                  yield resume yieldPred, unwind err
+      //              bb2:
+      //                  br noYieldPred
+      //              yieldPred:
+      //                  br mergePoint
+      //              noYieldPred:
+      //                  br mergePoint
+      //              mergePoint:
+      //                  ...
+      // Intuitively, the conflicting branch is the instruction where
+      // 'yieldPred' and 'noYieldPred' meet when traversing the CFG in the
+      // reverse order. More formally, the branching instruction is a
+      // "lowest common ancestor" of 'yieldPred' and 'noYieldPred' in the
+      // the DAG obtained from the CFG by ignoring back edges of loops.
+      //
+      // Note that the lowest common ancestor may not be unique in a DAG.
+      // But, any such ancestor could be considered as the conflicting branch as
+      // 'yieldPred' and 'noYieldPred' will belong to two different branches of
+      // every such ancestor.
+
+      // Find all transitive predecessors of 'yieldPred' that were visited
+      // during the analysis
+      SmallPtrSet<SILBasicBlock *, 8> predecessorsOfYieldPred;
+      for (auto *predBB :
+           llvm::breadth_first<llvm::Inverse<SILBasicBlock *>>(yieldPred)) {
+        if (visitedBBs.count(predBB)) {
+          predecessorsOfYieldPred.insert(predBB);
+        }
+      }
+
+      // Find the first predecessor of 'noYieldPred' that is also a predecessor
+      // of 'yieldPred', in the breadth-first search order of the reversed CFG.
+      SILBasicBlock *lowestCommonAncestorBB = nullptr;
+      SmallPtrSet<SILBasicBlock *, 8> predecessorsOfNoYieldPred;
+      for (auto *pred :
+           llvm::breadth_first<llvm::Inverse<SILBasicBlock *>>(noYieldPred)) {
+        if (!visitedBBs.count(pred)) {
+          continue;
+        }
+        if (predecessorsOfYieldPred.count(pred)) {
+          lowestCommonAncestorBB = pred;
+          break;
+        }
+        predecessorsOfNoYieldPred.insert(pred);
+      }
+      assert(lowestCommonAncestorBB);
+
+      auto *conflictingBranch = lowestCommonAncestorBB->getTerminator();
+
+      // Step 2: Find the target basic block of the 'conflictingBranch' that
+      // doesn't yield.
+      SILBasicBlock *noYieldTarget = nullptr;
+      for (auto *succ : lowestCommonAncestorBB->getSuccessorBlocks()) {
+        if (predecessorsOfNoYieldPred.count(succ)) {
+          noYieldTarget = succ;
+          break;
+        }
+      }
+      assert(noYieldTarget);
+
+      // Step 3: Report specialized diagnostics for each kind of conflicting
+      // branch.
+
+      // For conditional-branch instructions, which correspond to the
+      // conditions of 'if', 'where' or 'guard' statements, report for what
+      // truth value the branch doesn't yield.
+      if (auto *condbr = dyn_cast<CondBranchInst>(conflictingBranch)) {
+        diagnose(astCtx, condbr->getLoc().getSourceLoc(),
+                 diag::branch_doesnt_yield,
+                 /*non-yielding branch*/ condbr->getTrueBB() == noYieldTarget);
+        return;
+      }
+
+      // For switch_enum instructions, report the case that doesn't yield.
+      enum SwitchCaseKind { Default, OptionNil, OptionSome };
+
+      if (auto *switchEnum = dyn_cast<SwitchEnumInstBase>(conflictingBranch)) {
+        auto enumCaseLoc = noYieldTarget->begin()->getLoc().getSourceLoc();
+
+        if (switchEnum->hasDefault() &&
+            switchEnum->getDefaultBB() == noYieldTarget) {
+          diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, Default);
+          return;
+        }
+
+        // Find the case identifier that doesn't yield.
+        NullablePtr<EnumElementDecl> enumElemDecl =
+            switchEnum->getUniqueCaseForDestination(noYieldTarget);
+        assert(enumElemDecl.isNonNull());
+
+        // Specialize diagnostics for cases of an optional.
+        if (enumElemDecl.get() == astCtx.getOptionalNoneDecl()) {
+          diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, OptionNil);
+        } else if (enumElemDecl.get() == astCtx.getOptionalSomeDecl()) {
+          diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, OptionSome);
+        } else {
+          diagnose(astCtx, enumCaseLoc, diag::named_case_doesnt_yield,
+                   enumElemDecl.get()->getName());
+        }
+        return;
+      }
+
+      // For switch_value instructions, report the case number that doesn't
+      // yield.
+      if (auto *switchValue = dyn_cast<SwitchValueInst>(conflictingBranch)) {
+        auto enumCaseLoc = noYieldTarget->begin()->getLoc().getSourceLoc();
+
+        if (switchValue->hasDefault() &&
+            switchValue->getDefaultBB() == noYieldTarget) {
+          diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, Default);
+          return;
+        }
+        // Find the case that doesn't yield.
+        Optional<unsigned> caseNumberOpt =
+            switchValue->getUniqueCaseForDestination(noYieldTarget);
+        assert(caseNumberOpt.hasValue());
+
+        auto caseNumber = caseNumberOpt.getValue();
+        diagnose(
+            astCtx, enumCaseLoc, diag::switch_value_case_doesnt_yield,
+            (Twine(caseNumber) + llvm::getOrdinalSuffix(caseNumber)).str());
+        return;
+      }
+
+      // For try-apply instructions, report whether throwing or non-throwing
+      // case doesn't yield.
+      if (auto *tryApply = dyn_cast<TryApplyInst>(conflictingBranch)) {
+        diagnose(astCtx, tryApply->getLoc().getSourceLoc(),
+                 diag::try_branch_doesnt_yield,
+                 /*does error case not yield?*/ tryApply->getErrorBB() ==
+                     noYieldTarget);
+        return;
+      }
+
+      llvm_unreachable("unexpected branch resulting in conflicting yield "
+                       "states found in generalized accessor");
     }
     }
   }