[ASTWalker] Re-introduce Action::SkipChildren

Now with the behavior of resuming at the post-walk.

Author: hamishknight

include/swift/AST/ASTWalker.h

@@ -163,7 +163,7 @@ class ASTWalker {
 
     // The 'Action' set of types, which do not take a payload.
     struct ContinueWalkAction {};
-    struct SkipChildrenIfWalkAction { bool Cond; };
+    struct SkipChildrenIfWalkAction { bool Cond; bool SkipPostWalk; };
     struct StopIfWalkAction { bool Cond; };
     struct StopWalkAction {};
 
@@ -216,13 +216,30 @@ class ASTWalker {
       return {Continue(), std::move(node)};
     }
 
-    /// Skips visiting both the node's children and its post-visitation.
+    /// Continue the current walk, replacing the current node with \p node.
+    /// However, skip visiting the children of \p node, and resume at its
+    /// post-walk.
+    template <typename T>
+    static _Detail::SkipChildrenIfWalkResult<T> SkipChildren(T node) {
+      return SkipChildrenIf(true, std::move(node));
+    }
+
+    /// Similar to \c Action::SkipChildren, but also skips the call to the
+    /// post-visitation method.
     template <typename T>
     static _Detail::SkipChildrenIfWalkResult<T>
     SkipNode(T node) {
       return SkipNodeIf(true, std::move(node));
     }
 
+    /// If \p cond is true, this is equivalent to \c Action::SkipChildren(node).
+    /// Otherwise, it is equivalent to \c Action::Continue(node).
+    template <typename T>
+    static _Detail::SkipChildrenIfWalkResult<T>
+    SkipChildrenIf(bool cond, T node) {
+      return {SkipChildrenIf(cond), std::move(node)};
+    }
+
     /// If \p cond is true, this is equivalent to \c Action::SkipNode(node).
     /// Otherwise, it is equivalent to \c Action::Continue(node).
     template <typename T>
@@ -231,6 +248,14 @@ class ASTWalker {
       return {SkipNodeIf(cond), std::move(node)};
     }
 
+    /// If \p cond is true, this is equivalent to \c Action::Continue(node).
+    /// Otherwise, it is equivalent to \c Action::SkipChildren(node).
+    template <typename T>
+    static _Detail::SkipChildrenIfWalkResult<T>
+    VisitChildrenIf(bool cond, T node) {
+      return SkipChildrenIf(!cond, std::move(node));
+    }
+
     /// If \p cond is true, this is equivalent to \c Action::Continue(node).
     /// Otherwise, it is equivalent to \c Action::SkipNode(node).
     template <typename T>
@@ -249,16 +274,35 @@ class ASTWalker {
     /// Continue the current walk.
     static _Detail::ContinueWalkAction Continue() { return {}; }
 
-    /// Skips visiting both the node's children and its post-visitation.
+    /// Continue the current walk, but do not visit the children of the current
+    /// node, resuming at its post-walk.
+    static _Detail::SkipChildrenIfWalkAction SkipChildren() {
+      return SkipChildrenIf(true);
+    }
+
+    /// Similar to \c Action::SkipChildren, but also skips the call to the
+    /// post-visitation method.
     static _Detail::SkipChildrenIfWalkAction SkipNode() {
       return SkipNodeIf(true);
     }
 
+    /// If \p cond is true, this is equivalent to \c Action::SkipChildren().
+    /// Otherwise, it is equivalent to \c Action::Continue().
+    static _Detail::SkipChildrenIfWalkAction SkipChildrenIf(bool cond) {
+      return {cond, /*SkipPostWalk*/ false};
+    }
+
     /// If \p cond is true, this is equivalent to \c Action::SkipNode().
     /// Otherwise, it is equivalent to \c Action::Continue().
     static _Detail::SkipChildrenIfWalkAction
     SkipNodeIf(bool cond) {
-      return {cond};
+      return {cond, /*SkipPostWalk*/ true};
+    }
+
+    /// If \p cond is true, this is equivalent to \c Action::Continue().
+    /// Otherwise, it is equivalent to \c Action::SkipChildren().
+    static _Detail::SkipChildrenIfWalkAction VisitChildrenIf(bool cond) {
+      return SkipChildrenIf(!cond);
     }
 
     /// If \p cond is true, this is equivalent to \c Action::Continue().
@@ -281,14 +325,19 @@ class ASTWalker {
   /// A pre-visitation action for AST nodes that do not support being replaced
   /// while walking.
   struct PreWalkAction {
-    enum Kind { Stop, SkipNode, Continue };
+    enum Kind { Stop, SkipChildren, SkipNode, Continue };
     Kind Action;
 
     PreWalkAction(_Detail::ContinueWalkAction) : Action(Continue) {}
     PreWalkAction(_Detail::StopWalkAction) : Action(Stop) {}
 
-    PreWalkAction(_Detail::SkipChildrenIfWalkAction action)
-        : Action(action.Cond ? SkipNode : Continue) {}
+    PreWalkAction(_Detail::SkipChildrenIfWalkAction action) {
+      if (action.Cond) {
+        Action = action.SkipPostWalk ? SkipNode : SkipChildren;
+      } else {
+        Action = Continue;
+      }
+    }
 
     PreWalkAction(_Detail::StopIfWalkAction action)
         : Action(action.Cond ? Stop : Continue) {}

lib/AST/ASTWalker.cpp

@@ -1457,11 +1457,13 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       return true;
     case PreWalkAction::SkipNode:
       return false;
+    case PreWalkAction::SkipChildren:
+      break;
     case PreWalkAction::Continue:
+      if (VisitChildren())
+        return true;
       break;
     }
-    if (VisitChildren())
-      return true;
     switch (WalkPost().Action) {
     case PostWalkAction::Stop:
       return true;
@@ -1476,21 +1478,25 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   T *traverse(PreWalkResult<T *> Pre,
               llvm::function_ref<T *(T *)> VisitChildren,
               llvm::function_ref<PostWalkResult<T *>(T *)> WalkPost) {
+    auto Node = Pre.Value;
+    assert(!Node || *Node && "Use Action::Stop instead of returning nullptr");
     switch (Pre.Action.Action) {
     case PreWalkAction::Stop:
       return nullptr;
     case PreWalkAction::SkipNode:
-      assert(*Pre.Value && "Use Action::Stop instead of returning nullptr");
-      return *Pre.Value;
-    case PreWalkAction::Continue:
+      return *Node;
+    case PreWalkAction::SkipChildren:
       break;
-    }
-    assert(*Pre.Value && "Use Action::Stop instead of returning nullptr");
-    auto Value = VisitChildren(*Pre.Value);
-    if (!Value)
-      return nullptr;
+    case PreWalkAction::Continue: {
+      auto NewNode = VisitChildren(*Node);
+      if (!NewNode)
+        return nullptr;
 
-    auto Post = WalkPost(Value);
+      Node = NewNode;
+      break;
+    }
+    }
+    auto Post = WalkPost(*Node);
     switch (Post.Action.Action) {
     case PostWalkAction::Stop:
       return nullptr;