[CSGen] NFC: Extract var reference collector to be used in more places

Author: xedin

lib/Sema/CSGen.cpp

@@ -848,6 +848,61 @@ 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));
+  }
+};
+
   class ConstraintGenerator : public ExprVisitor<ConstraintGenerator, Type> {
     ConstraintSystem &CS;
     DeclContext *CurDC;
@@ -2902,76 +2957,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));
-        }
-      } 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::FallbackType, closureType,