[ConstraintSystem] Add a new locator element to denote branches of ternary operator

`TernaryBranch` with a boolean flag to identify "then" (true) or
"else" (false) branches of the ternary operator expression.

Author: xedin

lib/Sema/ConstraintLocator.cpp

@@ -61,16 +61,17 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
     case ConditionalRequirement:
     case TypeParameterRequirement:
     case ContextualType:
-    case SynthesizedArgument: {
+    case SynthesizedArgument:
+    case TernaryBranch: {
       auto numValues = numNumericValuesInPathElement(elt.getKind());
       for (unsigned i = 0; i < numValues; ++i)
         id.AddInteger(elt.getValue(i));
       break;
     }
 #define SIMPLE_LOCATOR_PATH_ELT(Name) case Name :
 #include "ConstraintLocatorPathElts.def"
-      // Nothing to do for simple locator elements.
-      break;
+    // Nothing to do for simple locator elements.
+    break;
     }
   }
 }
@@ -115,6 +116,7 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   case ConstraintLocator::Condition:
   case ConstraintLocator::DynamicCallable:
   case ConstraintLocator::ImplicitCallAsFunction:
+  case ConstraintLocator::TernaryBranch:
     return 0;
 
   case ConstraintLocator::FunctionArgument:
@@ -463,6 +465,12 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) const {
     case ImplicitCallAsFunction:
       out << "implicit reference to callAsFunction";
       break;
+
+    case TernaryBranch:
+      auto branchElt = elt.castTo<LocatorPathElt::TernaryBranch>();
+      out << (branchElt.forThen() ? "'then'" : "'else'")
+          << " branch of a ternary operator";
+      break;
     }
   }
   out << ']';

lib/Sema/ConstraintLocator.h

@@ -78,6 +78,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case KeyPathComponent:
     case SynthesizedArgument:
     case KeyPathDynamicMember:
+    case TernaryBranch:
       return 1;
 
     case TypeParameterRequirement:
@@ -782,6 +783,20 @@ class LocatorPathElt::KeyPathDynamicMember final : public LocatorPathElt {
   }
 };
 
+class LocatorPathElt::TernaryBranch final : public LocatorPathElt {
+public:
+  TernaryBranch(bool side)
+      : LocatorPathElt(ConstraintLocator::TernaryBranch, side) {}
+
+  bool forThen() const { return bool(getValue(0)); }
+
+  bool forElse() const { return !bool(getValue(0)); }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == ConstraintLocator::TernaryBranch;
+  }
+};
+
 /// A simple stack-only builder object that constructs a
 /// constraint locator without allocating memory.
 ///

lib/Sema/ConstraintLocatorPathElts.def

@@ -172,6 +172,9 @@ SIMPLE_LOCATOR_PATH_ELT(Condition)
 
 SIMPLE_LOCATOR_PATH_ELT(DynamicCallable)
 
+/// The 'true' or 'false' branch of a ternary operator.
+CUSTOM_LOCATOR_PATH_ELT(TernaryBranch)
+
 #undef LOCATOR_PATH_ELT
 #undef CUSTOM_LOCATOR_PATH_ELT
 #undef SIMPLE_LOCATOR_PATH_ELT

lib/Sema/ConstraintSystem.cpp

@@ -3313,6 +3313,15 @@ void constraints::simplifyLocator(Expr *&anchor,
       continue;
     }
 
+    case ConstraintLocator::TernaryBranch: {
+      auto branch = path[0].castTo<LocatorPathElt::TernaryBranch>();
+      auto *ifExpr = cast<IfExpr>(anchor);
+
+      anchor = branch.forThen() ? ifExpr->getThenExpr() : ifExpr->getElseExpr();
+      path = path.slice(1);
+      continue;
+    }
+
     default:
       // FIXME: Lots of other cases to handle.
       break;