Allow SemaAnnotator to handle patterns

Add the necessary walking hooks, and fix
ReferenceCollector to use it.

Author: hamishknight

include/swift/IDE/SourceEntityWalker.h

@@ -73,6 +73,9 @@ class SourceEntityWalker {
   /// Walks the provided Expr.
   /// \returns true if traversal was aborted, false otherwise.
   bool walk(Expr *E);
+  /// Walks the provided Pattern.
+  /// \returns true if traversal was aborted, false otherwise.
+  bool walk(Pattern *P);
   /// Walks the provided ASTNode.
   /// \returns true if traversal was aborted, false otherwise.
   bool walk(ASTNode N);
@@ -101,6 +104,14 @@ class SourceEntityWalker {
   /// returns false, the remaining traversal is terminated and returns failure.
   virtual bool walkToExprPost(Expr *E) { return true; }
 
+  /// This method is called when first visiting a pattern, before walking
+  /// into its children. If it returns false, the subtree is skipped.
+  virtual bool walkToPatternPre(Pattern *P) { return true; }
+
+  /// This method is called after visiting the children of a pattern. If it
+  /// returns false, the remaining traversal is terminated and returns failure.
+  virtual bool walkToPatternPost(Pattern *P) { return true; }
+
   /// This method is called when a ValueDecl is referenced in source. If it
   /// returns false, the remaining traversal is terminated and returns failure.
   ///

lib/IDE/Refactoring.cpp

@@ -5474,10 +5474,10 @@ class ReferenceCollector : private SourceEntityWalker {
     return shouldWalkInto(S->getSourceRange());
   }
 
-  std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) {
+  bool walkToPatternPre(Pattern *P) override {
     if (P == Target.dyn_cast<Pattern *>())
       AfterTarget = true;
-    return { shouldWalkInto(P->getSourceRange()), P };
+    return shouldWalkInto(P->getSourceRange());
   }
 
   bool shouldWalkInto(SourceRange Range) {

lib/IDE/SourceEntityWalker.cpp

@@ -66,6 +66,7 @@ class SemaAnnotator : public ASTWalker {
   Stmt *walkToStmtPost(Stmt *S) override;
 
   std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override;
+  Pattern *walkToPatternPost(Pattern *P) override;
 
   bool handleImports(ImportDecl *Import);
   bool handleCustomAttributes(Decl *D);
@@ -587,6 +588,9 @@ std::pair<bool, Pattern *> SemaAnnotator::walkToPatternPre(Pattern *P) {
     return { false, nullptr };
   }
 
+  if (!SEWalker.walkToPatternPre(P))
+    return { false, P };
+
   if (P->isImplicit())
     return { true, P };
 
@@ -609,6 +613,16 @@ std::pair<bool, Pattern *> SemaAnnotator::walkToPatternPre(Pattern *P) {
   return { false, P };
 }
 
+Pattern *SemaAnnotator::walkToPatternPost(Pattern *P) {
+  if (isDone())
+     return nullptr;
+
+  bool Continue = SEWalker.walkToPatternPost(P);
+  if (!Continue)
+    Cancelled = true;
+  return Continue ? P : nullptr;
+}
+
 bool SemaAnnotator::handleCustomAttributes(Decl *D) {
   // CustomAttrs of non-param VarDecls are handled when this method is called
   // on their containing PatternBindingDecls (see below).
@@ -824,6 +838,11 @@ bool SourceEntityWalker::walk(Expr *E) {
   return performWalk(Annotator, [&]() { return E->walk(Annotator); });
 }
 
+bool SourceEntityWalker::walk(Pattern *P) {
+  SemaAnnotator Annotator(*this);
+  return performWalk(Annotator, [&]() { return P->walk(Annotator); });
+}
+
 bool SourceEntityWalker::walk(Decl *D) {
   SemaAnnotator Annotator(*this);
   return performWalk(Annotator, [&]() { return D->walk(Annotator); });