Merge pull request #4773 from practicalswift/add-missing-crash-cases-and-regressions

diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
index 33204c3..0de84dc 100644
--- a/lib/Sema/CSSolver.cpp
+++ b/lib/Sema/CSSolver.cpp
@@ -1342,6 +1342,13 @@
   // Allocate new constraint system for sub-expression.
   ConstraintSystem cs(TC, DC, None);
 
+  // Set contextual type if present. This is done before constraint generation
+  // to give a "hint" to that operation about possible optimizations.
+  auto CT = IsPrimary ? CS.getContextualType() : CS.getContextualType(E);
+  if (!CT.isNull())
+    cs.setContextualType(E, CS.getContextualTypeLoc(),
+                         CS.getContextualTypePurpose());
+
   // Generate constraints for the new system.
   if (auto generatedExpr = cs.generateConstraints(E)) {
     E = generatedExpr;
@@ -1355,7 +1362,7 @@
   // constraint to the system.
   if (!CT.isNull()) {
     auto constraintKind = ConstraintKind::Conversion;
-    if (CTP == CTP_CallArgument)
+    if (CS.getContextualTypePurpose() == CTP_CallArgument)
       constraintKind = ConstraintKind::ArgumentConversion;
 
     cs.addConstraint(constraintKind, E->getType(), CT,
@@ -1469,21 +1476,52 @@
     std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
       // A dictionary expression is just a set of tuples; try to solve ones
       // that have overload sets.
-      if (auto collectionExpr = dyn_cast<CollectionExpr>(expr)) {
-        visitCollectionExpr(collectionExpr);
+      if (auto dictionaryExpr = dyn_cast<DictionaryExpr>(expr)) {
+        bool isPrimaryExpr = expr == PrimaryExpr;
+        for (auto element : dictionaryExpr->getElements()) {
+          unsigned numOverloads = 0;
+          element->walk(OverloadSetCounter(numOverloads));
+
+          // There are no overload sets in the element; skip it.
+          if (numOverloads == 0)
+            continue;
+
+          // FIXME: Could we avoid creating a separate dictionary expression
+          // here by introducing a contextual type on the element?
+          auto dict = DictionaryExpr::create(CS.getASTContext(),
+                                             dictionaryExpr->getLBracketLoc(),
+                                             { element },
+                                             dictionaryExpr->getRBracketLoc(),
+                                             dictionaryExpr->getType());
+
+          // Make each of the dictionary elements an independent dictionary,
+          // such makes it easy to type-check everything separately.
+          Candidates.push_back(Candidate(CS, dict, isPrimaryExpr));
+        }
+
         // Don't try to walk into the dictionary.
+        return { false, expr };
+      }
+
+      // Consider all of the collections to be candidates,
+      // FIXME: try to split collections into parts for simplified solving.
+      if (isa<CollectionExpr>(expr)) {
+        Candidates.push_back(Candidate(CS, expr, false));
         return {false, expr};
       }
 
       // Let's not attempt to type-check closures, which has already been
       // type checked anyway.
       if (isa<ClosureExpr>(expr)) {
-        return {false, expr};
+        return { false, expr };
       }
 
+      // Coerce to type expressions are only viable if they have
+      // a single child expression.
       if (auto coerceExpr = dyn_cast<CoerceExpr>(expr)) {
-        visitCoerceExpr(coerceExpr);
-        return {false, expr};
+        if (!coerceExpr->getSubExpr()) {
+          return { false, expr };
+        }
       }
 
       if (auto OSR = dyn_cast<OverloadSetRefExpr>(expr)) {
@@ -1501,24 +1539,12 @@
     }
 
     Expr *walkToExprPost(Expr *expr) override {
-      if (expr == PrimaryExpr) {
-        // If this is primary expression and there are no candidates
-        // to be solved, let's not record it, because it's going to be
-        // solved irregardless.
-        if (Candidates.empty())
-          return expr;
-
-        auto contextualType = CS.getContextualType();
-        // If there is a contextual type set for this expression.
-        if (!contextualType.isNull()) {
-          Candidates.push_back(Candidate(CS, expr, contextualType,
-                                         CS.getContextualTypePurpose()));
-          return expr;
-        }
-
-        // Or it's a function application with other candidates present.
-        if (isa<ApplyExpr>(expr)) {
-          Candidates.push_back(Candidate(CS, expr));
+      // If there are sub-expressions to consider and
+      // contextual type is involved, let's add top-most expression
+      // to the queue just to make sure that we didn't miss any solutions.
+      if (expr == PrimaryExpr && !Candidates.empty()) {
+        if (!CS.getContextualType().isNull()) {
+          Candidates.push_back(Candidate(CS, expr, true));
           return expr;
         }
       }
@@ -1548,157 +1574,10 @@
       // there is no point of solving this expression,
       // because we won't be able to reduce its domain.
       if (numOverloadSets > 1)
-        Candidates.push_back(Candidate(CS, expr));
+        Candidates.push_back(Candidate(CS, expr, expr == PrimaryExpr));
 
       return expr;
     }
-
-  private:
-    /// \brief Extract type of the element from given collection type.
-    ///
-    /// \param collection The type of the collection container.
-    ///
-    /// \returns ErrorType on failure, properly constructed type otherwise.
-    Type extractElementType(Type collection) {
-      auto &ctx = CS.getASTContext();
-      if (collection.isNull() || collection->is<ErrorType>())
-        return ErrorType::get(ctx);
-
-      auto base = collection.getPointer();
-      auto isInvalidType = [](Type type) -> bool {
-        return type.isNull() || type->hasUnresolvedType() ||
-               type->is<ErrorType>();
-      };
-
-      // Array type.
-      if (auto array = dyn_cast<ArraySliceType>(base)) {
-        auto elementType = array->getBaseType();
-        // If base type is invalid let's return error type.
-        return isInvalidType(elementType) ? ErrorType::get(ctx) : elementType;
-      }
-
-      // Map or Set or any other associated collection type.
-      if (auto boundGeneric = dyn_cast<BoundGenericType>(base)) {
-        if (boundGeneric->hasUnresolvedType())
-          return ErrorType::get(ctx);
-
-        llvm::SmallVector<TupleTypeElt, 2> params;
-        for (auto &type : boundGeneric->getGenericArgs()) {
-          // One of the generic arguments in invalid or unresolved.
-          if (isInvalidType(type))
-            return ErrorType::get(ctx);
-
-          params.push_back(type);
-        }
-
-        // If there is just one parameter, let's return it directly.
-        if (params.size() == 1)
-          return params[0].getType();
-
-        return TupleType::get(params, ctx);
-      }
-
-      return ErrorType::get(ctx);
-    }
-
-    bool isSuitableCollection(TypeRepr *collectionTypeRepr) {
-      // Only generic identifier, array or dictionary.
-      switch (collectionTypeRepr->getKind()) {
-      case TypeReprKind::GenericIdent:
-      case TypeReprKind::Array:
-      case TypeReprKind::Dictionary:
-        return true;
-
-      default:
-        return false;
-      }
-    }
-
-    void visitCoerceExpr(CoerceExpr *coerceExpr) {
-      auto subExpr = coerceExpr->getSubExpr();
-      // Coerce expression is valid only if it has sub-expression.
-      if (!subExpr) return;
-
-      unsigned numOverloadSets = 0;
-      subExpr->forEachChildExpr([&](Expr *childExpr) -> Expr * {
-        if (isa<OverloadSetRefExpr>(childExpr)) {
-          ++numOverloadSets;
-          return childExpr;
-        }
-
-        if (auto nestedCoerceExpr = dyn_cast<CoerceExpr>(childExpr)) {
-          visitCoerceExpr(nestedCoerceExpr);
-          // Don't walk inside of nested coercion expression directly,
-          // that is be done by recursive call to visitCoerceExpr.
-          return nullptr;
-        }
-
-        // If sub-expression we are trying to coerce to type is a collection,
-        // let's allow collector discover it with assigned contextual type
-        // of coercion, which allows collections to be solved in parts.
-        if (auto collectionExpr = dyn_cast<CollectionExpr>(childExpr)) {
-          auto castTypeLoc = coerceExpr->getCastTypeLoc();
-          auto typeRepr = castTypeLoc.getTypeRepr();
-
-          if (typeRepr && isSuitableCollection(typeRepr)) {
-            // Clone representative to avoid modifying in-place,
-            // FIXME: We should try and silently resolve the type here,
-            // instead of cloning representative.
-            auto coercionRepr = typeRepr->clone(CS.getASTContext());
-            // Let's try to resolve coercion type from cloned representative.
-            auto coercionType = CS.TC.resolveType(coercionRepr, CS.DC,
-                                                  TypeResolutionOptions());
-
-            // Looks like coercion type is invalid, let's skip this sub-tree.
-            if (coercionType->is<ErrorType>())
-              return nullptr;
-
-            // Visit collection expression inline.
-            visitCollectionExpr(collectionExpr, coercionType,
-                                CTP_CoerceOperand);
-          }
-        }
-
-        return childExpr;
-      });
-
-      // It's going to be inefficient to try and solve
-      // coercion in parts, so let's just make it a candidate directly,
-      // if it contains at least a single overload set.
-
-      if (numOverloadSets > 0)
-        Candidates.push_back(Candidate(CS, coerceExpr));
-    }
-
-    void visitCollectionExpr(CollectionExpr *collectionExpr,
-                             Type contextualType = Type(),
-                             ContextualTypePurpose CTP = CTP_Unused) {
-      // If there is a contextual type set for this collection,
-      // let's propagate it to the candidate.
-      if (!contextualType.isNull()) {
-        auto elementType = extractElementType(contextualType);
-        // If we couldn't deduce element type for the collection, let's
-        // not attempt to solve it.
-        if (elementType->is<ErrorType>())
-          return;
-
-        contextualType = elementType;
-      }
-
-      for (auto element : collectionExpr->getElements()) {
-        unsigned numOverloads = 0;
-        element->walk(OverloadSetCounter(numOverloads));
-
-        // There are no overload sets in the element; skip it.
-        if (numOverloads == 0)
-          continue;
-
-        // Record each of the collection elements, which passed
-        // number of overload sets rule, as a candidate for solving
-        // with contextual type of the collection.
-        Candidates.push_back(Candidate(CS, element, contextualType, CTP));
-      }
-    }
   };
 
   ExprCollector collector(expr, *this, domains);
diff --git a/lib/Sema/ConstraintSystem.h b/lib/Sema/ConstraintSystem.h
index 148982f..7d5620f 100644
--- a/lib/Sema/ConstraintSystem.h
+++ b/lib/Sema/ConstraintSystem.h
@@ -1022,17 +1022,15 @@
   /// to reduce scopes of the overload sets (disjunctions) in the system.
   class Candidate {
     Expr *E;
+    bool IsPrimary;
+
+    ConstraintSystem &CS;
     TypeChecker &TC;
     DeclContext *DC;
 
-    // Contextual Information.
-    Type CT;
-    ContextualTypePurpose CTP;
-
   public:
-    Candidate(ConstraintSystem &cs, Expr *expr, Type ct = Type(),
-              ContextualTypePurpose ctp = ContextualTypePurpose::CTP_Unused)
-        : E(expr), TC(cs.TC), DC(cs.DC), CT(ct), CTP(ctp) {}
+    Candidate(ConstraintSystem &cs, Expr *expr, bool primaryExpr)
+        : E(expr), IsPrimary(primaryExpr), CS(cs), TC(cs.TC), DC(cs.DC) {}
 
     /// \brief Return underlaying expression.
     Expr *getExpr() const { return E; }
diff --git a/test/Sema/complex_expressions.swift b/test/Sema/complex_expressions.swift
index bf1ff48..681fc7a 100644
--- a/test/Sema/complex_expressions.swift
+++ b/test/Sema/complex_expressions.swift
@@ -87,12 +87,3 @@
 func sr1794(pt: P, p0: P, p1: P) -> Bool {
   return (pt.x - p0.x) * (p1.y - p0.y) - (pt.y - p0.y) * (p1.x - p0.x) < 0.0
 }
-
-// Tests for partial contextual type application in sub-expressions
-
-let v1 = (1 - 2 / 3 * 6) as UInt
-let v2 = (([1 + 2 * 3, 4, 5])) as [UInt]
-let v3 = ["hello": 1 + 2, "world": 3 + 4 + 5 * 3] as Dictionary<String, UInt>
-let v4 = [1 + 2 + 3, 4] as [UInt32] + [2 * 3] as [UInt32]
-let v5 = ([1 + 2 + 3, 4] as [UInt32]) + ([2 * 3] as [UInt32])
-let v6 = [1 + 2 + 3, 4] as Set<UInt32>