From 17c1ed948e3a6394b7b758b7f8f1a330cedfcdb5 Mon Sep 17 00:00:00 2001
From: Benjamin Eberlei <kontakt@beberlei.de>
Date: Sat, 5 Feb 2011 09:31:40 +0100
Subject: [PATCH] [DDC-250] Initial untested support for @ManyToMany(indexBy)
 and @OneToMany(indexBy) option.

---
 .../ORM/Mapping/ClassMetadataInfo.php         | 10 +++
 .../ORM/Mapping/Driver/AnnotationDriver.php   |  2 +
 .../Mapping/Driver/DoctrineAnnotations.php    |  2 +
 lib/Doctrine/ORM/Mapping/Driver/XmlDriver.php |  8 +++
 .../ORM/Mapping/Driver/YamlDriver.php         |  8 +++
 .../ORM/Persisters/BasicEntityPersister.php   | 64 +++++++++++++------
 6 files changed, 75 insertions(+), 19 deletions(-)

diff --git a/lib/Doctrine/ORM/Mapping/ClassMetadataInfo.php b/lib/Doctrine/ORM/Mapping/ClassMetadataInfo.php
index 353f47128..1ec61b666 100644
--- a/lib/Doctrine/ORM/Mapping/ClassMetadataInfo.php
+++ b/lib/Doctrine/ORM/Mapping/ClassMetadataInfo.php
@@ -376,6 +376,11 @@ class ClassMetadataInfo
      * Only valid for many-to-many mappings. Note that one-to-many associations can be mapped
      * through a join table by simply mapping the association as many-to-many with a unique
      * constraint on the join table.
+     *
+     * - <b>indexBy</b> (string, optional, to-many only)
+     * Specification of a field on target-entity that is used to index the collection by.
+     * This field HAS to be either the primary key or a unique column. Otherwise the collection
+     * does not contain all the entities that are actually related.
      * 
      * A join table definition has the following structure:
      * <pre>
@@ -717,6 +722,11 @@ class ClassMetadataInfo
         }
         $mapping['isOwningSide'] = true; // assume owning side until we hit mappedBy
 
+        // unset optional indexBy attribute if its empty
+        if (isset($mapping['indexBy']) && !$mapping['indexBy']) {
+            unset($mapping['indexBy']);
+        }
+
         // If targetEntity is unqualified, assume it is in the same namespace as
         // the sourceEntity.
         $mapping['sourceEntity'] = $this->name;
diff --git a/lib/Doctrine/ORM/Mapping/Driver/AnnotationDriver.php b/lib/Doctrine/ORM/Mapping/Driver/AnnotationDriver.php
index 81dcc90da..9eb83c2b8 100644
--- a/lib/Doctrine/ORM/Mapping/Driver/AnnotationDriver.php
+++ b/lib/Doctrine/ORM/Mapping/Driver/AnnotationDriver.php
@@ -304,6 +304,7 @@ class AnnotationDriver implements Driver
                 $mapping['mappedBy'] = $oneToManyAnnot->mappedBy;
                 $mapping['targetEntity'] = $oneToManyAnnot->targetEntity;
                 $mapping['cascade'] = $oneToManyAnnot->cascade;
+                $mapping['indexBy'] = $oneToManyAnnot->indexBy;
                 $mapping['orphanRemoval'] = $oneToManyAnnot->orphanRemoval;
                 $mapping['fetch'] = constant('Doctrine\ORM\Mapping\ClassMetadata::FETCH_' . $oneToManyAnnot->fetch);
 
@@ -362,6 +363,7 @@ class AnnotationDriver implements Driver
                 $mapping['mappedBy'] = $manyToManyAnnot->mappedBy;
                 $mapping['inversedBy'] = $manyToManyAnnot->inversedBy;
                 $mapping['cascade'] = $manyToManyAnnot->cascade;
+                $mapping['indexBy'] = $manyToManyAnnot->indexBy;
                 $mapping['fetch'] = constant('Doctrine\ORM\Mapping\ClassMetadata::FETCH_' . $manyToManyAnnot->fetch);
 
                 if ($orderByAnnot = $this->_reader->getPropertyAnnotation($property, 'Doctrine\ORM\Mapping\OrderBy')) {
diff --git a/lib/Doctrine/ORM/Mapping/Driver/DoctrineAnnotations.php b/lib/Doctrine/ORM/Mapping/Driver/DoctrineAnnotations.php
index 435676566..ef566a083 100644
--- a/lib/Doctrine/ORM/Mapping/Driver/DoctrineAnnotations.php
+++ b/lib/Doctrine/ORM/Mapping/Driver/DoctrineAnnotations.php
@@ -80,6 +80,7 @@ final class OneToMany extends Annotation {
     public $cascade;
     public $fetch = 'LAZY';
     public $orphanRemoval = false;
+    public $indexBy;
 }
 final class ManyToOne extends Annotation {
     public $targetEntity;
@@ -93,6 +94,7 @@ final class ManyToMany extends Annotation {
     public $inversedBy;
     public $cascade;
     public $fetch = 'LAZY';
+    public $indexBy;
 }
 final class ElementCollection extends Annotation {
     public $tableName;
diff --git a/lib/Doctrine/ORM/Mapping/Driver/XmlDriver.php b/lib/Doctrine/ORM/Mapping/Driver/XmlDriver.php
index cd84849ae..3ec712b70 100644
--- a/lib/Doctrine/ORM/Mapping/Driver/XmlDriver.php
+++ b/lib/Doctrine/ORM/Mapping/Driver/XmlDriver.php
@@ -309,6 +309,10 @@ class XmlDriver extends AbstractFileDriver
                     $mapping['orderBy'] = $orderBy;
                 }
 
+                if (isset($oneToManyElement->{'index-by'})) {
+                    $mapping['indexBy'] = (string)$oneToManyElement->{'index-by'};
+                }
+
                 $metadata->mapOneToMany($mapping);
             }
         }
