Merge pull request #70600 from hamishknight/postal-service

[ASTWalker] Continue post-walking with `Action::SkipChildren`

Author: hamishknight

include/swift/AST/ASTWalker.h

@@ -163,23 +163,24 @@ 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 {};
 
     // The 'Result' set of types, which do take a payload.
     template <typename T>
     struct ContinueWalkResult {
+      ContinueWalkAction Action;
       T Value;
     };
     template <typename T>
     struct SkipChildrenIfWalkResult {
-      bool Cond;
+      SkipChildrenIfWalkAction Action;
       T Value;
     };
     template <typename T>
     struct StopIfWalkResult {
-      bool Cond;
+      StopIfWalkAction Action;
       T Value;
     };
   };
@@ -212,23 +213,39 @@ class ASTWalker {
     /// Continue the current walk, replacing the current node with \p node.
     template <typename T>
     static _Detail::ContinueWalkResult<T> Continue(T node) {
-      return {std::move(node)};
+      return {Continue(), std::move(node)};
     }
 
     /// Continue the current walk, replacing the current node with \p node.
-    /// However, skip visiting the children of \p node, and instead resume the
-    /// walk of the parent 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 {cond, std::move(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>
+    static _Detail::SkipChildrenIfWalkResult<T>
+    SkipNodeIf(bool cond, T node) {
+      return {SkipNodeIf(cond), std::move(node)};
     }
 
     /// If \p cond is true, this is equivalent to \c Action::Continue(node).
@@ -239,26 +256,47 @@ class ASTWalker {
       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>
+    static _Detail::SkipChildrenIfWalkResult<T>
+    VisitNodeIf(bool cond, T node) {
+      return SkipNodeIf(!cond, std::move(node));
+    }
+
     /// If \p cond is true, this is equivalent to \c Action::Stop().
     /// Otherwise, it is equivalent to \c Action::Continue(node).
     template <typename T>
     static _Detail::StopIfWalkResult<T> StopIf(bool cond, T node) {
-      return {cond, std::move(node)};
+      return {StopIf(cond), std::move(node)};
     }
 
     /// Continue the current walk.
     static _Detail::ContinueWalkAction Continue() { return {}; }
 
     /// Continue the current walk, but do not visit the children of the current
-    /// node. Instead, resume at the parent's post-walk.
+    /// 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};
+      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, /*SkipPostWalk*/ true};
     }
 
     /// If \p cond is true, this is equivalent to \c Action::Continue().
@@ -267,6 +305,13 @@ class ASTWalker {
       return SkipChildrenIf(!cond);
     }
 
+    /// If \p cond is true, this is equivalent to \c Action::Continue().
+    /// Otherwise, it is equivalent to \c Action::SkipNode().
+    static _Detail::SkipChildrenIfWalkAction
+    VisitNodeIf(bool cond) {
+      return SkipNodeIf(!cond);
+    }
+
     /// Terminate the walk, returning without visiting any other nodes.
     static _Detail::StopWalkAction Stop() { return {}; }
 
