Add max depth parameter to breadth first search

src/Search/BreadthFirst.php

@@ -9,25 +9,43 @@
 class BreadthFirst extends Base
 {
     /**
-     *
+     * @param int $maxDepth
      * @return Vertices
      */
-    public function getVertices()
+    public function getVertices($maxDepth = PHP_INT_MAX)
     {
         $queue = array($this->vertex);
         // to not add vertices twice in array visited
         $mark = array($this->vertex->getId() => true);
         // visited vertices
         $visited = array();
 
+        // keep track of depth
+        $currentDepth = 0;
+        $nodesThisLevel = 1;
+        $nodesNextLevel = 0;
+
         do {
             // get first from queue
             $t = array_shift($queue);
             // save as visited
-            $visited[$t->getId()]= $t;
+            $visited[$t->getId()] = $t;
 
             // get next vertices
-            foreach ($this->getVerticesAdjacent($t)->getMap() as $id => $vertex) {
+            $children = $this->getVerticesAdjacent($t);
+
+            // track depth
+            $nodesNextLevel = $children->count();
+            if (--$nodesThisLevel === 0) {
+                if (++$currentDepth > $maxDepth) {
+                    return new Vertices($visited);
+                }
+                $nodesThisLevel = $nodesNextLevel;
+                $nodesNextLevel = 0;
+            }
+
+            // process next vertices
+            foreach ($children->getMap() as $id => $vertex) {
                 // if not "touched" before
                 if (!isset($mark[$id])) {
                     // add to queue
@@ -37,7 +55,7 @@ public function getVertices()
                 }
             }
 
-        // untill queue is empty
+            // until queue is empty
         } while ($queue);
 
         return new Vertices($visited);

tests/Search/BreadthFirstTest.php

@@ -0,0 +1,55 @@
+<?php
+
+use Fhaculty\Graph\Graph;
+use Graphp\Algorithms\Search\BreadthFirst;
+
+class BreadthFirstSearchTest extends TestCase
+{
+    public function providerMaxDepth()
+    {
+        return array(
+            "simple path (no limit)" => array(
+                "edges" => array(
+                    array(1, 2), array(2, 3), array(3, 4), array(4, 5),
+                ),
+                "subject" => 1,
+                "maxDepth" => null,
+                "expected" => array(1, 2, 3, 4, 5),
+            ),
+            "simple path (limit = 0)" => array(
+                "edges" => array(
+                    array(1, 2), array(2, 3), array(3, 4), array(4, 5),
+                ),
+                "subject" => 1,
+                "maxDepth" => 0,
+                "expected" => array(1),
+            ),
+            "simple path (limit = 1)" => array(
+                "edges" => array(
+                    array(1, 2), array(2, 3), array(3, 4), array(4, 5),
+                ),
+                "subject" => 1,
+                "maxDepth" => 1,
+                "expected" => array(1, 2),
+            ),
+        );
+    }
+
+    /**
+     * @dataProvider providerMaxDepth
+     */
+    public function testMaxDepth(array $edges, $subject, $maxDepth, array $expected)
+    {
+        $g = new Graph();
+        foreach ($edges as $e) {
+            $g->createVertex($e[0], true)->createEdgeTo($g->createVertex($e[1], true));
+        }
+        $a = new BreadthFirst($g->getVertex($subject));
+        if ($maxDepth !== null) {
+            $v = $a->getVertices($maxDepth);
+        } else {
+            $v = $a->getVertices(); // Simulate default
+        }
+        $this->assertSame($expected, $v->getIds());
+    }
+}