From 523697d0b663221bc91d39c93b4317d164726816 Mon Sep 17 00:00:00 2001
From: Sander Marechal <s.marechal@jejik.com>
Date: Thu, 20 Jun 2013 09:29:56 +0200
Subject: [PATCH] [DDC-1398] Extra-lazy get for indexed associations

If an association is EXTRA_LAZY and has an indexBy, then
you can call get() without loading the entire collection.
---
 lib/Doctrine/ORM/PersistentCollection.php     |  8 ++++
 .../ORM/Persisters/ManyToManyPersister.php    | 18 ++++++++
 .../ORM/Persisters/OneToManyPersister.php     | 18 ++++++++
 .../Functional/ExtraLazyCollectionTest.php    | 41 ++++++++++++++++++-
 4 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/lib/Doctrine/ORM/PersistentCollection.php b/lib/Doctrine/ORM/PersistentCollection.php
index 53ba1c1f2..d5a922aaf 100644
--- a/lib/Doctrine/ORM/PersistentCollection.php
+++ b/lib/Doctrine/ORM/PersistentCollection.php
@@ -517,6 +517,14 @@ final class PersistentCollection implements Collection, Selectable
      */
     public function get($key)
     {
+        if ( ! $this->initialized
+            && $this->association['fetch'] === Mapping\ClassMetadataInfo::FETCH_EXTRA_LAZY
+            && isset($this->association['indexBy'])
+        ) {
+            $persister = $this->em->getUnitOfWork()->getCollectionPersister($this->association);
+            return $persister->get($this, $key);
+        }
+
         $this->initialize();
 
         return $this->coll->get($key);
diff --git a/lib/Doctrine/ORM/Persisters/ManyToManyPersister.php b/lib/Doctrine/ORM/Persisters/ManyToManyPersister.php
index d9f7e30ca..4ad32e8d8 100644
--- a/lib/Doctrine/ORM/Persisters/ManyToManyPersister.php
+++ b/lib/Doctrine/ORM/Persisters/ManyToManyPersister.php
@@ -33,6 +33,24 @@ use Doctrine\ORM\UnitOfWork;
  */
 class ManyToManyPersister extends AbstractCollectionPersister
 {
+    /**
+     * {@inheritdoc}
+     *
+     * @override
+     */
+    public function get(PersistentCollection $coll, $index)
+    {
+        $mapping   = $coll->getMapping();
+        $uow       = $this->em->getUnitOfWork();
+        $persister = $uow->getEntityPersister($mapping['targetEntity']);
+
+        if (!isset($mapping['indexBy'])) {
+            throw new \BadMethodCallException("Selecting a collection by index is only supported on indexed collections.");
+        }
+
+        return current($persister->load(array($mapping['indexBy'] => $index), null, null, array(), 0, 1));
+    }
+
     /**
      * {@inheritdoc}
      *
diff --git a/lib/Doctrine/ORM/Persisters/OneToManyPersister.php b/lib/Doctrine/ORM/Persisters/OneToManyPersister.php
index 911d699f6..041dfe5e2 100644
--- a/lib/Doctrine/ORM/Persisters/OneToManyPersister.php
+++ b/lib/Doctrine/ORM/Persisters/OneToManyPersister.php
@@ -32,6 +32,24 @@ use Doctrine\ORM\UnitOfWork;
  */
 class OneToManyPersister extends AbstractCollectionPersister
 {
+    /**
+     * {@inheritdoc}
+     *
+     * @override
+     */
+    public function get(PersistentCollection $coll, $index)
+    {
+        $mapping   = $coll->getMapping();
+        $uow       = $this->em->getUnitOfWork();
+        $persister = $uow->getEntityPersister($mapping['targetEntity']);
+
+        if (!isset($mapping['indexBy'])) {
+            throw new \BadMethodCallException("Selecting a collection by index is only supported on indexed collections.");
+        }
+
+        return current($persister->load(array($mapping['indexBy'] => $index), null, null, array(), 0, 1));
+    }
+
     /**
      * Generates the SQL UPDATE that updates a particular row's foreign
      * key to null.
diff --git a/tests/Doctrine/Tests/ORM/Functional/ExtraLazyCollectionTest.php b/tests/Doctrine/Tests/ORM/Functional/ExtraLazyCollectionTest.php
index b6c7a0f1e..aa6adb13a 100644
--- a/tests/Doctrine/Tests/ORM/Functional/ExtraLazyCollectionTest.php
+++ b/tests/Doctrine/Tests/ORM/Functional/ExtraLazyCollectionTest.php
@@ -24,7 +24,9 @@ class ExtraLazyCollectionTest extends \Doctrine\Tests\OrmFunctionalTestCase
 
         $class = $this->_em->getClassMetadata('Doctrine\Tests\Models\CMS\CmsUser');
         $class->associationMappings['groups']['fetch'] = ClassMetadataInfo::FETCH_EXTRA_LAZY;
+        $class->associationMappings['groups']['indexBy'] = 'id';
         $class->associationMappings['articles']['fetch'] = ClassMetadataInfo::FETCH_EXTRA_LAZY;
+        $class->associationMappings['articles']['indexBy'] = 'id';
 
         $class = $this->_em->getClassMetadata('Doctrine\Tests\Models\CMS\CmsGroup');
         $class->associationMappings['users']['fetch'] = ClassMetadataInfo::FETCH_EXTRA_LAZY;
@@ -40,6 +42,9 @@ class ExtraLazyCollectionTest extends \Doctrine\Tests\OrmFunctionalTestCase
         $class->associationMappings['groups']['fetch'] = ClassMetadataInfo::FETCH_LAZY;
         $class->associationMappings['articles']['fetch'] = ClassMetadataInfo::FETCH_LAZY;
 
+        unset($class->associationMappings['groups']['indexBy']);
+        unset($class->associationMappings['articles']['indexBy']);
+
         $class = $this->_em->getClassMetadata('Doctrine\Tests\Models\CMS\CmsGroup');
         $class->associationMappings['users']['fetch'] = ClassMetadataInfo::FETCH_LAZY;
     }
@@ -174,8 +179,8 @@ class ExtraLazyCollectionTest extends \Doctrine\Tests\OrmFunctionalTestCase
         $this->assertEquals($queryCount + 1, $this->getCurrentQueryCount());
 
         $this->assertEquals(2, count($someGroups));
-        $this->assertTrue($user->groups->contains($someGroups[0]));
-        $this->assertTrue($user->groups->contains($someGroups[1]));
+        $this->assertTrue($user->groups->contains(array_shift($someGroups)));
+        $this->assertTrue($user->groups->contains(array_shift($someGroups)));
     }
 
     /**
@@ -512,6 +517,38 @@ class ExtraLazyCollectionTest extends \Doctrine\Tests\OrmFunctionalTestCase
         $this->assertEquals($qc + 1, $this->getCurrentQueryCount());
     }
 
+    /**
+     * @group DDC-1398
+     */
+    public function testGetIndexByOneToMany()
+    {
+        $user = $this->_em->find('Doctrine\Tests\Models\CMS\CmsUser', $this->userId);
+        /* @var $user CmsUser */
+
+        $queryCount = $this->getCurrentQueryCount();
+
+        $user->articles->get($this->articleId);
+
+        $this->assertFalse($user->articles->isInitialized());
+        $this->assertEquals($queryCount + 1, $this->getCurrentQueryCount());
+    }
+
+    /**
+     * @group DDC-1398
+     */
+    public function testGetIndexByManyToMany()
+    {
+        $user = $this->_em->find('Doctrine\Tests\Models\CMS\CmsUser', $this->userId);
+        /* @var $user CmsUser */
+
+        $queryCount = $this->getCurrentQueryCount();
+
+        $user->groups->get($this->groupId);
+
+        $this->assertFalse($user->groups->isInitialized());
+        $this->assertEquals($queryCount + 1, $this->getCurrentQueryCount());
+    }
+
     private function loadFixture()
     {
         $user1 = new \Doctrine\Tests\Models\CMS\CmsUser();