@@ -280,16 +325,19 @@ class ASTWalker {
   /// A pre-visitation action for AST nodes that do not support being replaced
   /// while walking.
   struct PreWalkAction {
-    enum Kind { Stop, SkipChildren, Continue };
+    enum Kind { Stop, SkipChildren, SkipNode, Continue };
     Kind Action;
 
-    PreWalkAction(Kind action) : Action(action) {}
-
     PreWalkAction(_Detail::ContinueWalkAction) : Action(Continue) {}
     PreWalkAction(_Detail::StopWalkAction) : Action(Stop) {}
 
-    PreWalkAction(_Detail::SkipChildrenIfWalkAction action)
-        : Action(action.Cond ? SkipChildren : 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) {}
@@ -303,8 +351,6 @@ class ASTWalker {
     enum Kind { Stop, Continue };
     Kind Action;
 
-    PostWalkAction(Kind action) : Action(action) {}
-
     PostWalkAction(_Detail::ContinueWalkAction) : Action(Continue) {}
     PostWalkAction(_Detail::StopWalkAction) : Action(Stop) {}
 
@@ -340,21 +386,18 @@ class ASTWalker {
 
     template <typename U>
     PreWalkResult(_Detail::ContinueWalkResult<U> Result)
-        : Action(PreWalkAction::Continue), Value(std::move(Result.Value)) {}
+        : Action(Result.Action), Value(std::move(Result.Value)) {}
 
     template <typename U>
     PreWalkResult(_Detail::SkipChildrenIfWalkResult<U> Result)
-        : Action(Result.Cond ? PreWalkAction::SkipChildren
-                             : PreWalkAction::Continue),
-          Value(std::move(Result.Value)) {}
+        : Action(Result.Action), Value(std::move(Result.Value)) {}
 
     template <typename U>
     PreWalkResult(_Detail::StopIfWalkResult<U> Result)
-        : Action(Result.Cond ? PreWalkAction::Stop : PreWalkAction::Continue),
-          Value(std::move(Result.Value)) {}
+        : Action(Result.Action), Value(std::move(Result.Value)) {}
 
-    PreWalkResult(_Detail::StopWalkAction)
-        : Action(PreWalkAction::Stop), Value(llvm::None) {}
+    PreWalkResult(_Detail::StopWalkAction Action)
+        : Action(Action), Value(llvm::None) {}
   };
 
   /// Do not construct directly, use \c Action::<action> instead.
@@ -385,15 +428,14 @@ class ASTWalker {
 
     template <typename U>
     PostWalkResult(_Detail::ContinueWalkResult<U> Result)
-        : Action(PostWalkAction::Continue), Value(std::move(Result.Value)) {}
+        : Action(Result.Action), Value(std::move(Result.Value)) {}
 
     template <typename U>
     PostWalkResult(_Detail::StopIfWalkResult<U> Result)
-        : Action(Result.Cond ? PostWalkAction::Stop : PostWalkAction::Continue),
-          Value(std::move(Result.Value)) {}
+        : Action(Result.Action), Value(std::move(Result.Value)) {}
 
-    PostWalkResult(_Detail::StopWalkAction)
-        : Action(PostWalkAction::Stop), Value(llvm::None) {}
+    PostWalkResult(_Detail::StopWalkAction Action)
+        : Action(Action), Value(llvm::None) {}
   };
 
   /// This method is called when first visiting an expression

include/swift/AST/TypeWalker.h

@@ -22,7 +22,7 @@ class TypeWalker {
 public:
   enum class Action {
     Continue,
-    SkipChildren,
+    SkipNode,
     Stop
   };
 

include/swift/Sema/ConstraintSystem.h

@@ -2685,12 +2685,12 @@ class ConstraintSystem {
 
     /// Ignore statements.
     PreWalkResult<Stmt *> walkToStmtPre(Stmt *stmt) override {
-      return Action::SkipChildren(stmt);
+      return Action::SkipNode(stmt);
     }
 
     /// Ignore declarations.
     PreWalkAction walkToDeclPre(Decl *decl) override {
-      return Action::SkipChildren();
+      return Action::SkipNode();
     }
   };
 
@@ -6389,7 +6389,7 @@ class TypeVarRefCollector : public ASTWalker {
   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));
+    return Action::VisitNodeIf(isa<PatternBindingDecl>(D));
   }
 
   ArrayRef<TypeVariableType *> getTypeVars() const {

lib/APIDigester/ModuleAnalyzerNodes.cpp

@@ -2445,10 +2445,10 @@ class ConstExtractor: public ASTWalker {
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     if (E->isSemanticallyConstExpr()) {
       record(E, E);
-      return Action::SkipChildren(E);
+      return Action::SkipNode(E);
     }
     if (handleSimpleReference(E)) {
-      return Action::SkipChildren(E);
+      return Action::SkipNode(E);
     }
     return Action::Continue(E);
   }

lib/AST/ASTScopeCreation.cpp

@@ -115,47 +115,45 @@ class ScopeCreator final : public ASTAllocated<ScopeCreator> {
           scopeCreator
               .constructExpandAndInsert<ClosureParametersScope>(
                   parent, closure);
-          return Action::SkipChildren(E);
+          return Action::SkipNode(E);
         }
         if (auto *capture = dyn_cast<CaptureListExpr>(E)) {
           scopeCreator
               .constructExpandAndInsert<CaptureListScope>(
                   parent, capture);
-          return Action::SkipChildren(E);
+          return Action::SkipNode(E);
         }
 
         // If we have a single value statement expression, we need to add any
         // scopes in the underlying statement.
         if (auto *SVE = dyn_cast<SingleValueStmtExpr>(E)) {
           scopeCreator.addToScopeTree(SVE->getStmt(), parent);
-          return Action::SkipChildren(E);
+          return Action::SkipNode(E);
         }
 
         // If we have a try/try!/try?, we need to add a scope for it
         if (auto anyTry = dyn_cast<AnyTryExpr>(E)) {
           scopeCreator.constructExpandAndInsert<TryScope>(parent, anyTry);
-          return Action::SkipChildren(E);
+          return Action::SkipNode(E);
         }
 
         return Action::Continue(E);
       }
       PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
-        if (isa<BraceStmt>(S)) { // closures hidden in here
-          return Action::Continue(S);
-        }
-        return Action::SkipChildren(S);
+        // Closures can occur in BraceStmts.
+        return Action::VisitNodeIf(isa<BraceStmt>(S), S);
       }
       PreWalkResult<Pattern *> walkToPatternPre(Pattern *P) override {
-        return Action::SkipChildren(P);
+        return Action::SkipNode(P);
       }
       PreWalkAction walkToDeclPre(Decl *D) override {
-        return Action::SkipChildren();
+        return Action::SkipNode();
       }
       PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
-        return Action::SkipChildren();
+        return Action::SkipNode();
       }
       PreWalkAction walkToParameterListPre(ParameterList *PL) override {
-        return Action::SkipChildren();
+        return Action::SkipNode();
       }
 
       MacroWalking getMacroWalkingBehavior() const override {

lib/AST/ASTVerifier.cpp

@@ -125,7 +125,7 @@ ASTWalker::PreWalkResult<Expr *> dispatchVisitPreExprHelper(
     return ASTWalker::Action::Continue(node);
   }
   V.cleanup(node);
-  return ASTWalker::Action::SkipChildren(node);
+  return ASTWalker::Action::SkipNode(node);
 }
 
 template <typename Verifier, typename Kind>
