[Constraint solver] Compute common apply result type in the solver.

DougGregor · DougGregor · commit 5f99d91ea85f · 2019-02-27T23:30:03.000-08:00
Constraint generation for function application expressions contains a simple
hack to try to find the common result type for an overload set containing
callable things. Instead, perform this “common result type” computation
when simplifying an applicable function constraint, so it is more
widely applicable.
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -2516,8 +2516,6 @@ namespace {
     }
 
     Type visitApplyExpr(ApplyExpr *expr) {
-      Type outputTy;
-
       auto fnExpr = expr->getFn();
 
       if (auto *UDE = dyn_cast<UnresolvedDotExpr>(fnExpr)) {
@@ -2526,65 +2524,9 @@ namespace {
           return resultOfTypeOperation(typeOperation, expr->getArg());
       }
 
-      if (auto fnType = CS.getType(fnExpr)->getAs<AnyFunctionType>()) {
-        outputTy = fnType->getResult();
-      } else if (auto OSR = dyn_cast<OverloadedDeclRefExpr>(fnExpr)) {
-        // Determine if the overloads are all functions that share a common
-        // return type.
-        Type commonType;
-        for (auto OD : OSR->getDecls()) {
-          auto OFD = dyn_cast<AbstractFunctionDecl>(OD);
-          if (!OFD) {
-            commonType = Type();
-            break;
-          }
-
-          auto OFT = OFD->getInterfaceType()->getAs<AnyFunctionType>();
-          if (!OFT) {
-            commonType = Type();
-            break;
-          }
-
-          // Look past the self parameter.
-          if (OFD->getDeclContext()->isTypeContext()) {
-            OFT = OFT->getResult()->getAs<AnyFunctionType>();
-            if (!OFT) {
-              commonType = Type();
-              break;
-            }
-          }
-
-          Type resultType = OFT->getResult();
-
-          // If there are any type parameters in the result,
-          if (resultType->hasTypeParameter()) {
-            commonType = Type();
-            break;
-          }
-
-          if (commonType.isNull()) {
-            commonType = resultType;
-          } else if (!commonType->isEqual(resultType)) {
-            commonType = Type();
-            break;
-          }
-        }
-
-        if (commonType) {
-          outputTy = commonType;
-        }
-      }
-      
-      // The function subexpression has some rvalue type T1 -> T2 for fresh
-      // variables T1 and T2.
-      if (outputTy.isNull()) {
-        outputTy = CS.createTypeVariable(
-            CS.getConstraintLocator(expr, ConstraintLocator::FunctionResult));
-      } else {
-        // Since we know what the output type is, we can set it as the favored
-        // type of this expression.
-        CS.setFavoredType(expr, outputTy.getPointer());
-      }
+      // The result type is a fresh type variable.
+      Type resultType = CS.createTypeVariable(
+          CS.getConstraintLocator(expr, ConstraintLocator::FunctionResult));
 
       // A direct call to a ClosureExpr makes it noescape.
       FunctionType::ExtInfo extInfo;
@@ -2598,11 +2540,20 @@ namespace {
       AnyFunctionType::decomposeInput(CS.getType(expr->getArg()), params);
 
       CS.addConstraint(ConstraintKind::ApplicableFunction,
-                       FunctionType::get(params, outputTy, extInfo),
+                       FunctionType::get(params, resultType, extInfo),
                        CS.getType(expr->getFn()),
         CS.getConstraintLocator(expr, ConstraintLocator::ApplyFunction));
 
-      return outputTy;
+      // If we ended up resolving the result type variable to a concrete type,
+      // set it as the favored type for this expression.
+      Type fixedType =
+          CS.getFixedTypeRecursive(resultType, /*wantRvalue=*/true);
+      if (!fixedType->isTypeVariableOrMember()) {
+        CS.setFavoredType(expr, fixedType.getPointer());
+        resultType = fixedType;
+      }
+
+      return resultType;
     }
 
     Type getSuperType(VarDecl *selfDecl,
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -4777,7 +4777,7 @@ ConstraintSystem::simplifyApplicableFnConstraint(
 
   // By construction, the left hand side is a type that looks like the
   // following: $T1 -> $T2.
-  assert(type1->is<FunctionType>());
+  auto func1 = type1->castTo<FunctionType>();
 
   // Let's check if this member couldn't be found and is fixed
   // to exist based on its usage.
@@ -4821,6 +4821,16 @@ ConstraintSystem::simplifyApplicableFnConstraint(
 
   };
 
+  // If the right-hand side is a type variable, try to find a common result
+  // type in the overload set.
+  if (auto typeVar = desugar2->getAs<TypeVariableType>()) {
+    auto choices = getUnboundBindOverloads(typeVar);
+    if (Type resultType = findCommonResultType(choices)) {
+      addConstraint(ConstraintKind::Bind, func1->getResult(), resultType,
+                    locator);
+    }
+  }
+
   // If right-hand side is a type variable, the constraint is unsolved.
   if (desugar2->isTypeVariableOrMember())
     return formUnsolved();
@@ -4856,7 +4866,6 @@ ConstraintSystem::simplifyApplicableFnConstraint(
   }
 
   // For a function, bind the output and convert the argument to the input.
-  auto func1 = type1->castTo<FunctionType>();
   if (auto func2 = dyn_cast<FunctionType>(desugar2)) {
     // The argument type must be convertible to the input type.
     if (::matchCallArguments(
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -1567,19 +1567,13 @@ bool ConstraintSystem::haveTypeInformationForAllArguments(
                       });
 }
 
-// Given a type variable representing the RHS of an ApplicableFunction
-// constraint, attempt to find the disjunction of bind overloads
-// associated with it. This may return null in cases where have not
-// yet created a disjunction because we need to resolve a base type,
-// e.g.: [1].map{ ... } does not have a disjunction until we decide on
-// a type for [1].
-static Constraint *getUnboundBindOverloadDisjunction(TypeVariableType *tyvar,
-                                                     ConstraintSystem &cs) {
-  auto *rep = cs.getRepresentative(tyvar);
-  assert(!cs.getFixedType(rep));
+Constraint *ConstraintSystem::getUnboundBindOverloadDisjunction(
+                                                  TypeVariableType *tyvar) {
+  auto *rep = getRepresentative(tyvar);
+  assert(!getFixedType(rep));
 
   llvm::SetVector<Constraint *> disjunctions;
-  cs.getConstraintGraph().gatherConstraints(
+  getConstraintGraph().gatherConstraints(
       rep, disjunctions, ConstraintGraph::GatheringKind::EquivalenceClass,
       [](Constraint *match) {
         return match->getKind() == ConstraintKind::Disjunction &&
@@ -1593,6 +1587,38 @@ static Constraint *getUnboundBindOverloadDisjunction(TypeVariableType *tyvar,
   return disjunctions[0];
 }
 
+/// solely resolved by an overload set.
+SmallVector<OverloadChoice, 2> ConstraintSystem::getUnboundBindOverloads(
+                                                     TypeVariableType *tyvar) {
+  // Always work on the representation.
+  tyvar = getRepresentative(tyvar);
+
+  SmallVector<OverloadChoice, 2> choices;
+
+  auto disjunction = getUnboundBindOverloadDisjunction(tyvar);
+  if (!disjunction) return choices;
+
+  for (auto constraint : disjunction->getNestedConstraints()) {
+    // We must have bind-overload constraints.
+    if (constraint->getKind() != ConstraintKind::BindOverload) {
+      choices.clear();
+      return choices;
+    }
+
+    // We must be binding the type variable (or a type variable equivalent to
+    // it).
+    auto boundTypeVar = constraint->getFirstType()->getAs<TypeVariableType>();
+    if (!boundTypeVar || getRepresentative(boundTypeVar) != tyvar) {
+      choices.clear();
+      return choices;
+    }
+
+    choices.push_back(constraint->getOverloadChoice());
+  }
+
+  return choices;
+}
+
 // Find a disjunction associated with an ApplicableFunction constraint
 // where we have some information about all of the types of in the
 // function application (even if we only know something about what the
@@ -1608,7 +1634,7 @@ Constraint *ConstraintSystem::selectApplyDisjunction() {
       auto *tyvar = applicable->getSecondType()->castTo<TypeVariableType>();
 
       // If we have created the disjunction for this apply, find it.
-      auto *disjunction = getUnboundBindOverloadDisjunction(tyvar, *this);
+      auto *disjunction = getUnboundBindOverloadDisjunction(tyvar);
       if (disjunction)
         return disjunction;
     }
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -1466,7 +1466,8 @@ static ArrayRef<OverloadChoice> partitionSIMDOperators(
 }
 
 /// Retrieve the type that will be used when matching the given overload.
-static Type getEffectiveOverloadType(const OverloadChoice &overload) {
+static Type getEffectiveOverloadType(const OverloadChoice &overload,
+                                     bool allowMembers) {
   switch (overload.getKind()) {
   case OverloadChoiceKind::Decl:
     // Declaration choices are handled below.
@@ -1506,9 +1507,19 @@ static Type getEffectiveOverloadType(const OverloadChoice &overload) {
                              genericFn->getExtInfo());
   }
 
-  // If this declaration is within a type context, bail out.
+  // If this declaration is within a type context, we might not be able
+  // to handle it.
   if (decl->getDeclContext()->isTypeContext()) {
-    return Type();
+    if (!allowMembers)
+      return Type();
+
+    // FIXME: This is overly restrictive for now.
+    if ((!isa<FuncDecl>(decl) && !isa<EnumElementDecl>(decl)) ||
+        (decl->isInstanceMember() &&
+         (!overload.getBaseType() || !overload.getBaseType()->getAnyNominal())))
+      return Type();
+
+    type = type->castTo<FunctionType>()->getResult();
   }
 
   return type;
@@ -1756,6 +1767,54 @@ class CommonTypeVisitor : public TypeVisitor<CommonTypeVisitor, Type, Type> {
 
 }
 
+Type ConstraintSystem::findCommonResultType(ArrayRef<OverloadChoice> choices) {
+  // Local function to consider this s new overload choice, updating the
+  // "common type". Returns true if this overload cannot be integrated into
+  // the common type, at which point there is no "common type".
+  Type commonType;
+  auto considerOverload = [&](const OverloadChoice &overload) -> bool {
+    // If we can't even get a type for the overload, there's nothing more to
+    // do.
+    Type overloadType =
+        getEffectiveOverloadType(overload, /*allowMembers=*/true);
+    if (!overloadType) {
+      return true;
+    }
+
+    auto functionType = overloadType->getAs<FunctionType>();
+    if (!functionType) {
+      return true;
+    }
+
+    auto resultType = functionType->getResult();
+    if (resultType->hasTypeParameter()) {
+      return true;
+    }
+
+    // If this is the first overload, record it's type as the common type.
+    if (!commonType) {
+      commonType = resultType;
+      return false;
+    }
+
+    // Find the common type between the current common type and the new
+    // overload's type.
+    commonType = CommonTypeVisitor().visit(commonType, resultType);
+    if (!commonType) {
+      return true;
+    }
+
+    return false;
+  };
+
+  for (const auto &choice : choices) {
+    if (considerOverload(choice))
+      return Type();
+  }
+
+  return commonType;
+}
+
 Type ConstraintSystem::findCommonOverloadType(
     ArrayRef<OverloadChoice> choices,
     ArrayRef<OverloadChoice> outerAlternatives,
@@ -1767,7 +1826,8 @@ Type ConstraintSystem::findCommonOverloadType(
   auto considerOverload = [&](const OverloadChoice &overload) -> bool {
     // If we can't even get a type for the overload, there's nothing more to
     // do.
-    Type overloadType = getEffectiveOverloadType(overload);
+    Type overloadType =
+      getEffectiveOverloadType(overload, /*allowMembers=*/false);
     if (!overloadType) {
       return true;
     }
diff --git a/lib/Sema/ConstraintSystem.h b/lib/Sema/ConstraintSystem.h
@@ -2326,6 +2326,12 @@ class ConstraintSystem {
                           const DeclRefExpr *base = nullptr,
                           OpenedTypeMap *replacements = nullptr);
 
+  /// Given a set of overload choices, try to find a common result type when
+  /// they are called.
+  ///
+  /// \returns the common type amongst the set of overload choices.
+  Type findCommonResultType(ArrayRef<OverloadChoice> choices);
+
   /// Given a set of overload choices, try to find a common structure amongst
   /// all of them.
   ///
@@ -3058,6 +3064,20 @@ class ConstraintSystem {
   /// Collect the current inactive disjunction constraints.
   void collectDisjunctions(SmallVectorImpl<Constraint *> &disjunctions);
 
+  // Given a type variable, attempt to find the disjunction of
+  // bind overloads associated with it. This may return null in cases where
+  // the disjunction has either not been created or binds the type variable
+  // in some manner other than by binding overloads.
+  Constraint *getUnboundBindOverloadDisjunction(TypeVariableType *tyvar);
+
+  /// Given a type variable that might represent an overload set, retrieve
+  ///
+  /// \returns the set of overload choices to which this type variable
+  /// could be bound, or an empty vector if the type variable is not
+  /// solely resolved by an overload set.
+  SmallVector<OverloadChoice, 2> getUnboundBindOverloads(
+                                                  TypeVariableType *tyvar);
+
   /// Solve the system of constraints after it has already been
   /// simplified.
   ///
