[ConstraintSystem] Implement tap expression checking in the solver

Generate a conjunction for each tap expression body as soon as it
gets a contextual type instead of separate post-factum type-checking
via `typeCheckTapBody`.

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -472,6 +472,9 @@ class TypeVariableType::Implementation {
   /// Determine whether this type variable represents a closure type.
   bool isClosureType() const;
 
+  /// Determine whether this type variable represents a type of tap expression.
+  bool isTapType() const;
+
   /// Determine whether this type variable represents one of the
   /// parameter types associated with a closure.
   bool isClosureParameterType() const;
@@ -3951,6 +3954,20 @@ class ConstraintSystem {
   /// \returns `true` if pack expansion has been resolved, `false` otherwise.
   bool resolvePackExpansion(TypeVariableType *typeVar, Type contextualType);
 
+  /// Bind tap expression to the given contextual type and generate
+  /// constraints for its body.
+  ///
+  /// \param typeVar The type variable representing the tap expression.
+  /// \param contextualType The contextual type this tap expression
+  /// would be bound to.
+  /// \param locator The locator associated with contextual type.
+  ///
+  /// \returns `true` if it was possible to generate constraints for
+  /// the body and assign fixed type to the tap expression, `false`
+  /// otherwise.
+  bool resolveTapBody(TypeVariableType *typeVar, Type contextualType,
+                      ConstraintLocatorBuilder locator);
+
   /// Assign a fixed type to the given type variable.
   ///
   /// \param typeVar The type variable to bind.
@@ -4354,6 +4371,14 @@ class ConstraintSystem {
                       FreeTypeVariableBinding allowFreeTypeVariables =
                           FreeTypeVariableBinding::Disallow);
 
+  /// Generate constraints for the body of the given tap expression.
+  ///
+  /// \param tap the tap expression
+  ///
+  /// \returns \c true if constraint generation failed, \c false otherwise
+  LLVM_NODISCARD
+  bool generateConstraints(TapExpr *tap);
+
   /// Generate constraints for the body of the given function or closure.
   ///
   /// \param fn The function or closure expression
@@ -5285,6 +5310,21 @@ class ConstraintSystem {
           llvm::Optional<SyntacticElementTarget>(SyntacticElementTarget)>
           rewriteTarget);
 
+  /// Apply the given solution to the given tap expression.
+  ///
+  /// \param solution The solution to apply.
+  /// \param tapExpr The tap expression to which the solution is being applied.
+  /// \param currentDC The declaration context in which transformations
+  /// will be applied.
+  /// \param rewriteTarget Function that performs a rewrite of any
+  /// solution application target within the context.
+  ///
+  /// \returns true if solution cannot be applied.
+  bool applySolutionToBody(
+      Solution &solution, TapExpr *tapExpr, DeclContext *&currentDC,
+      std::function<Optional<SyntacticElementTarget>(SyntacticElementTarget)>
+          rewriteTarget);
+
   /// Reorder the disjunctive clauses for a given expression to
   /// increase the likelihood that a favored constraint will be successfully
   /// resolved before any others.

lib/Sema/CSApply.cpp

@@ -8638,7 +8638,17 @@ namespace {
         for (const auto &tuple : TapsToTypeCheck) {
           auto tap = std::get<0>(tuple);
           auto tapDC = std::get<1>(tuple);
-          hadError |= TypeChecker::typeCheckTapBody(tap, tapDC);
+
+          hadError |= cs.applySolutionToBody(
+              solution, tap, tapDC, [&](SolutionApplicationTarget target) {
+                auto resultTarget = rewriteTarget(target);
+                if (resultTarget) {
+                  if (auto expr = resultTarget->getAsExpr())
+                    solution.setExprTypes(expr);
+                }
+
+                return resultTarget;
+              });
         }
         TapsToTypeCheck.clear();
       }

lib/Sema/CSGen.cpp

@@ -1236,15 +1236,34 @@ namespace {
           return nullptr;
         }
 
-        auto interpolationTV = DependentMemberType::get(tv, associatedTypeDecl);
+        auto interpolationTV =
+            CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
+        auto interpolationType =
+            DependentMemberType::get(tv, associatedTypeDecl);
+
+        CS.addConstraint(ConstraintKind::Equal, interpolationTV,
+                         interpolationType, locator);
 
         auto appendingExprType = CS.getType(appendingExpr);
         auto appendingLocator = CS.getConstraintLocator(appendingExpr);
 
-        // Must be Conversion; if it's Equal, then in semi-rare cases, the 
+        SmallVector<TypeVariableType *, 2> referencedVars;
+
+        // If tap expression is located in a closure, let's connect them
+        // because interpolation could use parameters.
+        if (auto *tap = getAsExpr<TapExpr>(appendingExpr)) {
+          auto *tapDC = tap->getVar()->getDeclContext();
+          if (auto *closure = dyn_cast<ClosureExpr>(tapDC)) {
+            referencedVars.push_back(
+                CS.getType(closure)->castTo<TypeVariableType>());
+          }
+        }
+
+        // Must be Conversion; if it's Equal, then in semi-rare cases, the
         // interpolation temporary variable cannot be @lvalue.
-        CS.addConstraint(ConstraintKind::Conversion, appendingExprType,
-                         interpolationTV, appendingLocator);
+        CS.addUnsolvedConstraint(Constraint::create(
+            CS, ConstraintKind::Conversion, appendingExprType, interpolationTV,
+            appendingLocator, referencedVars));
       }
 
       return tv;

lib/Sema/CSSimplify.cpp

@@ -4413,6 +4413,12 @@ ConstraintSystem::matchTypesBindTypeVar(
                : getTypeMatchFailure(locator);
   }
 
+  if (typeVar->getImpl().isTapType()) {
+    return resolveTapBody(typeVar, type, locator)
+               ? getTypeMatchSuccess()
+               : getTypeMatchFailure(locator);
+  }
+
   assignFixedType(typeVar, type, /*updateState=*/true,
                   /*notifyInference=*/!flags.contains(TMF_BindingTypeVariable));
 
@@ -11508,6 +11514,23 @@ bool ConstraintSystem::resolvePackExpansion(TypeVariableType *typeVar,
   return true;
 }
 
+bool ConstraintSystem::resolveTapBody(TypeVariableType *typeVar,
+                                      Type contextualType,
+                                      ConstraintLocatorBuilder locator) {
+  auto *tapLoc = typeVar->getImpl().getLocator();
+  auto *tapExpr = castToExpr<TapExpr>(tapLoc->getAnchor());
+
+  // Assign a type to tap expression itself.
+  assignFixedType(typeVar, contextualType, getConstraintLocator(locator));
+  // Set type to `$interpolation` variable declared in the body of tap
+  // expression.
+  setType(tapExpr->getVar(), contextualType);
+
+  // With all of the contextual information recorded in the constraint system,
+  // it's time to generate constraints for the body of this tap expression.
+  return !generateConstraints(tapExpr);
+}
+
 ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyDynamicTypeOfConstraint(
                                         Type type1, Type type2,

lib/Sema/CSSyntacticElement.cpp

@@ -359,6 +359,17 @@ static void createConjunction(ConstraintSystem &cs,
     isIsolated = true;
   }
 
+  if (locator->directlyAt<TapExpr>()) {
+    auto *tap = castToExpr<TapExpr>(locator->getAnchor());
+    auto *tapDC = tap->getVar()->getDeclContext();
+
+    referencedVars.push_back(cs.getType(tap)->castTo<TypeVariableType>());
+
+    if (auto *closure = dyn_cast<ClosureExpr>(tapDC)) {
+      referencedVars.push_back(cs.getType(closure)->castTo<TypeVariableType>());
+    }
+  }
+
   UnresolvedVarCollector paramCollector(cs);
 
   for (const auto &entry : elements) {
@@ -1375,6 +1386,22 @@ class SyntacticElementConstraintGenerator
 };
 }
 
+bool ConstraintSystem::generateConstraints(TapExpr *tap) {
+  SyntacticElementConstraintGenerator generator(
+      *this, SyntacticElementContext::forTapExpr(tap),
+      getConstraintLocator(tap));
+
+  auto *body = tap->getBody();
+
+  if (!body) {
+    assert(tap->getSubExpr());
+    return false;
+  }
+
+  generator.visit(tap->getBody());
+  return generator.hadError;
+}
+
 bool ConstraintSystem::generateConstraints(AnyFunctionRef fn, BraceStmt *body) {
   NullablePtr<ConstraintLocator> locator;
 
@@ -2538,8 +2565,23 @@ bool ConstraintSystem::applySolutionToBody(Solution &solution,
   return false;
 }
 
+bool ConstraintSystem::applySolutionToBody(Solution &solution, TapExpr *tapExpr,
+                                           DeclContext *&currentDC,
+                                           RewriteTargetFn rewriteTarget) {
+  SyntacticElementSolutionApplication application(
+      solution, SyntacticElementContext::forTapExpr(tapExpr), rewriteTarget);
+
+  auto body = application.apply();
+
+  if (!body || application.hadError)
+    return true;
+
+  tapExpr->setBody(castToStmt<BraceStmt>(body));
+  return false;
+}
+
 bool ConjunctionElement::mightContainCodeCompletionToken(
-    const ConstraintSystem &cs) const {
+  const ConstraintSystem &cs) const {
   if (Element->getKind() == ConstraintKind::SyntacticElement) {
     if (Element->getSyntacticElement().getSourceRange().isInvalid()) {
       return true;

lib/Sema/TypeCheckConstraints.cpp

@@ -120,6 +120,10 @@ bool TypeVariableType::Implementation::isClosureType() const {
   return isExpr<ClosureExpr>(locator->getAnchor()) && locator->getPath().empty();
 }
 
+bool TypeVariableType::Implementation::isTapType() const {
+  return locator && locator->directlyAt<TapExpr>();
+}
+
 bool TypeVariableType::Implementation::isClosureParameterType() const {
   if (!(locator && locator->getAnchor()))
     return false;