@@ -141,7 +141,7 @@ ASTWalker::PreWalkResult<Expr *> dispatchVisitPreExprHelper(
     return ASTWalker::Action::Continue(node);
   }
   V.cleanup(node);
-  return ASTWalker::Action::SkipChildren(node);
+  return ASTWalker::Action::SkipNode(node);
 }
 
 template <typename Verifier, typename Kind>
@@ -157,7 +157,7 @@ ASTWalker::PreWalkResult<Expr *> dispatchVisitPreExprHelper(
     return ASTWalker::Action::Continue(node);
   }
   V.cleanup(node);
-  return ASTWalker::Action::SkipChildren(node);
+  return ASTWalker::Action::SkipNode(node);
 }
 
 template <typename Verifier, typename Kind>
@@ -170,7 +170,7 @@ ASTWalker::PreWalkResult<Expr *> dispatchVisitPreExprHelper(
     return ASTWalker::Action::Continue(node);
   }
   V.cleanup(node);
-  return ASTWalker::Action::SkipChildren(node);
+  return ASTWalker::Action::SkipNode(node);
 }
 
 namespace {
@@ -401,7 +401,7 @@ class Verifier : public ASTWalker {
       if (shouldVerify(node))
         return Action::Continue();
       cleanup(node);
-      return Action::SkipChildren();
+      return Action::SkipNode();
     }
 
     /// Helper template for dispatching pre-visitation.
@@ -419,7 +419,7 @@ class Verifier : public ASTWalker {
       if (shouldVerify(node))
         return Action::Continue(node);
       cleanup(node);
-      return Action::SkipChildren(node);
+      return Action::SkipNode(node);
     }
 
     /// Helper template for dispatching pre-visitation.
@@ -430,7 +430,7 @@ class Verifier : public ASTWalker {
       if (shouldVerify(node))
         return Action::Continue(node);
       cleanup(node);
-      return Action::SkipChildren(node);
+      return Action::SkipNode(node);
     }
 
     /// Helper template for dispatching post-visitation.