[TypeChecker] Re-work type-checking of string interpolations

include/swift/Sema/ConstraintSystem.h

@@ -466,6 +466,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;
@@ -3865,6 +3868,20 @@ class ConstraintSystem {
   bool resolveClosure(TypeVariableType *typeVar, Type contextualType,
                       ConstraintLocatorBuilder locator);
 
+  /// 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.
@@ -4259,6 +4276,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
+  [[nodiscard]]
+  bool generateConstraints(TapExpr *tap);
+
   /// Generate constraints for the body of the given function or closure.
   ///
   /// \param fn The function or closure expression
@@ -5174,6 +5199,21 @@ class ConstraintSystem {
       std::function<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

@@ -8451,15 +8451,11 @@ namespace {
   class ExprWalker : public ASTWalker {
     ExprRewriter &Rewriter;
     SmallVector<ClosureExpr *, 4> ClosuresToTypeCheck;
-    SmallVector<std::pair<TapExpr *, DeclContext *>, 4> TapsToTypeCheck;
 
   public:
     ExprWalker(ExprRewriter &Rewriter) : Rewriter(Rewriter) { }
 
-    ~ExprWalker() {
-      assert(ClosuresToTypeCheck.empty());
-      assert(TapsToTypeCheck.empty());
-    }
+    ~ExprWalker() { assert(ClosuresToTypeCheck.empty()); }
 
     bool shouldWalkIntoPropertyWrapperPlaceholderValue() override {
       // Property wrapper placeholder underlying values are filled in
@@ -8468,9 +8464,7 @@ namespace {
     }
 
     /// Check if there are any closures or tap expressions left to process separately.
-    bool hasDelayedTasks() {
-      return !ClosuresToTypeCheck.empty() || !TapsToTypeCheck.empty();
-    }
+    bool hasDelayedTasks() { return !ClosuresToTypeCheck.empty(); }
 
     /// Process delayed closure bodies and `Tap` expressions.
     ///
@@ -8513,17 +8507,6 @@ namespace {
         hadError |= TypeChecker::typeCheckClosureBody(closure);
       }
 
-      // Tap expressions too; they should or should not be
-      // type-checked under the same conditions as closure bodies.
-      {
-        for (const auto &tuple : TapsToTypeCheck) {
-          auto tap = std::get<0>(tuple);
-          auto tapDC = std::get<1>(tuple);
-          hadError |= TypeChecker::typeCheckTapBody(tap, tapDC);
-        }
-        TapsToTypeCheck.clear();
-      }
-
       return hadError;
     }
 
@@ -8548,10 +8531,9 @@ namespace {
         return Action::SkipChildren(SVE);
       }
 
-      if (auto tap = dyn_cast<TapExpr>(expr)) {
-        // We remember the DeclContext because the code to handle
-        // single-expression-body closures above changes it.
-        TapsToTypeCheck.push_back(std::make_pair(tap, Rewriter.dc));
+      if (auto tap = dyn_cast_or_null<TapExpr>(expr)) {
+        rewriteTapExpr(tap);
+        return Action::SkipChildren(tap);
       }
 
       if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
@@ -8694,6 +8676,26 @@ namespace {
             if (auto expr = resultTarget->getAsExpr())
               solution.setExprTypes(expr);
 
+            return resultTarget;
+         });
+    }
+
+    void rewriteTapExpr(TapExpr *tap) {
+      auto &solution = Rewriter.solution;
+
+      // First, let's visit the tap expression itself
+      // and set all of the inferred types.
+      Rewriter.visitTapExpr(tap);
+
+      // Now, let's apply solution to the body
+      (void)Rewriter.cs.applySolutionToBody(
+          solution, tap, Rewriter.dc, [&](SyntacticElementTarget target) {
+            auto resultTarget = rewriteTarget(target);
+            if (resultTarget) {
+              if (auto expr = resultTarget->getAsExpr())
+                solution.setExprTypes(expr);
+            }
+
             return resultTarget;
           });
     }

lib/Sema/CSGen.cpp

@@ -848,6 +848,65 @@ namespace {
 
 namespace {
 
+// Collect any variable references whose types involve type variables,
+// because there will be a dependency on those type variables once we have
+// generated constraints for the closure/tap body. This includes references
+// to other closure params such as in `{ x in { x }}` where the inner
+// closure is dependent on the outer closure's param type, as well as
+// cases like `for i in x where bar({ i })` where there's a dependency on
+// the type variable for the pattern `i`.
+struct VarRefCollector : public ASTWalker {
+  ConstraintSystem &cs;
+  llvm::SmallPtrSet<TypeVariableType *, 4> varRefs;
+
+  VarRefCollector(ConstraintSystem &cs) : cs(cs) {}
+
+  bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Arguments;
+  }
+
+  PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
+    // Retrieve type variables from references to var decls.
+    if (auto *declRef = dyn_cast<DeclRefExpr>(expr)) {
+      if (auto *varDecl = dyn_cast<VarDecl>(declRef->getDecl())) {
+        if (auto varType = cs.getTypeIfAvailable(varDecl)) {
+          varType->getTypeVariables(varRefs);
+        }
+      }
+    }
+
+    // FIXME: We can see UnresolvedDeclRefExprs here because we have
+    // not yet run preCheckExpression() on the entire closure body
+    // yet.
+    //
+    // We could consider pre-checking more eagerly.
+    if (auto *declRef = dyn_cast<UnresolvedDeclRefExpr>(expr)) {
+      auto name = declRef->getName();
+      auto loc = declRef->getLoc();
+      if (name.isSimpleName() && loc.isValid()) {
+        auto *varDecl =
+            dyn_cast_or_null<VarDecl>(ASTScope::lookupSingleLocalDecl(
+                cs.DC->getParentSourceFile(), name.getFullName(), loc));
+        if (varDecl)
+          if (auto varType = cs.getTypeIfAvailable(varDecl))
+            varType->getTypeVariables(varRefs);
+      }
+    }
+
+    return Action::Continue(expr);
+  }
+
+  PreWalkAction walkToDeclPre(Decl *D) override {
+    return Action::VisitChildrenIf(isa<PatternBindingDecl>(D));
+  }
+
+  PreWalkResult<Pattern *> walkToPatternPre(Pattern *P) override {
+    return Action::SkipChildren(P);
+  }
+};
+
   class ConstraintGenerator : public ExprVisitor<ConstraintGenerator, Type> {
     ConstraintSystem &CS;
     DeclContext *CurDC;
@@ -1199,15 +1258,38 @@ 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 (auto *tap = getAsExpr<TapExpr>(appendingExpr)) {
+          // Collect all of the variable references that appear
+          // in the tap body, otherwise tap expression is going
+          // to get disconnected from the context.
+          if (auto *body = tap->getBody()) {
+            VarRefCollector refCollector(CS);
+
+            body->walk(refCollector);
+
+            referencedVars.append(refCollector.varRefs.begin(),
+                                  refCollector.varRefs.end());
+          }
+        }
+
+        // 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;
@@ -2934,80 +3016,22 @@ namespace {
       auto *locator = CS.getConstraintLocator(closure);
       auto closureType = CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
 
-      // Collect any variable references whose types involve type variables,
-      // because there will be a dependency on those type variables once we have
-      // generated constraints for the closure body. This includes references
-      // to other closure params such as in `{ x in { x }}` where the inner
-      // closure is dependent on the outer closure's param type, as well as
-      // cases like `for i in x where bar({ i })` where there's a dependency on
-      // the type variable for the pattern `i`.
-      struct CollectVarRefs : public ASTWalker {
-        ConstraintSystem &cs;
-        llvm::SmallPtrSet<TypeVariableType *, 4> varRefs;
-
-        CollectVarRefs(ConstraintSystem &cs) : cs(cs) { }
-
-        bool shouldWalkCaptureInitializerExpressions() override { return true; }
-
-        MacroWalking getMacroWalkingBehavior() const override {
-          return MacroWalking::Arguments;
-        }
-
-        PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
-          // Retrieve type variables from references to var decls.
-          if (auto *declRef = dyn_cast<DeclRefExpr>(expr)) {
-            if (auto *varDecl = dyn_cast<VarDecl>(declRef->getDecl())) {
-              if (auto varType = cs.getTypeIfAvailable(varDecl)) {
-                varType->getTypeVariables(varRefs);
-              }
-            }
-          }
-
-          // FIXME: We can see UnresolvedDeclRefExprs here because we have
-          // not yet run preCheckExpression() on the entire closure body
-          // yet.
-          //
-          // We could consider pre-checking more eagerly.
-          if (auto *declRef = dyn_cast<UnresolvedDeclRefExpr>(expr)) {
-            auto name = declRef->getName();
-            auto loc = declRef->getLoc();
-            if (name.isSimpleName() && loc.isValid()) {
-              auto *varDecl = dyn_cast_or_null<VarDecl>(
-                ASTScope::lookupSingleLocalDecl(cs.DC->getParentSourceFile(),
-                                                name.getFullName(), loc));
-              if (varDecl)
-                if (auto varType = cs.getTypeIfAvailable(varDecl))
-                  varType->getTypeVariables(varRefs);
-            }
-          }
-
-          return Action::Continue(expr);
-        }
-
-        PreWalkAction walkToDeclPre(Decl *D) override {
-          return Action::VisitChildrenIf(isa<PatternBindingDecl>(D));
-        }
-
-        PreWalkResult<Pattern *> walkToPatternPre(Pattern *P) override {
-          return Action::SkipChildren(P);
-        }
-      } collectVarRefs(CS);
-
+      VarRefCollector refCollector(CS);
       // Walk the capture list if this closure has one,  because it could
       // reference declarations from the outer closure.
       if (auto *captureList =
               getAsExpr<CaptureListExpr>(CS.getParentExpr(closure))) {
-        captureList->walk(collectVarRefs);
+        captureList->walk(refCollector);
       } else {
-        closure->walk(collectVarRefs);
+        closure->walk(refCollector);
       }
 
       auto inferredType = inferClosureType(closure);
       if (!inferredType || inferredType->hasError())
         return Type();
 
       SmallVector<TypeVariableType *, 4> referencedVars{
-          collectVarRefs.varRefs.begin(), collectVarRefs.varRefs.end()};
+          refCollector.varRefs.begin(), refCollector.varRefs.end()};
 
       CS.addUnsolvedConstraint(Constraint::create(
           CS, ConstraintKind::DefaultClosureType, closureType, inferredType,

lib/Sema/CSSimplify.cpp

@@ -4379,6 +4379,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));
 
@@ -11222,6 +11228,23 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
   return !generateConstraints(AnyFunctionRef{closure}, closure->getBody());
 }
 
+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,