[CS] Merge VarRefCollector & UnresolvedVarCollector

These now do basically the same thing, merge their
implementations.

Author: hamishknight

include/swift/Sema/ConstraintSystem.h

@@ -6193,6 +6193,39 @@ class ConjunctionElementProducer : public BindingProducer<ConjunctionElement> {
   }
 };
 
+/// Find any references to not yet resolved outer VarDecls (including closure
+/// parameters) used in the body of a conjunction element (e.g closures, taps,
+/// if/switch expressions). This is required because isolated conjunctions, just
+/// like single-expression closures, have to be connected to type variables they
+/// are going to use, otherwise they'll get placed in a separate solver
+/// component and would never produce a solution.
+class VarRefCollector : public ASTWalker {
+  ConstraintSystem &CS;
+  llvm::SmallSetVector<TypeVariableType *, 4> TypeVars;
+
+public:
+  VarRefCollector(ConstraintSystem &cs) : CS(cs) {}
+
+  /// Infer the referenced type variables from a given decl.
+  void inferTypeVars(Decl *D);
+
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Arguments;
+  }
+
+  PreWalkResult<Expr *> walkToExprPre(Expr *expr) override;
+
+  PreWalkAction walkToDeclPre(Decl *D) override {
+    // We only need to walk into PatternBindingDecls, other kinds of decls
+    // cannot reference outer vars.
+    return Action::VisitChildrenIf(isa<PatternBindingDecl>(D));
+  }
+
+  ArrayRef<TypeVariableType *> getTypeVars() const {
+    return TypeVars.getArrayRef();
+  }
+};
+
 /// Determine whether given type is a known one
 /// for a key path `{Writable, ReferenceWritable}KeyPath`.
 bool isKnownKeyPathType(Type type);

lib/Sema/CSGen.cpp

@@ -846,63 +846,42 @@ namespace {
   };
 } // end anonymous 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; }
+void VarRefCollector::inferTypeVars(Decl *D) {
+  // We're only interested in VarDecls.
+  if (!isa_and_nonnull<VarDecl>(D))
+    return;
 
-  MacroWalking getMacroWalkingBehavior() const override {
-    return MacroWalking::Arguments;
-  }
+  auto ty = CS.getTypeIfAvailable(D);
+  if (!ty)
+    return;
 
-  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);
-        }
-      }
-    }
+  SmallPtrSet<TypeVariableType *, 4> typeVars;
+  ty->getTypeVariables(typeVars);
+  TypeVars.insert(typeVars.begin(), typeVars.end());
+}
 
-    // 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);
-      }
+ASTWalker::PreWalkResult<Expr *>
+VarRefCollector::walkToExprPre(Expr *expr) {
+  if (auto *DRE = dyn_cast<DeclRefExpr>(expr))
+    inferTypeVars(DRE->getDecl());
+
+  // FIXME: We can see UnresolvedDeclRefExprs here because we don't walk into
+  // patterns when running preCheckExpression, since we don't resolve patterns
+  // until CSGen. We ought to consider moving pattern resolution into
+  // pre-checking, which would allow us to pre-check patterns normally.
+  if (auto *declRef = dyn_cast<UnresolvedDeclRefExpr>(expr)) {
+    auto name = declRef->getName();
+    auto loc = declRef->getLoc();
+    if (name.isSimpleName() && loc.isValid()) {
+      auto *SF = CS.DC->getParentSourceFile();
+      auto *D = ASTScope::lookupSingleLocalDecl(SF, name.getFullName(), loc);
+      inferTypeVars(D);
     }
-
-    return Action::Continue(expr);
-  }
-
-  PreWalkAction walkToDeclPre(Decl *D) override {
-    return Action::VisitChildrenIf(isa<PatternBindingDecl>(D));
   }
-};
+  return Action::Continue(expr);
+}
 
