[IDE] Allow SourceEntityWalker to keep track of its ASTWalker

This will allow it to be queried for things like
the parent ASTNode in the current traversal.

Author: hamishknight

include/swift/IDE/SourceEntityWalker.h

@@ -14,6 +14,7 @@
 #define SWIFT_IDE_SOURCE_ENTITY_WALKER_H
 
 #include "swift/AST/ASTWalker.h"
+#include "swift/Basic/Defer.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/SourceLoc.h"
 #include "llvm/ADT/PointerUnion.h"
@@ -176,6 +177,24 @@ class SourceEntityWalker {
   virtual ~SourceEntityWalker() {}
 
   virtual void anchor();
+
+  /// Retrieve the current ASTWalker being used to traverse the AST.
+  const ASTWalker &getWalker() const {
+    assert(Walker && "Not walking!");
+    return *Walker;
+  }
+
+private:
+  ASTWalker *Walker = nullptr;
+
+  /// Utility that lets us keep track of an ASTWalker when walking.
+  bool performWalk(ASTWalker &W, llvm::function_ref<bool(void)> DoWalk) {
+    Walker = &W;
+    SWIFT_DEFER {
+      Walker = nullptr;
+    };
+    return DoWalk();
+  }
 };
 
 } // namespace swift

lib/IDE/SourceEntityWalker.cpp

@@ -806,32 +806,32 @@ bool SemaAnnotator::shouldIgnore(Decl *D) {
 
 bool SourceEntityWalker::walk(SourceFile &SrcFile) {
   SemaAnnotator Annotator(*this);
-  return SrcFile.walk(Annotator);
+  return performWalk(Annotator, [&]() { return SrcFile.walk(Annotator); });
 }
 
 bool SourceEntityWalker::walk(ModuleDecl &Mod) {
   SemaAnnotator Annotator(*this);
-  return Mod.walk(Annotator);
+  return performWalk(Annotator, [&]() { return Mod.walk(Annotator); });
 }
 
 bool SourceEntityWalker::walk(Stmt *S) {
   SemaAnnotator Annotator(*this);
-  return S->walk(Annotator);
+  return performWalk(Annotator, [&]() { return S->walk(Annotator); });
 }
 
 bool SourceEntityWalker::walk(Expr *E) {
   SemaAnnotator Annotator(*this);
-  return E->walk(Annotator);
+  return performWalk(Annotator, [&]() { return E->walk(Annotator); });
 }
 
 bool SourceEntityWalker::walk(Decl *D) {
   SemaAnnotator Annotator(*this);
-  return D->walk(Annotator);
+  return performWalk(Annotator, [&]() { return D->walk(Annotator); });
 }
 
 bool SourceEntityWalker::walk(DeclContext *DC) {
   SemaAnnotator Annotator(*this);
-  return DC->walkContext(Annotator);
+  return performWalk(Annotator, [&]() { return DC->walkContext(Annotator); });
 }
 
 bool SourceEntityWalker::walk(ASTNode N) {