Add new constant to be returned from enterNode() to not traverse current and child nodes (#536)

* Add new constant to be returned from enterNode() to not travers current node for subsequent visitors and skip children traversing

* Allow visitors to replace nodes in leaveNode() when DONT_TRAVERSE_CURRENT_AND_CHILDREN is used

lib/PhpParser/NodeTraverser.php

@@ -30,6 +30,15 @@ class NodeTraverser implements NodeTraverserInterface
      */
     const REMOVE_NODE = 3;
 
+    /**
+     * If NodeVisitor::enterNode() returns DONT_TRAVERSE_CURRENT_AND_CHILDREN, child nodes
+     * of the current node will not be traversed for any visitors.
+     *
+     * For subsequent visitors enterNode() will not be called as well.
+     * leaveNode() will be invoked for visitors that has enterNode() method invoked.
+     */
+    const DONT_TRAVERSE_CURRENT_AND_CHILDREN = 4;
+
     /** @var NodeVisitor[] Visitors */
     protected $visitors = [];
 
@@ -108,7 +117,9 @@ class NodeTraverser implements NodeTraverserInterface
                 }
             } elseif ($subNode instanceof Node) {
                 $traverseChildren = true;
-                foreach ($this->visitors as $visitor) {
+                $breakVisitorIndex = null;
+
+                foreach ($this->visitors as $visitorIndex => $visitor) {
                     $return = $visitor->enterNode($subNode);
                     if (null !== $return) {
                         if ($return instanceof Node) {
@@ -116,6 +127,10 @@ class NodeTraverser implements NodeTraverserInterface
                             $subNode = $return;
                         } elseif (self::DONT_TRAVERSE_CHILDREN === $return) {
                             $traverseChildren = false;
+                        } elseif (self::DONT_TRAVERSE_CURRENT_AND_CHILDREN === $return) {
+                            $traverseChildren = false;
+                            $breakVisitorIndex = $visitorIndex;
+                            break;
                         } elseif (self::STOP_TRAVERSAL === $return) {
                             $this->stopTraversal = true;
                             break 2;
@@ -134,8 +149,9 @@ class NodeTraverser implements NodeTraverserInterface
                     }
                 }
 
-                foreach ($this->visitors as $visitor) {
+                foreach ($this->visitors as $visitorIndex => $visitor) {
                     $return = $visitor->leaveNode($subNode);
+
                     if (null !== $return) {
                         if ($return instanceof Node) {
                             $this->ensureReplacementReasonable($subNode, $return);
@@ -154,6 +170,10 @@ class NodeTraverser implements NodeTraverserInterface
                             );
                         }
                     }
+
+                    if ($breakVisitorIndex === $visitorIndex) {
+                        break;
+                    }
                 }
             }
         }
@@ -174,7 +194,9 @@ class NodeTraverser implements NodeTraverserInterface
         foreach ($nodes as $i => &$node) {
             if ($node instanceof Node) {
                 $traverseChildren = true;
-                foreach ($this->visitors as $visitor) {
+                $breakVisitorIndex = null;
+
+                foreach ($this->visitors as $visitorIndex => $visitor) {
                     $return = $visitor->enterNode($node);
                     if (null !== $return) {
                         if ($return instanceof Node) {
@@ -182,6 +204,10 @@ class NodeTraverser implements NodeTraverserInterface
                             $node = $return;
                         } elseif (self::DONT_TRAVERSE_CHILDREN === $return) {
                             $traverseChildren = false;
+                        } elseif (self::DONT_TRAVERSE_CURRENT_AND_CHILDREN === $return) {
+                            $traverseChildren = false;
+                            $breakVisitorIndex = $visitorIndex;
+                            break;
                         } elseif (self::STOP_TRAVERSAL === $return) {
                             $this->stopTraversal = true;
                             break 2;
@@ -200,8 +226,9 @@ class NodeTraverser implements NodeTraverserInterface
                     }
                 }
 
-                foreach ($this->visitors as $visitor) {
+                foreach ($this->visitors as $visitorIndex => $visitor) {
                     $return = $visitor->leaveNode($node);
+
                     if (null !== $return) {
                         if ($return instanceof Node) {
                             $this->ensureReplacementReasonable($node, $return);
@@ -226,6 +253,10 @@ class NodeTraverser implements NodeTraverserInterface
                             );
                         }
                     }
+
+                    if ($breakVisitorIndex === $visitorIndex) {
+                        break;
+                    }
                 }
             } elseif (\is_array($node)) {
                 throw new \LogicException('Invalid node structure: Contains nested arrays');

test/PhpParser/NodeTraverserTest.php

@@ -165,6 +165,42 @@ class NodeTraverserTest extends TestCase
         $this->assertEquals($stmts, $traverser->traverse($stmts));
     }
 
+    public function testDontTraverseCurrentAndChildren() {
+        // print 'str'; -($foo * $foo);
+        $strNode = new String_('str');
+        $printNode = new Expr\Print_($strNode);
+        $varNode = new Expr\Variable('foo');
+        $mulNode = new Expr\BinaryOp\Mul($varNode, $varNode);
+        $divNode = new Expr\BinaryOp\Div($varNode, $varNode);
+        $negNode = new Expr\UnaryMinus($mulNode);
+        $stmts = [$printNode, $negNode];
+
+        $visitor1 = $this->getMockBuilder(NodeVisitor::class)->getMock();
+        $visitor2 = $this->getMockBuilder(NodeVisitor::class)->getMock();
+
+        $visitor1->expects($this->at(1))->method('enterNode')->with($printNode)
+            ->will($this->returnValue(NodeTraverser::DONT_TRAVERSE_CURRENT_AND_CHILDREN));
+        $visitor1->expects($this->at(2))->method('leaveNode')->with($printNode);
+
+        $visitor1->expects($this->at(3))->method('enterNode')->with($negNode);
+        $visitor2->expects($this->at(1))->method('enterNode')->with($negNode);
+
+        $visitor1->expects($this->at(4))->method('enterNode')->with($mulNode)
+            ->will($this->returnValue(NodeTraverser::DONT_TRAVERSE_CURRENT_AND_CHILDREN));
+        $visitor1->expects($this->at(5))->method('leaveNode')->with($mulNode)->willReturn($divNode);
+
+        $visitor1->expects($this->at(6))->method('leaveNode')->with($negNode);
+        $visitor2->expects($this->at(2))->method('leaveNode')->with($negNode);
+
+        $traverser = new NodeTraverser;
+        $traverser->addVisitor($visitor1);
+        $traverser->addVisitor($visitor2);
+
+        $resultStmts = $traverser->traverse($stmts);
+
+        $this->assertInstanceOf(Expr\BinaryOp\Div::class, $resultStmts[1]->expr);
+    }
+
     public function testStopTraversal() {
         $varNode1 = new Expr\Variable('a');
         $varNode2 = new Expr\Variable('b');