tree repository childcount for table inherited nodes, issue7 still affects other places

lib/Gedmo/Tree/Repository/TreeNodeRepository.php

@@ -4,6 +4,7 @@ namespace Gedmo\Tree\Repository;
 
 use Doctrine\ORM\EntityRepository,
     Doctrine\ORM\Query,
+    Doctrine\ORM\Mapping\ClassMetadataInfo,
     Gedmo\Tree\Node;
 
 /**
@@ -76,7 +77,7 @@ class TreeNodeRepository extends EntityRepository
                 $id = $meta->getReflectionProperty($nodeId)->getValue($node);
                 $qb = $this->_em->createQueryBuilder();
                 $qb->select('COUNT(node.' . $nodeId . ')')
-                    ->from($this->_entityName, 'node')
+                    ->from($this->_getBaseEntityClass($meta), 'node')
                     ->where('node.' . $config['parent'] . ' = ' . $id);
                     
                 $q = $qb->getQuery();
@@ -89,7 +90,7 @@ class TreeNodeRepository extends EntityRepository
                 }
             }
         } else {
-            $dql = "SELECT COUNT(node.{$nodeId}) FROM {$this->_entityName} node";
+            $dql = "SELECT COUNT(node.{$nodeId}) FROM " . $this->_getBaseEntityClass($meta) . " node";
             if ($direct) {
                 $dql .= ' WHERE node.' . $config['parent'] . ' IS NULL';
             }
@@ -682,4 +683,29 @@ class TreeNodeRepository extends EntityRepository
             }
         }
     }
+    
+    /**
+     * If entity is mapped through table inheritance
+     * SINGLE_TABLE or JOINED this function checks
+     * recursively for base class in order to produce
+     * correct results
+     * 
+     * @param ClassMetadataInfo $meta
+     * @return string
+     */
+    private function _getBaseEntityClass(ClassMetadataInfo $meta)
+    {
+        $className = $meta->name;
+        if ($meta->discriminatorMap) {
+            // lookup for base class in discriminator map
+            $reflClass = $meta->getReflectionClass();
+            if ($reflParent = $reflClass->getParentClass()) {
+                if (in_array($reflParent->getName(), $meta->discriminatorMap)) {
+                    $meta = $this->_em->getClassMetadata($reflParent->getName());
+                    $className = $this->_getBaseEntityClass($meta);
+                }
+            }
+        }
+        return $className;
+    }
 }

tests/Gedmo/Tree/Fixture/BaseNode.php

@@ -53,6 +53,11 @@ class BaseNode extends ANode
         $this->parent = $parent;    
     }
     
+    public function getChildren()
+    {
+        return $this->children;
+    }
+    
     public function getParent()
     {
         return $this->parent;    

tests/Gedmo/Tree/MultiInheritanceTest.php

@@ -81,6 +81,29 @@ class MultiInheritanceTest extends \PHPUnit_Framework_TestCase
         $this->assertEquals(2, count($translations['en_us']));
     }
     
+    /**
+     * Test case for github issue#7
+     * Child count is invalid resulting in SINGLE_TABLE and JOINED 
+     * inheritance mapping
+     */
+    public function testCaseGithubIssue7()
+    {
+        $repo = $this->em->getRepository(self::TEST_ENTITY_CLASS);
+        $vegies = $repo->findOneByTitle('Vegitables');
+        
+        $count = $repo->childCount($vegies, true/*direct*/);
+        $this->assertEquals(3, $count);
+        
+        $children = $vegies->getChildren();
+        //$children = $repo->children($vegies, true); // fails
+        $this->assertEquals(3, count($children));
+        
+        // @todo: $repo->children will fail
+        // also. Structural changes needed. Maybe
+        // persistentCollection usage instead of query builder
+        // $repo->getPath should also fail
+    }
+    
     private function _populate()
     {
         $root = new \Tree\Fixture\Node();