@@ -415,6 +419,10 @@ class XmlDriver extends AbstractFileDriver
                     $mapping['orderBy'] = $orderBy;
                 }
 
+                if (isset($manyToManyElement->{'index-by'})) {
+                    $mapping['indexBy'] = (string)$manyToManyElement->{'index-by'};
+                }
+
                 $metadata->mapManyToMany($mapping);
             }
         }
diff --git a/lib/Doctrine/ORM/Mapping/Driver/YamlDriver.php b/lib/Doctrine/ORM/Mapping/Driver/YamlDriver.php
index 0a6c6d0bd..0f88474f1 100644
--- a/lib/Doctrine/ORM/Mapping/Driver/YamlDriver.php
+++ b/lib/Doctrine/ORM/Mapping/Driver/YamlDriver.php
@@ -301,6 +301,10 @@ class YamlDriver extends AbstractFileDriver
                     $mapping['orderBy'] = $oneToManyElement['orderBy'];
                 }
 
+                if (isset($oneToManyElement['indexBy'])) {
+                    $mapping['indexBy'] = $oneToManyElement['indexBy'];
+                }
+
                 $metadata->mapOneToMany($mapping);
             }
         }
@@ -404,6 +408,10 @@ class YamlDriver extends AbstractFileDriver
                     $mapping['orderBy'] = $manyToManyElement['orderBy'];
                 }
 
+                if (isset($manyToManyElement['indexBy'])) {
+                    $mapping['indexBy'] = $manyToManyElement['indexBy'];
+                }
+
                 $metadata->mapManyToMany($mapping);
             }
         }
diff --git a/lib/Doctrine/ORM/Persisters/BasicEntityPersister.php b/lib/Doctrine/ORM/Persisters/BasicEntityPersister.php
index 242a500df..0cfd5e8cd 100644
--- a/lib/Doctrine/ORM/Persisters/BasicEntityPersister.php
+++ b/lib/Doctrine/ORM/Persisters/BasicEntityPersister.php
@@ -750,15 +750,55 @@ class BasicEntityPersister
     public function getManyToManyCollection(array $assoc, $sourceEntity, $offset = null, $limit = null)
     {
         $stmt = $this->getManyToManyStatement($assoc, $sourceEntity, $offset, $limit);
+        return $this->loadArrayFromStatement($assoc, $stmt);
+    }
 
+    /**
+     * Load an array of entities from a given dbal statement.
+     * 
+     * @param array $assoc
+     * @param Doctrine\DBAL\Statement $stmt
+     * @return array
+     */
+    private function loadArrayFromStatement($assoc, $stmt)
+    {
         $entities = array();
-        while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
-            $entities[] = $this->_createEntity($result);
+        if ($assoc['indexBy']) {
+            while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
+                $entity = $this->_createEntity($result);
+                $entities[$this->_class->reflFields[$assoc['indexBy']]->getValue($entity)] = $entity;
+            }
+        } else {
+            while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
+                $entities[] = $this->_createEntity($result);
+            }
         }
         $stmt->closeCursor();
         return $entities;
     }
 
+    /**
+     * Hydrate a collection from a given dbal statement.
+     * 
+     * @param array $assoc
+     * @param Doctrine\DBAL\Statement $stmt
+     * @param PersistentCollection $coll
+     */
+    private function loadCollectionFromStatement($assoc, $stmt, $coll)
+    {
+        if ($assoc['indexBy']) {
+            while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
+                $entity = $this->_createEntity($result);
+                $coll->hydrateSet($this->_class->reflFields[$assoc['indexBy']]->getValue($entity), $entity);
+            }
+        } else {
+            while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
+                $coll->hydrateAdd($this->_createEntity($result));
+            }
+        }
+        $stmt->closeCursor();
+    }
+
     /**
      * Loads a collection of entities of a many-to-many association.
      *
@@ -772,11 +812,7 @@ class BasicEntityPersister
     public function loadManyToManyCollection(array $assoc, $sourceEntity, PersistentCollection $coll)
     {
         $stmt = $this->getManyToManyStatement($assoc, $sourceEntity);
-
-        while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
-            $coll->hydrateAdd($this->_createEntity($result));
-        }
-        $stmt->closeCursor();
+        return $this->loadCollectionFromStatement($assoc, $stmt, $coll);
     }
 
     private function getManyToManyStatement(array $assoc, $sourceEntity, $offset = null, $limit = null)
@@ -1238,13 +1274,7 @@ class BasicEntityPersister
     public function getOneToManyCollection(array $assoc, $sourceEntity, $offset = null, $limit = null)
     {
         $stmt = $this->getOneToManyStatement($assoc, $sourceEntity, $offset, $limit);
-
-        $entities = array();
-        while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
-            $entities[] = $this->_createEntity($result);
-        }
-        $stmt->closeCursor();
-        return $entities;
+        return $this->loadArrayFromStatement($assoc, $stmt);
     }
 
     /**
@@ -1259,11 +1289,7 @@ class BasicEntityPersister
     public function loadOneToManyCollection(array $assoc, $sourceEntity, PersistentCollection $coll)
     {
         $stmt = $this->getOneToManyStatement($assoc, $sourceEntity);
-
-        while ($result = $stmt->fetch(PDO::FETCH_ASSOC)) {
-            $coll->hydrateAdd($this->_createEntity($result));
-        }
-        $stmt->closeCursor();
+        $this->loadCollectionFromStatement($assoc, $stmt, $coll);
     }
 
     /**