Merge pull request #37629 from hamishknight/ghost-of-objc-past

lib/IDE/Refactoring.cpp

@@ -16,6 +16,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticsRefactoring.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/ForeignAsyncConvention.h"
 #include "swift/AST/GenericParamList.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/Pattern.h"
@@ -4418,44 +4419,114 @@ struct AsyncHandlerParamDesc : public AsyncHandlerDesc {
   }
 };
 
-enum class ConditionType { INVALID, NIL, NOT_NIL };
+/// The type of a condition in a conditional statement.
+enum class ConditionType {
+  NIL,             // == nil
+  NOT_NIL,         // != nil
+  IS_TRUE,         // if b
+  IS_FALSE,        // if !b
+  SUCCESS_PATTERN, // case .success
+  FAILURE_PATTEN   // case .failure
+};
+
+/// Indicates whether a condition describes a success or failure path. For
+/// example, a check for whether an error parameter is present is a failure
+/// path. A check for a nil error parameter is a success path. This is distinct
+/// from ConditionType, as it relies on contextual information about what values
+/// need to be checked for success or failure.
+enum class ConditionPath { SUCCESS, FAILURE };
+
+static ConditionPath flippedConditionPath(ConditionPath Path) {
+  switch (Path) {
+  case ConditionPath::SUCCESS:
+    return ConditionPath::FAILURE;
+  case ConditionPath::FAILURE:
+    return ConditionPath::SUCCESS;
+  }
+  llvm_unreachable("Unhandled case in switch!");
+}
 
 /// Finds the `Subject` being compared to in various conditions. Also finds any
 /// pattern that may have a bound name.
 struct CallbackCondition {
-  ConditionType Type = ConditionType::INVALID;
+  Optional<ConditionType> Type;
   const Decl *Subject = nullptr;
   const Pattern *BindPattern = nullptr;
-  // Bit of a hack. When the `Subject` is a `Result` type we use this to
-  // distinguish between the `.success` and `.failure` case (as opposed to just
-  // checking whether `Subject` == `TheErrDecl`)
-  bool ErrorCase = false;
-
-  CallbackCondition() = default;
 
   /// Initializes a `CallbackCondition` with a `!=` or `==` comparison of
-  /// an `Optional` typed `Subject` to `nil`, ie.
+  /// an `Optional` typed `Subject` to `nil`, or a `Bool` typed `Subject` to a
+  /// boolean literal, ie.
   ///   - `<Subject> != nil`
   ///   - `<Subject> == nil`
+  ///   - `<Subject> != true`
+  ///   - `<Subject> == false`
   CallbackCondition(const BinaryExpr *BE, const FuncDecl *Operator) {
     bool FoundNil = false;
+    BooleanLiteralExpr *FoundBool = nullptr;
+    bool DidUnwrapOptional = false;
+
     for (auto *Operand : {BE->getLHS(), BE->getRHS()}) {
+      Operand = Operand->getSemanticsProvidingExpr();
+      if (auto *IIOE = dyn_cast<InjectIntoOptionalExpr>(Operand)) {
+        Operand = IIOE->getSubExpr()->getSemanticsProvidingExpr();
+        DidUnwrapOptional = true;
+      }
       if (isa<NilLiteralExpr>(Operand)) {
         FoundNil = true;
+      } else if (auto *BLE = dyn_cast<BooleanLiteralExpr>(Operand)) {
+        FoundBool = BLE;
       } else if (auto *DRE = dyn_cast<DeclRefExpr>(Operand)) {
         Subject = DRE->getDecl();
       }
     }
 
-    if (Subject && FoundNil) {
+    if (!Subject)
+      return;
+
+    if (FoundNil) {
       if (Operator->getBaseName() == "==") {
         Type = ConditionType::NIL;
       } else if (Operator->getBaseName() == "!=") {
         Type = ConditionType::NOT_NIL;
       }
+    } else if (FoundBool) {
+      if (Operator->getBaseName() == "==") {
+        Type = FoundBool->getValue() ? ConditionType::IS_TRUE
+                                     : ConditionType::IS_FALSE;
+      } else if (Operator->getBaseName() == "!=" && !DidUnwrapOptional) {
+        // Note that we don't consider this case if we unwrapped an optional,
+        // as e.g optBool != false is a check for true *or* nil.
+        Type = FoundBool->getValue() ? ConditionType::IS_FALSE
+                                     : ConditionType::IS_TRUE;
+      }
     }
   }
 
+  /// A bool condition expression.
+  explicit CallbackCondition(const Expr *E) {
+    if (!E->getType()->isBool())
+      return;
+
+    auto CondType = ConditionType::IS_TRUE;
+    E = E->getSemanticsProvidingExpr();
+
+    // If we have a prefix negation operator, this is a check for false.
+    if (auto *PrefixOp = dyn_cast<PrefixUnaryExpr>(E)) {
+      auto *Callee = PrefixOp->getCalledValue();
+      if (Callee && Callee->isOperator() && Callee->getBaseName() == "!") {
+        CondType = ConditionType::IS_FALSE;
+        E = PrefixOp->getArg()->getSemanticsProvidingExpr();
+      }
+    }
+
+    auto *DRE = dyn_cast<DeclRefExpr>(E);
+    if (!DRE)
+      return;
+
+    Subject = DRE->getDecl();
+    Type = CondType;
+  }
+
   /// Initializes a `CallbackCondition` with binding of an `Optional` or
   /// `Result` typed `Subject`, ie.
   ///   - `let bind = <Subject>`
@@ -4495,67 +4566,19 @@ struct CallbackCondition {
     }
   }
 
-  bool isValid() const { return Type != ConditionType::INVALID; }
-
-  /// Given an `if` condition `Cond` and a set of `Decls`, find any
-  /// `CallbackCondition`s in `Cond` that use one of those `Decls` and add them
-  /// to the map `AddTo`. Return `true` if all elements in the condition are
-  /// "handled", ie. every condition can be mapped to a single `Decl` in
-  /// `Decls`.
-  static bool all(StmtCondition Cond, llvm::DenseSet<const Decl *> Decls,
-                  llvm::DenseMap<const Decl *, CallbackCondition> &AddTo) {
-    bool Handled = true;
-    for (auto &CondElement : Cond) {
-      if (auto *BoolExpr = CondElement.getBooleanOrNull()) {
-        SmallVector<Expr *, 1> Exprs;
-        Exprs.push_back(BoolExpr);
-
-        while (!Exprs.empty()) {
-          auto *Next = Exprs.pop_back_val();
-          if (auto *ACE = dyn_cast<AutoClosureExpr>(Next))
-            Next = ACE->getSingleExpressionBody();
-
-          if (auto *BE = dyn_cast_or_null<BinaryExpr>(Next)) {
-            auto *Operator = isOperator(BE);
-            if (Operator) {
-              if (Operator->getBaseName() == "&&") {
-                Exprs.push_back(BE->getLHS());
-                Exprs.push_back(BE->getRHS());
-              } else {
-                addCond(CallbackCondition(BE, Operator), Decls, AddTo, Handled);
-              }
-              continue;
-            }
-          }
-
-          Handled = false;
-        }
-      } else if (auto *P = CondElement.getPatternOrNull()) {
-        addCond(CallbackCondition(P, CondElement.getInitializer()), Decls,
-                AddTo, Handled);
-      }
-    }
-    return Handled && !AddTo.empty();
-  }
+  bool isValid() const { return Type.hasValue(); }
 
 private:
-  static void addCond(const CallbackCondition &CC,
-                      llvm::DenseSet<const Decl *> Decls,
-                      llvm::DenseMap<const Decl *, CallbackCondition> &AddTo,
-                      bool &Handled) {
-    if (!CC.isValid() || !Decls.count(CC.Subject) ||
-        !AddTo.try_emplace(CC.Subject, CC).second)
-      Handled = false;
-  }
-
   void initFromEnumPattern(const Decl *D, const EnumElementPattern *EEP) {
     if (auto *EED = EEP->getElementDecl()) {
       auto eedTy = EED->getParentEnum()->getDeclaredType();
       if (!eedTy || !eedTy->isResult())
         return;
-      if (EED->getNameStr() == StringRef("failure"))
-        ErrorCase = true;
-      Type = ConditionType::NOT_NIL;
+      if (EED->getNameStr() == StringRef("failure")) {
+        Type = ConditionType::FAILURE_PATTEN;
+      } else {
+        Type = ConditionType::SUCCESS_PATTERN;
+      }
       Subject = D;
       BindPattern = EEP->getSubPattern();
     }
@@ -4589,6 +4612,32 @@ private:
   }
 };
 
