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
This commit is contained in:
Maks Rafalko 2018-10-08 23:26:00 +03:00 committed by Nikita Popov
parent 674c5610fb
commit dc323458b4
2 changed files with 71 additions and 4 deletions

View File

@ -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');

View File

@ -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');