[Tree] Refactored "childCount" functionality, and added to RepositoryInterface. Included this missing functionality to Materialized Path strategies too.

lib/Gedmo/Tree/Document/MongoDB/Repository/MaterializedPathRepository.php

@@ -81,6 +81,32 @@ class MaterializedPathRepository extends AbstractTreeRepository
         return $this->getRootNodesQuery($sortByField, $direction)->execute();
     }
 
+    /**
+     * {@inheritDoc}
+     */
+    public function childCount($node = null, $direct = false)
+    {
+        $meta = $this->getClassMetadata();
+
+        if (is_object($node)) {
+            if (!($node instanceof $meta->name)) {
+                throw new InvalidArgumentException("Node is not related to this repository");
+            }
+
+            $wrapped = new MongoDocumentWrapper($node, $this->dm);
+
+            if (!$wrapped->hasValidIdentifier()) {
+                throw new InvalidArgumentException("Node is not managed by UnitOfWork");
+            }
+        }
+
+        $qb = $this->getChildrenQueryBuilder($node, $direct);
+
+        $qb->count();
+
+        return (int) $qb->getQuery()->execute();
+    }
+
     /**
      * {@inheritDoc}
      */

lib/Gedmo/Tree/Entity/Repository/AbstractTreeRepository.php

@@ -8,7 +8,8 @@ use Doctrine\ORM\EntityRepository,
     Gedmo\Tool\Wrapper\EntityWrapper,
     Gedmo\Tree\RepositoryUtils,
     Gedmo\Tree\RepositoryUtilsInterface,
-    Gedmo\Tree\RepositoryInterface;
+    Gedmo\Tree\RepositoryInterface,
+    Gedmo\Exception\InvalidArgumentException;
 
 abstract class AbstractTreeRepository extends EntityRepository implements RepositoryInterface
 {
@@ -79,6 +80,36 @@ abstract class AbstractTreeRepository extends EntityRepository implements Reposi
         return $this->repoUtils;
     }
 
+    /**
+     * {@inheritDoc}
+     */
+    public function childCount($node = null, $direct = false)
+    {
+        $meta = $this->getClassMetadata();
+
+        if (is_object($node)) {
+            if (!($node instanceof $meta->name)) {
+                throw new InvalidArgumentException("Node is not related to this repository");
+            }
+
+            $wrapped = new EntityWrapper($node, $this->_em);
+
+            if (!$wrapped->hasValidIdentifier()) {
+                throw new InvalidArgumentException("Node is not managed by UnitOfWork");
+            }
+        }
+
+        $qb = $this->getChildrenQueryBuilder($node, $direct);
+        $aliases = $qb->getRootAliases();
+        $alias = $aliases[0];
+
+        $qb->select('COUNT('.$alias.')');
+
+        $result = $qb->getQuery()->getScalarResult();
+
+        return (int) $result[0][1];
+    }
+
     /**
      * @see \Gedmo\Tree\RepositoryUtilsInterface::childrenHierarchy
      */

lib/Gedmo/Tree/Entity/Repository/ClosureTreeRepository.php

@@ -61,54 +61,6 @@ class ClosureTreeRepository extends AbstractTreeRepository
         return $this->getRootNodesQuery($sortByField, $direction)->getResult();
     }
 