+/// A CallbackCondition with additional semantic information about whether it
+/// is for a success path or failure path.
+struct ClassifiedCondition : public CallbackCondition {
+  ConditionPath Path;
+
+  /// Whether this represents an Obj-C style boolean flag check for success.
+  bool IsObjCStyleFlagCheck;
+
+  explicit ClassifiedCondition(CallbackCondition Cond, ConditionPath Path,
+                               bool IsObjCStyleFlagCheck)
+      : CallbackCondition(Cond), Path(Path),
+        IsObjCStyleFlagCheck(IsObjCStyleFlagCheck) {}
+};
+
+/// A wrapper for a map of parameter decls to their classified conditions, or
+/// \c None if they are not present in any conditions.
+struct ClassifiedCallbackConditions final
+    : llvm::MapVector<const Decl *, ClassifiedCondition> {
+  Optional<ClassifiedCondition> lookup(const Decl *D) const {
+    auto Res = find(D);
+    if (Res == end())
+      return None;
+    return Res->second;
+  }
+};
+
 /// A list of nodes to print, along with a list of locations that may have
 /// preceding comments attached, which also need printing. For example:
 ///
@@ -4734,7 +4783,7 @@ public:
     Nodes.addNode(Node);
   }
 
-  void addBinding(const CallbackCondition &FromCondition,
+  void addBinding(const ClassifiedCondition &FromCondition,
                   DiagnosticEngine &DiagEngine) {
     if (!FromCondition.BindPattern)
       return;
@@ -4758,9 +4807,8 @@ public:
     BoundNames.try_emplace(FromCondition.Subject, Name);
   }
 
-  void addAllBindings(
-      const llvm::DenseMap<const Decl *, CallbackCondition> &FromConditions,
-      DiagnosticEngine &DiagEngine) {
+  void addAllBindings(const ClassifiedCallbackConditions &FromConditions,
+                      DiagnosticEngine &DiagEngine) {
     for (auto &Entry : FromConditions) {
       addBinding(Entry.second, DiagEngine);
       if (DiagEngine.hadAnyError())
@@ -4769,6 +4817,23 @@ public:
   }
 };
 
+/// Whether or not the given statement starts a new scope. Note that most
+/// statements are handled by the \c BraceStmt check. The others listed are
+/// a somewhat special case since they can also declare variables in their
+/// condition.
+static bool startsNewScope(Stmt *S) {
+  switch (S->getKind()) {
+  case StmtKind::Brace:
+  case StmtKind::If:
+  case StmtKind::While:
+  case StmtKind::ForEach:
+  case StmtKind::Case:
+    return true;
+  default:
+    return false;
+  }
+}
+
 struct ClassifiedBlocks {
   ClassifiedBlock SuccessBlock;
   ClassifiedBlock ErrorBlock;
@@ -4789,21 +4854,24 @@ struct CallbackClassifier {
   /// Updates the success and error block of `Blocks` with nodes and bound
   /// names from `Body`. Errors are added through `DiagEngine`, possibly
   /// resulting in partially filled out blocks.
-  static void classifyInto(ClassifiedBlocks &Blocks,
+  static void classifyInto(ClassifiedBlocks &Blocks, const FuncDecl *Callee,
+                           ArrayRef<const ParamDecl *> SuccessParams,
                            llvm::DenseSet<SwitchStmt *> &HandledSwitches,
                            DiagnosticEngine &DiagEngine,
                            llvm::DenseSet<const Decl *> UnwrapParams,
                            const ParamDecl *ErrParam, HandlerType ResultType,
                            BraceStmt *Body) {
     assert(!Body->getElements().empty() && "Cannot classify empty body");
-    CallbackClassifier Classifier(Blocks, HandledSwitches, DiagEngine,
-                                  UnwrapParams, ErrParam,
-                                  ResultType == HandlerType::RESULT);
+    CallbackClassifier Classifier(Blocks, Callee, SuccessParams,
+                                  HandledSwitches, DiagEngine, UnwrapParams,
+                                  ErrParam, ResultType == HandlerType::RESULT);
     Classifier.classifyNodes(Body->getElements(), Body->getRBraceLoc());
   }
 
 private:
   ClassifiedBlocks &Blocks;
+  const FuncDecl *Callee;
+  ArrayRef<const ParamDecl *> SuccessParams;
   llvm::DenseSet<SwitchStmt *> &HandledSwitches;
   DiagnosticEngine &DiagEngine;
   ClassifiedBlock *CurrentBlock;
@@ -4811,15 +4879,16 @@ private:
   const ParamDecl *ErrParam;
   bool IsResultParam;
 
-  CallbackClassifier(ClassifiedBlocks &Blocks,
+  CallbackClassifier(ClassifiedBlocks &Blocks, const FuncDecl *Callee,
+                     ArrayRef<const ParamDecl *> SuccessParams,
                      llvm::DenseSet<SwitchStmt *> &HandledSwitches,
                      DiagnosticEngine &DiagEngine,
                      llvm::DenseSet<const Decl *> UnwrapParams,
                      const ParamDecl *ErrParam, bool IsResultParam)
-      : Blocks(Blocks), HandledSwitches(HandledSwitches),
-        DiagEngine(DiagEngine), CurrentBlock(&Blocks.SuccessBlock),
-        UnwrapParams(UnwrapParams), ErrParam(ErrParam),
-        IsResultParam(IsResultParam) {}
+      : Blocks(Blocks), Callee(Callee), SuccessParams(SuccessParams),
+        HandledSwitches(HandledSwitches), DiagEngine(DiagEngine),
+        CurrentBlock(&Blocks.SuccessBlock), UnwrapParams(UnwrapParams),
+        ErrParam(ErrParam), IsResultParam(IsResultParam) {}
 
   void classifyNodes(ArrayRef<ASTNode> Nodes, SourceLoc endCommentLoc) {
     for (auto I = Nodes.begin(), E = Nodes.end(); I < E; ++I) {
@@ -4849,12 +4918,269 @@ private:
     CurrentBlock->addPossibleCommentLoc(endCommentLoc);
   }
 
+  /// Whether any of the provided ASTNodes have a child expression that force
+  /// unwraps the error parameter. Note that this doesn't walk into new scopes.
+  bool hasForceUnwrappedErrorParam(ArrayRef<ASTNode> Nodes) {
+    if (IsResultParam || !ErrParam)
+      return false;
+
+    class ErrUnwrapFinder : public ASTWalker {
+      const ParamDecl *ErrParam;
+      bool FoundUnwrap = false;
+
+    public:
+      explicit ErrUnwrapFinder(const ParamDecl *ErrParam)
+          : ErrParam(ErrParam) {}
+      bool foundUnwrap() const { return FoundUnwrap; }
+
+      std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+        // Don't walk into ternary conditionals as they may have additional
+        // conditions such as err != nil that make a force unwrap now valid.
+        if (isa<IfExpr>(E))
+          return {false, E};
+
+        auto *FVE = dyn_cast<ForceValueExpr>(E);
+        if (!FVE)
+          return {true, E};
+
+        auto *DRE = dyn_cast<DeclRefExpr>(FVE->getSubExpr());
+        if (!DRE)
+          return {true, E};
+
+        if (DRE->getDecl() != ErrParam)
+          return {true, E};
+
+        // If we find the node we're looking for, make a note of it, and abort
+        // the walk.
+        FoundUnwrap = true;
+        return {false, nullptr};
+      }
+
+      std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+        // Don't walk into new explicit scopes, we only want to consider force
+        // unwraps in the immediate conditional body.
+        if (!S->isImplicit() && startsNewScope(S))
+          return {false, S};
+        return {true, S};
+      }
+
+      bool walkToDeclPre(Decl *D) override {
+        // Don't walk into new explicit DeclContexts.
+        return D->isImplicit() || !isa<DeclContext>(D);
+      }
+    };
+    for (auto Node : Nodes) {
+      ErrUnwrapFinder walker(ErrParam);
+      Node.walk(walker);
+      if (walker.foundUnwrap())
+        return true;
+    }
+    return false;
+  }
+
+  /// Given a callback condition, classify it as a success or failure path.
+  Optional<ClassifiedCondition>
+  classifyCallbackCondition(const CallbackCondition &Cond,
+                            const NodesToPrint &SuccessNodes, Stmt *ElseStmt) {
+    if (!Cond.isValid())
+      return None;
+
+    // For certain types of condition, they need to appear in certain lists.
+    auto CondType = *Cond.Type;
+    switch (CondType) {
+    case ConditionType::NOT_NIL:
+    case ConditionType::NIL:
+      if (!UnwrapParams.count(Cond.Subject))
+        return None;
+      break;
+    case ConditionType::IS_TRUE:
+    case ConditionType::IS_FALSE:
+      if (!llvm::is_contained(SuccessParams, Cond.Subject))
+        return None;
+      break;
+    case ConditionType::SUCCESS_PATTERN:
+    case ConditionType::FAILURE_PATTEN:
+      if (!IsResultParam || Cond.Subject != ErrParam)
+        return None;
+      break;
+    }
+
+    // Let's start with a success path, and flip any negative conditions.
+    auto Path = ConditionPath::SUCCESS;
+
+    // If it's an error param, that's a flip.
+    if (Cond.Subject == ErrParam && !IsResultParam)
+      Path = flippedConditionPath(Path);
+
+    // If we have a nil, false, or failure condition, that's a flip.
+    switch (CondType) {
+    case ConditionType::NIL:
+    case ConditionType::IS_FALSE:
+    case ConditionType::FAILURE_PATTEN:
+      Path = flippedConditionPath(Path);
+      break;
+    case ConditionType::IS_TRUE:
+    case ConditionType::NOT_NIL:
+    case ConditionType::SUCCESS_PATTERN:
+      break;
+    }
+
+    // If we have a bool condition, it could be an Obj-C style flag check, which
+    // we do some extra checking for. Otherwise, we're done.
+    if (CondType != ConditionType::IS_TRUE &&
+        CondType != ConditionType::IS_FALSE) {
+      return ClassifiedCondition(Cond, Path, /*ObjCFlagCheck*/ false);
+    }
+
+    // Check to see if the async alternative function has a convention that
+    // specifies where the flag is and what it indicates.
+    Optional<std::pair</*Idx*/ unsigned, /*SuccessFlag*/ bool>> CustomFlag;
+    if (auto *AsyncAlt = Callee->getAsyncAlternative()) {
+      if (auto Conv = AsyncAlt->getForeignAsyncConvention()) {
+        if (auto Idx = Conv->completionHandlerFlagParamIndex()) {
+          auto IsSuccessFlag = Conv->completionHandlerFlagIsErrorOnZero();
+          CustomFlag = std::make_pair(*Idx, IsSuccessFlag);
+        }
+      }
+    }
+    if (CustomFlag) {
+      auto Idx = CustomFlag->first;
+      if (Idx < 0 || Idx >= SuccessParams.size())
+        return None;
+
+      if (SuccessParams[Idx] != Cond.Subject)
+        return None;
+
+      // The path may need to be flipped depending on whether the flag indicates
+      // success.
+      auto IsSuccessFlag = CustomFlag->second;
+      if (!IsSuccessFlag)
+        Path = flippedConditionPath(Path);
+
+      return ClassifiedCondition(Cond, Path, /*ObjCStyleFlagCheck*/ true);
+    }
+
+    // If we've reached here, we have a bool flag check that isn't specified in
+    // the async convention. We apply a heuristic to see if the error param is
+    // force unwrapped in the conditional body. In that case, the user is
+    // expecting it to be the error path, and it's more likely than not that the
+    // flag value conveys no more useful information in the error block.
+
+    // First check the success block.
+    auto FoundInSuccessBlock =
+        hasForceUnwrappedErrorParam(SuccessNodes.getNodes());
+
+    // Then check the else block if we have it.
+    if (ASTNode ElseNode = ElseStmt) {
+      // Unwrap the BraceStmt of the else clause if needed. This is needed as
+      // we won't walk into BraceStmts by default as they introduce new
+      // scopes.
+      ArrayRef<ASTNode> Nodes;
+      if (auto *BS = dyn_cast<BraceStmt>(ElseStmt)) {
+        Nodes = BS->getElements();
+      } else {
+        Nodes = llvm::makeArrayRef(ElseNode);
+      }
+      if (hasForceUnwrappedErrorParam(Nodes)) {
+        // If we also found an unwrap in the success block, we don't know what's
+        // happening here.
+        if (FoundInSuccessBlock)
+          return None;
+
+        // Otherwise we can determine this as a success condition. Note this is
+        // flipped as if the error is present in the else block, this condition
+        // is for success.
+        return ClassifiedCondition(Cond, ConditionPath::SUCCESS,
+                                   /*ObjCStyleFlagCheck*/ true);
+      }
+    }
+
+    if (FoundInSuccessBlock) {
+      // Note that the path is flipped as if the error is present in the success
+      // block, this condition is for failure.
+      return ClassifiedCondition(Cond, ConditionPath::FAILURE,
+                                 /*ObjCStyleFlagCheck*/ true);
+    }
+
+    // Otherwise we can't classify this.
+    return None;
+  }
+
+  /// Classifies all the conditions present in a given StmtCondition, taking
+  /// into account its success body and failure body. Returns \c true if there
+  /// were any conditions that couldn't be classified, \c false otherwise.
+  bool classifyConditionsOf(StmtCondition Cond,
+                            const NodesToPrint &ThenNodesToPrint,
+                            Stmt *ElseStmt,
+                            ClassifiedCallbackConditions &Conditions) {
+    bool UnhandledConditions = false;
+    Optional<ClassifiedCondition> ObjCFlagCheck;
+    auto TryAddCond = [&](CallbackCondition CC) {
+      auto Classified =
+          classifyCallbackCondition(CC, ThenNodesToPrint, ElseStmt);
+
+      // If we couldn't classify this, or if there are multiple Obj-C style flag
+      // checks, this is unhandled.
+      if (!Classified || (ObjCFlagCheck && Classified->IsObjCStyleFlagCheck)) {
+        UnhandledConditions = true;
+        return;
+      }
+
+      // If we've seen multiple conditions for the same subject, don't handle
+      // this.
+      if (!Conditions.insert({CC.Subject, *Classified}).second) {
+        UnhandledConditions = true;
+        return;
+      }
+
+      if (Classified->IsObjCStyleFlagCheck)
+        ObjCFlagCheck = Classified;
+    };
+
+    for (auto &CondElement : Cond) {
+      if (auto *BoolExpr = CondElement.getBooleanOrNull()) {
+        SmallVector<Expr *, 1> Exprs;
+        Exprs.push_back(BoolExpr);
+
+        while (!Exprs.empty()) {
+          auto *Next = Exprs.pop_back_val()->getSemanticsProvidingExpr();
+          if (auto *ACE = dyn_cast<AutoClosureExpr>(Next))
+            Next = ACE->getSingleExpressionBody()->getSemanticsProvidingExpr();
+
+          if (auto *BE = dyn_cast_or_null<BinaryExpr>(Next)) {
+            auto *Operator = isOperator(BE);
+            if (Operator) {
+              // If we have an && operator, decompose its arguments.
+              if (Operator->getBaseName() == "&&") {
+                Exprs.push_back(BE->getLHS());
+                Exprs.push_back(BE->getRHS());
+              } else {
+                // Otherwise check to see if we have an == nil or != nil
+                // condition.
+                TryAddCond(CallbackCondition(BE, Operator));
+              }
+              continue;
+            }
+          }
+
+          // Check to see if we have a lone bool condition.
+          TryAddCond(CallbackCondition(Next));
+        }
+      } else if (auto *P = CondElement.getPatternOrNull()) {
+        TryAddCond(CallbackCondition(P, CondElement.getInitializer()));
+      }
+    }
+    return UnhandledConditions || Conditions.empty();
+  }
+
+  /// Classifies the conditions of a conditional statement, and adds the
+  /// necessary nodes to either the success or failure block.
   void classifyConditional(Stmt *Statement, StmtCondition Condition,
                            NodesToPrint ThenNodesToPrint, Stmt *ElseStmt) {
-    llvm::DenseMap<const Decl *, CallbackCondition> CallbackConditions;
-    bool UnhandledConditions =
-        !CallbackCondition::all(Condition, UnwrapParams, CallbackConditions);
-    CallbackCondition ErrCondition = CallbackConditions.lookup(ErrParam);
+    ClassifiedCallbackConditions CallbackConditions;
+    bool UnhandledConditions = classifyConditionsOf(
+        Condition, ThenNodesToPrint, ElseStmt, CallbackConditions);
+    auto ErrCondition = CallbackConditions.lookup(ErrParam);
 
     if (UnhandledConditions) {
       // Some unknown conditions. If there's an else, assume we can't handle
@@ -4870,12 +5196,11 @@ private:
       } else if (ElseStmt) {
         DiagEngine.diagnose(Statement->getStartLoc(),
                             diag::unknown_callback_conditions);
-      } else if (ErrCondition.isValid() &&
-                 ErrCondition.Type == ConditionType::NOT_NIL) {
+      } else if (ErrCondition && ErrCondition->Path == ConditionPath::FAILURE) {
         Blocks.ErrorBlock.addNode(Statement);
       } else {
         for (auto &Entry : CallbackConditions) {
-          if (Entry.second.Type == ConditionType::NIL) {
+          if (Entry.second.Path == ConditionPath::FAILURE) {
             Blocks.ErrorBlock.addNode(Statement);
             return;
           }
@@ -4885,42 +5210,36 @@ private:
       return;
     }
 
-    ClassifiedBlock *ThenBlock = &Blocks.SuccessBlock;
-    ClassifiedBlock *ElseBlock = &Blocks.ErrorBlock;
+    // If all the conditions were classified, make sure they're all consistently
+    // on the success or failure path.
+    Optional<ConditionPath> Path;
+    for (auto &Entry : CallbackConditions) {
+      auto &Cond = Entry.second;
+      if (!Path) {
+        Path = Cond.Path;
+      } else if (*Path != Cond.Path) {
+        // Similar to the unknown conditions case. Add the whole if unless
+        // there's an else, in which case use the fallback instead.
+        // TODO: Split the `if` statement
 
-    if (ErrCondition.isValid() && (!IsResultParam || ErrCondition.ErrorCase) &&
-        ErrCondition.Type == ConditionType::NOT_NIL) {
-      ClassifiedBlock *TempBlock = ThenBlock;
-      ThenBlock = ElseBlock;
-      ElseBlock = TempBlock;
-    } else {
-      ConditionType CondType = ConditionType::INVALID;
-      for (auto &Entry : CallbackConditions) {
-        if (IsResultParam || Entry.second.Subject != ErrParam) {
-          if (CondType == ConditionType::INVALID) {
-            CondType = Entry.second.Type;
-          } else if (CondType != Entry.second.Type) {
-            // Similar to the unknown conditions case. Add the whole if unless
-            // there's an else, in which case use the fallback instead.
-            // TODO: Split the `if` statement
-
-            if (ElseStmt) {
-              DiagEngine.diagnose(Statement->getStartLoc(),
-                                  diag::mixed_callback_conditions);
-            } else {
-              CurrentBlock->addNode(Statement);
-            }
-            return;
-          }
+        if (ElseStmt) {
+          DiagEngine.diagnose(Statement->getStartLoc(),
+                              diag::mixed_callback_conditions);
+        } else {
+          CurrentBlock->addNode(Statement);
         }
-      }
-
-      if (CondType == ConditionType::NIL) {
-        ClassifiedBlock *TempBlock = ThenBlock;
-        ThenBlock = ElseBlock;
-        ElseBlock = TempBlock;
+        return;
       }
     }
+    assert(Path && "Didn't classify a path?");
+
+    auto *ThenBlock = &Blocks.SuccessBlock;
+    auto *ElseBlock = &Blocks.ErrorBlock;
+
+    // If the condition is for a failure path, the error block is ThenBlock, and
+    // the success block is ElseBlock.
+    if (*Path == ConditionPath::FAILURE)
+      std::swap(ThenBlock, ElseBlock);
 
     // We'll be dropping the statement, but make sure to keep any attached
     // comments.
@@ -4991,13 +5310,16 @@ private:
         return;
       }
 
-      CallbackCondition CC(ErrParam, &Items[0]);
-      ClassifiedBlock *Block = &Blocks.SuccessBlock;
-      ClassifiedBlock *OtherBlock = &Blocks.ErrorBlock;
-      if (CC.ErrorCase) {
-        Block = &Blocks.ErrorBlock;
-        OtherBlock = &Blocks.SuccessBlock;
-      }
+      auto *Block = &Blocks.SuccessBlock;
+      auto *OtherBlock = &Blocks.ErrorBlock;
+      auto SuccessNodes = NodesToPrint::inBraceStmt(CS->getBody());
+
+      // Classify the case pattern.
+      auto CC = classifyCallbackCondition(
+          CallbackCondition(ErrParam, &Items[0]), SuccessNodes,
+          /*elseStmt*/ nullptr);
+      if (CC && CC->Path == ConditionPath::FAILURE)
+        std::swap(Block, OtherBlock);
 
       // We'll be dropping the case, but make sure to keep any attached
       // comments. Because these comments will effectively be part of the
@@ -5008,8 +5330,9 @@ private:
       if (CS == Cases.back())
         Block->addPossibleCommentLoc(SS->getRBraceLoc());
 
-      setNodes(Block, OtherBlock, NodesToPrint::inBraceStmt(CS->getBody()));
-      Block->addBinding(CC, DiagEngine);
+      setNodes(Block, OtherBlock, std::move(SuccessNodes));
+      if (CC)
+        Block->addBinding(*CC, DiagEngine);
       if (DiagEngine.hadAnyError())
         return;
     }
@@ -5018,23 +5341,6 @@ private:
   }
 };
 
-/// Whether or not the given statement starts a new scope. Note that most
-/// statements are handled by the \c BraceStmt check. The others listed are
-/// a somewhat special case since they can also declare variables in their
-/// condition.
-static bool startsNewScope(Stmt *S) {
-  switch (S->getKind()) {
-  case StmtKind::Brace:
-  case StmtKind::If:
-  case StmtKind::While:
-  case StmtKind::ForEach:
-  case StmtKind::Case:
-    return true;
-  default:
-    return false;
-  }
-}
-
 /// Name of a decl if it has one, an empty \c Identifier otherwise.
 static Identifier getDeclName(const Decl *D) {
   if (auto *VD = dyn_cast<ValueDecl>(D)) {
@@ -6070,6 +6376,9 @@ private:
                                  const AsyncHandlerDesc &HandlerDesc,
                                  const ClosureExpr *Callback,
                                  PtrArrayRef<Expr *> ArgList) {
+    auto *Callee = getUnderlyingFunc(CE->getFn());
+    assert(Callee);
+
     ArrayRef<const ParamDecl *> CallbackParams =
         Callback->getParameters()->getArray();
     auto CallbackBody = Callback->getBody();
@@ -6101,9 +6410,9 @@ private:
       }
       if (ErrParam)
         UnwrapParams.insert(ErrParam);
-      CallbackClassifier::classifyInto(Blocks, HandledSwitches, DiagEngine,
-                                       UnwrapParams, ErrParam, HandlerDesc.Type,
-                                       CallbackBody);
+      CallbackClassifier::classifyInto(
+          Blocks, Callee, SuccessParams, HandledSwitches, DiagEngine,
+          UnwrapParams, ErrParam, HandlerDesc.Type, CallbackBody);
     }
 
     if (DiagEngine.hadAnyError()) {