[Type Checker] Extend Path Consistency algorithm to cover collections

xedin · xedin · commit f590a1ba0332 · 2016-09-13T01:34:14.000-07:00
ExprCollector is extended to cover all generic collections instead of
only dictionary expressions. Contextual type propagation is extended
to support partial solving of collections embedded into coerce expressions.
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -1342,13 +1342,6 @@ bool ConstraintSystem::Candidate::solve() {
   // 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;
@@ -1362,7 +1355,7 @@ bool ConstraintSystem::Candidate::solve() {
   // constraint to the system.
   if (!CT.isNull()) {
     auto constraintKind = ConstraintKind::Conversion;
-    if (CS.getContextualTypePurpose() == CTP_CallArgument)
+    if (CTP == CTP_CallArgument)
       constraintKind = ConstraintKind::ArgumentConversion;
 
     cs.addConstraint(constraintKind, E->getType(), CT,
@@ -1476,52 +1469,21 @@ void ConstraintSystem::shrink(Expr *expr) {
     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 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));
-        }
-
+      if (auto collectionExpr = dyn_cast<CollectionExpr>(expr)) {
+        visitCollectionExpr(collectionExpr);
         // 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)) {
-        if (!coerceExpr->getSubExpr()) {
-          return { false, expr };
-        }
+        visitCoerceExpr(coerceExpr);
+        return {false, expr};
       }
 
       if (auto OSR = dyn_cast<OverloadSetRefExpr>(expr)) {
@@ -1539,12 +1501,24 @@ void ConstraintSystem::shrink(Expr *expr) {
     }
 
     Expr *walkToExprPost(Expr *expr) override {
-      // 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));
+      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));
           return expr;
         }
       }
@@ -1574,10 +1548,157 @@ void ConstraintSystem::shrink(Expr *expr) {
       // 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, expr == PrimaryExpr));
+        Candidates.push_back(Candidate(CS, expr));
 
       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
@@ -1022,15 +1022,17 @@ class ConstraintSystem {
   /// 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, bool primaryExpr)
-        : E(expr), IsPrimary(primaryExpr), CS(cs), TC(cs.TC), DC(cs.DC) {}
+    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) {}
 
     /// \brief Return underlaying expression.
     Expr *getExpr() const { return E; }
diff --git a/test/Sema/complex_expressions.swift b/test/Sema/complex_expressions.swift
@@ -87,3 +87,12 @@ struct P {
 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>
