[Type Checker] Extend Path Consistency algorithm to cover collections

xedin · rudkx · commit 07424a2c5b41 · 2016-09-15T06:42:35.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. (cherry picked from commit 73318a2)
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -1358,13 +1358,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;
@@ -1378,7 +1371,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,
@@ -1492,52 +1485,22 @@ 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, CS.getContextualType(expr),
+                            CS.getContextualTypePurpose());
         // 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)) {
@@ -1555,12 +1518,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;
         }
       }
@@ -1590,10 +1565,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
@@ -1024,15 +1024,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,13 @@ 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>
+let v7: [UInt32] = [55 * 8, 0]