-    /**
-     * Counts the children of given TreeNode
-     *
-     * @param object $node - if null counts all records in tree
-     * @param boolean $direct - true to count only direct children
-     * @throws InvalidArgumentException - if input is not valid
-     * @return integer
-     */
-    public function childCount($node = null, $direct = false)
-    {
-        $count = 0;
-        $meta = $this->getClassMetadata();
-        $config = $this->listener->getConfiguration($this->_em, $meta->name);
-        if (null !== $node) {
-            if ($node instanceof $meta->name) {
-                if (!$this->_em->getUnitOfWork()->isInIdentityMap($node)) {
-                    throw new InvalidArgumentException("Node is not managed by UnitOfWork");
-                }
-                if ($direct) {
-                    $qb = $this->_em->createQueryBuilder();
-                    $qb->select('COUNT(node)')
-                        ->from($config['useObjectClass'], 'node')
-                        ->where('node.' . $config['parent'] . ' = :node');
-
-                    $q = $qb->getQuery();
-                } else {
-                    $closureMeta = $this->_em->getClassMetadata($config['closure']);
-                    $dql = "SELECT COUNT(c) FROM {$closureMeta->name} c";
-                    $dql .= " WHERE c.ancestor = :node";
-                    $dql .= " AND c.descendant <> :node";
-                    $q = $this->_em->createQuery($dql);
-                }
-                $q->setParameters(compact('node'));
-                $count = intval($q->getSingleScalarResult());
-            } else {
-                throw new InvalidArgumentException("Node is not related to this repository");
-            }
-        } else {
-            $dql = "SELECT COUNT(node) FROM " . $config['useObjectClass'] . " node";
-            if ($direct) {
-                $dql .= ' WHERE node.' . $config['parent'] . ' IS NULL';
-            }
-            $q = $this->_em->createQuery($dql);
-            $count = intval($q->getSingleScalarResult());
-        }
-        return $count;
-    }
-
     /**
      * Get the Tree path query by given $node
      *

lib/Gedmo/Tree/Entity/Repository/NestedTreeRepository.php

@@ -171,69 +171,6 @@ class NestedTreeRepository extends AbstractTreeRepository
         return $this->getPathQuery($node)->getResult();
     }
 
-    /**
-     * Counts the children of given TreeNode
-     *
-     * @param object $node - if null counts all records in tree
-     * @param boolean $direct - true to count only direct children
-     * @throws InvalidArgumentException - if input is not valid
-     * @return integer
-     */
-    public function childCount($node = null, $direct = false)
-    {
-        $count = 0;
-        $meta = $this->getClassMetadata();
-        $nodeId = $meta->getSingleIdentifierFieldName();
-        $config = $this->listener->getConfiguration($this->_em, $meta->name);
-        if (null !== $node) {
-            if ($node instanceof $meta->name) {
-                $wrapped = new EntityWrapper($node, $this->_em);
-                if (!$wrapped->hasValidIdentifier()) {
-                    throw new InvalidArgumentException("Node is not managed by UnitOfWork");
-                }
-                if ($direct) {
-                    $id = $wrapped->getIdentifier();
-                    $qb = $this->_em->createQueryBuilder();
-                    $qb->select($qb->expr()->count('node.' . $nodeId))
-                        ->from($config['useObjectClass'], 'node')
-                        ->where($id === null ?
-                            $qb->expr()->isNull('node.'.$config['parent']) :
-                            $qb->expr()->eq('node.'.$config['parent'], is_string($id) ? $qb->expr()->literal($id) : $id)
-                        )
-                    ;
-
-                    if (isset($config['root'])) {
-                        $rootId = $wrapped->getPropertyValue($config['root']);
-                        $qb->andWhere($rootId === null ?
-                            $qb->expr()->isNull('node.'.$config['root']) :
-                            $qb->expr()->eq('node.'.$config['root'], is_string($rootId) ? $qb->expr()->literal($rootId) : $rootId)
-                        );
-                    }
-                    $q = $qb->getQuery();
-                    $count = intval($q->getSingleScalarResult());
-                } else {
-                    $left = $wrapped->getPropertyValue($config['left']);
-                    $right = $wrapped->getPropertyValue($config['right']);
-                    if ($left && $right) {
-                        $count = ($right - $left - 1) / 2;
-                    }
-                }
-            } else {
-                throw new InvalidArgumentException("Node is not related to this repository");
-            }
-        } else {
-            $qb = $this->_em->createQueryBuilder();
-            $qb->select($qb->expr()->count('node.' . $nodeId))
-                ->from($config['useObjectClass'], 'node')
-            ;
-            if ($direct) {
-                $qb->where($qb->expr()->isNull('node.'.$config['parent']));
-            }
-            $count = intval($qb->getQuery()->getSingleScalarResult());
-        }
-        return $count;
-    }
-
     /**
      * @see getChildrenQueryBuilder
      */

lib/Gedmo/Tree/RepositoryInterface.php

@@ -45,4 +45,14 @@ interface RepositoryInterface extends RepositoryUtilsInterface
      * @return array - list of given $node children, null on failure
      */
     public function getChildren($node = null, $direct = false, $sortByField = null, $direction = 'ASC', $includeNode = false);
+
+    /**
+     * Counts the children of given TreeNode
+     *
+     * @param object $node - if null counts all records in tree
+     * @param boolean $direct - true to count only direct children
+     * @throws InvalidArgumentException - if input is not valid
+     * @return integer
+     */
+    public function childCount($node = null, $direct = false);
 }

tests/Gedmo/Tree/ClosureTreeRepositoryTest.php