+namespace {
   class ConstraintGenerator : public ExprVisitor<ConstraintGenerator, Type> {
     ConstraintSystem &CS;
     DeclContext *CurDC;
@@ -1309,8 +1288,8 @@ struct VarRefCollector : public ASTWalker {
 
             body->walk(refCollector);
 
-            referencedVars.append(refCollector.varRefs.begin(),
-                                  refCollector.varRefs.end());
+            auto vars = refCollector.getTypeVars();
+            referencedVars.append(vars.begin(), vars.end());
           }
         }
 
@@ -2971,9 +2950,7 @@ struct VarRefCollector : public ASTWalker {
       if (!inferredType || inferredType->hasError())
         return Type();
 
-      SmallVector<TypeVariableType *, 4> referencedVars{
-          refCollector.varRefs.begin(), refCollector.varRefs.end()};
-
+      auto referencedVars = refCollector.getTypeVars();
       CS.addUnsolvedConstraint(
           Constraint::create(CS, ConstraintKind::FallbackType, closureType,
                              inferredType, locator, referencedVars));

lib/Sema/CSSyntacticElement.cpp

@@ -239,71 +239,6 @@ class TypeVariableRefFinder : public ASTWalker {
   }
 };
 
-/// Find any references to not yet resolved outer VarDecls (including closure
-/// parameters) used in the body of the inner closure. This is required because
-/// isolated conjunctions, just like single-expression closures, have
-/// to be connected to type variables they are going to use, otherwise
-/// they'll get placed in a separate solver component and would never
-/// produce a solution.
-class UnresolvedVarCollector : public ASTWalker {
-  ConstraintSystem &CS;
-
-  llvm::SmallSetVector<TypeVariableType *, 4> Vars;
-
-public:
-  UnresolvedVarCollector(ConstraintSystem &cs) : CS(cs) {}
-
-  MacroWalking getMacroWalkingBehavior() const override {
-    return MacroWalking::Arguments;
-  }
-
-  PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
-    if (auto *DRE = dyn_cast<DeclRefExpr>(expr)) {
-      auto *decl = DRE->getDecl();
-      if (isa<VarDecl>(decl)) {
-        if (auto type = CS.getTypeIfAvailable(decl)) {
-          if (auto *typeVar = type->getAs<TypeVariableType>()) {
-            Vars.insert(typeVar);
-          } else if (type->hasTypeVariable()) {
-            // Parameter or result type could be only partially
-            // resolved e.g. `{ (x: X) -> Void in ... }` where
-            // `X` is a generic type.
-            SmallPtrSet<TypeVariableType *, 4> typeVars;
-            type->getTypeVariables(typeVars);
-            Vars.insert(typeVars.begin(), typeVars.end());
-          }
-        }
-      }
-    }
-
-    // FIXME: We can see UnresolvedDeclRefExprs here because we don't walk into
-    // patterns when running preCheckExpression, since we don't resolve patterns
-    // until CSGen. We ought to consider moving pattern resolution into
-    // pre-checking, which would allow us to pre-check patterns normally.
-    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)) {
-            SmallPtrSet<TypeVariableType *, 4> typeVars;
-            varType->getTypeVariables(typeVars);
-            Vars.insert(typeVars.begin(), typeVars.end());
-          }
-        }
-      }
-    }
-    return Action::Continue(expr);
-  }
-
-  ArrayRef<TypeVariableType *> getVariables() const {
-    return Vars.getArrayRef();
-  }
-};
-
 // MARK: Constraint generation
 
 /// Check whether it makes sense to convert this element into a constraint.
@@ -387,7 +322,7 @@ static void createConjunction(ConstraintSystem &cs,
     isIsolated = true;
   }
 
-  UnresolvedVarCollector paramCollector(cs);
+  VarRefCollector paramCollector(cs);
 
   for (const auto &entry : elements) {
     ASTNode element = std::get<0>(entry);
@@ -415,7 +350,7 @@ static void createConjunction(ConstraintSystem &cs,
   if (constraints.empty())
     return;
 
-  for (auto *externalVar : paramCollector.getVariables())
+  for (auto *externalVar : paramCollector.getTypeVars())
     referencedVars.push_back(externalVar);
 
   cs.addUnsolvedConstraint(Constraint::createConjunction(