Add a new level parameter to the PreOrder and PostOrder visitors.

Author: drupol

spec/drupol/phptree/Visitor/PostOrderVisitorSpec.php

@@ -4,8 +4,13 @@
 
 namespace spec\drupol\phptree\Visitor;
 
+use drupol\phptree\Node\NaryNode;
 use drupol\phptree\Node\Node;
+use drupol\phptree\Node\ValueNode;
+use drupol\phptree\Render\GraphViz;
+use drupol\phptree\Visitor\BreadthFirstVisitor;
 use drupol\phptree\Visitor\PostOrderVisitor;
+use Fhaculty\Graph\Graph;
 use PhpSpec\ObjectBehavior;
 
 class PostOrderVisitorSpec extends ObjectBehavior
@@ -34,4 +39,37 @@ public function it_can_traverse_a_tree()
             ->traverse($tree)
             ->shouldYield(new \ArrayIterator($nodes));
     }
+
+    public function it_can_traverse_a_tree_with_a_specific_level()
+    {
+        $tree = new NaryNode(2);
+
+        $data = \range('A', 'G');
+
+        $nodes = [];
+        foreach ($data as $key => $value) {
+            $nodes[$value] = new ValueNode($value);
+        }
+
+        $tree->add(...array_values($nodes));
+
+        $nodes['root'] = $tree;
+
+        $order = [
+            7 => $nodes['root'],
+        ];
+
+        $this
+            ->traverse($tree, 0)
+            ->shouldYield(new \ArrayIterator($order));
+
+        $order = [
+            3 => $nodes['A'],
+            6 => $nodes['B'],
+        ];
+
+        $this
+            ->traverse($tree, 1)
+            ->shouldYield(new \ArrayIterator($order));
+    }
 }

spec/drupol/phptree/Visitor/PreOrderVisitorSpec.php

@@ -4,7 +4,9 @@
 
 namespace spec\drupol\phptree\Visitor;
 
+use drupol\phptree\Node\NaryNode;
 use drupol\phptree\Node\Node;
+use drupol\phptree\Node\ValueNode;
 use drupol\phptree\Visitor\PreOrderVisitor;
 use PhpSpec\ObjectBehavior;
 
@@ -34,4 +36,37 @@ public function it_can_traverse_a_tree()
             ->traverse($tree)
             ->shouldYield(new \ArrayIterator($rootAndNodes));
     }
+
+    public function it_can_traverse_a_tree_with_a_specific_level()
+    {
+        $tree = new NaryNode(2);
+
+        $data = \range('A', 'G');
+
+        $nodes = [];
+        foreach ($data as $key => $value) {
+            $nodes[$value] = new ValueNode($value);
+        }
+
+        $tree->add(...array_values($nodes));
+
+        $nodes['root'] = $tree;
+
+        $order = [
+            0 => $nodes['root'],
+        ];
+
+        $this
+            ->traverse($tree, 0)
+            ->shouldYield(new \ArrayIterator($order));
+
+        $order = [
+            1 => $nodes['A'],
+            5 => $nodes['B'],
+        ];
+
+        $this
+            ->traverse($tree, 1)
+            ->shouldYield(new \ArrayIterator($order));
+    }
 }

src/Visitor/PostOrderVisitor.php

@@ -19,23 +19,25 @@ class PostOrderVisitor extends AbstractVisitor
     /**
      * {@inheritdoc}
      */
-    public function traverse(NodeInterface $node): \Traversable
+    public function traverse(NodeInterface $node, int $level = null): \Traversable
     {
         $this->index = 0;
 
-        return $this->doTraverse($node);
+        return $this->doTraverse($node, $level);
     }
 
     /**
      * {@inheritdoc}
      */
-    private function doTraverse(NodeInterface $node): \Traversable
+    private function doTraverse(NodeInterface $node, int $level = null): \Traversable
     {
         foreach ($node->children() as $child) {
-            yield from $this->doTraverse($child);
+            yield from $this->doTraverse($child, $level);
             $this->index++;
         }
 
-        yield $this->index => $node;
+        if (null === $level || $level === $node->depth()) {
+            yield $this->index => $node;
+        }
     }
 }

src/Visitor/PreOrderVisitor.php

@@ -19,23 +19,25 @@ class PreOrderVisitor extends AbstractVisitor
     /**
      * {@inheritdoc}
      */
-    public function traverse(NodeInterface $node): \Traversable
+    public function traverse(NodeInterface $node, int $level = null): \Traversable
     {
         $this->index = 0;
 
-        return $this->doTraverse($node);
+        return $this->doTraverse($node, $level);
     }
 
     /**
      * {@inheritdoc}
      */
-    private function doTraverse(NodeInterface $node): \Traversable
+    private function doTraverse(NodeInterface $node, int $level = null): \Traversable
     {
-        yield $this->index => $node;
+        if (null === $level || $level === $node->depth()) {
+            yield $this->index => $node;
+        }
 
         foreach ($node->children() as $child) {
             $this->index++;
-            yield from $this->doTraverse($child);
+            yield from $this->doTraverse($child, $level);
         }
     }
 }