@@ -45,15 +45,23 @@ class ClosureTreeRepositoryTest extends BaseTestCaseORM
         $repo = $this->em->getRepository(self::CATEGORY);
         $food = $repo->findOneByTitle('Food');
 
-        $directCount = $repo->childCount($food, true);
-        $this->assertEquals(3, $directCount);
+        // Count all
+        $count = $repo->childCount();
+        $this->assertEquals(15, $count);
 
-        $fruits = $repo->findOneByTitle('Fruits');
-        $count = $repo->childCount($fruits);
-        $this->assertEquals(4, $count);
+        // Count all, but only direct ones
+        $count = $repo->childCount(null, true);
+        $this->assertEquals(2, $count);
+
+        // Count food children
+        $food = $repo->findOneByTitle('Food');
+        $count = $repo->childCount($food);
+        $this->assertEquals(11, $count);
 
-        $rootCount = $repo->childCount(null, true);
-        $this->assertEquals(2, $rootCount);
+        // Count food children, but only direct ones
+        $food = $repo->findOneByTitle('Food');
+        $count = $repo->childCount($food, true);
+        $this->assertEquals(3, $count);
     }
 
     public function testPath()

tests/Gedmo/Tree/MaterializedPathODMMongoDBRepositoryTest.php

@@ -206,6 +206,33 @@ class MaterializedPathODMMongoDBRepositoryTest extends BaseTestCaseMongoODM
         $this->assertEquals('<ul><li>Drinks<ul><li>Whisky<ul><li>Best Whisky</li></ul></li></ul></li></ul>', $tree);
     }
 
+    public function testChildCount()
+    {
+        /** @var $repo \Gedmo\Tree\Document\MongoDB\Repository\MaterializedPathRepository */
+        $repo = $this->dm->getRepository(self::CATEGORY);
+
+        // Count all
+        $count = $repo->childCount();
+
+        $this->assertEquals(9, $count);
+
+        // Count all, but only direct ones
+        $count = $repo->childCount(null, true);
+
+        $this->assertEquals(3, $count);
+
+        // Count food children
+        $food = $repo->findOneByTitle('Food');
+        $count = $repo->childCount($food);
+
+        $this->assertEquals(4, $count);
+
+        // Count food children, but only direct ones
+        $count = $repo->childCount($food, true);
+
+        $this->assertEquals(2, $count);
+    }
+
     protected function getUsedEntityFixtures()
     {
         return array(

tests/Gedmo/Tree/MaterializedPathORMRepositoryTest.php

@@ -145,7 +145,7 @@ class MaterializedPathORMRepositoryTest extends BaseTestCaseORM
 
     public function testChildrenHierarchyMethod()
     {
-        /** @var $repo \Gedmo\Tree\Document\MongoDB\Repository\MaterializedPathRepository */
+        /** @var $repo \Gedmo\Tree\Entity\Repository\MaterializedPathRepository */
         $repo = $this->em->getRepository(self::CATEGORY);
         $tree = $repo->childrenHierarchy();
 
@@ -205,6 +205,33 @@ class MaterializedPathORMRepositoryTest extends BaseTestCaseORM
         $this->assertEquals('<ul><li>Drinks<ul><li>Whisky<ul><li>Best Whisky</li></ul></li></ul></li></ul>', $tree);
     }
 
+    public function testChildCount()
+    {
+        /** @var $repo \Gedmo\Tree\Entity\Repository\MaterializedPathRepository */
+        $repo = $this->em->getRepository(self::CATEGORY);
+
+        // Count all
+        $count = $repo->childCount();
+
+        $this->assertEquals(9, $count);
+
+        // Count all, but only direct ones
+        $count = $repo->childCount(null, true);
+
+        $this->assertEquals(3, $count);
+
+        // Count food children
+        $food = $repo->findOneByTitle('Food');
+        $count = $repo->childCount($food);
+
+        $this->assertEquals(4, $count);
+
+        // Count food children, but only direct ones
+        $count = $repo->childCount($food, true);
+
+        $this->assertEquals(2, $count);
+    }
+
     protected function getUsedEntityFixtures()
     {
         return array(

tests/Gedmo/Tree/NestedTreeRootRepositoryTest.php

@@ -219,6 +219,9 @@ class NestedTreeRootRepositoryTest extends BaseTestCaseORM
 
         $childCount = $repo->childCount();
         $this->assertEquals(6, $childCount);
+
+        $childCount = $repo->childCount(null, true);
+        $this->assertEquals(2, $childCount);
     }
 
     /**