diff --git a/src/Ting/UnitOfWork.php b/src/Ting/UnitOfWork.php
index 6ea37e8..b4f7bf4 100644
--- a/src/Ting/UnitOfWork.php
+++ b/src/Ting/UnitOfWork.php
@@ -114,7 +114,7 @@ class UnitOfWork implements PropertyListenerInterface
     {
         $hash = spl_object_hash($entity);
         if (isset($this->entitiesShouldBePersisted[$hash]) === true
-            && $this->entitiesShouldBePersisted[$hash] === self::STATE_NEW
+            && $this->entitiesShouldBePersisted[$hash]['state'] === self::STATE_NEW
         ) {
             return true;
         }
@@ -136,7 +136,7 @@ class UnitOfWork implements PropertyListenerInterface
         }
 
         $hash = spl_object_hash($entity);
-        $this->entitiesShouldBePersisted[$hash] = $state;
+        $this->entitiesShouldBePersisted[$hash] = ['state' => $state, 'entity' => $entity];
         $this->entities[$entity] = $entity;
 
         return $this;
@@ -227,7 +227,7 @@ class UnitOfWork implements PropertyListenerInterface
     public function pushDelete(NotifyPropertyInterface $entity)
     {
         $hash = spl_object_hash($entity);
-        $this->entitiesShouldBePersisted[$hash] = self::STATE_DELETE;
+        $this->entitiesShouldBePersisted[$hash] = ['state' => self::STATE_DELETE, 'entity' => $entity];
         $this->entities[$entity] = $entity;
 
         return $this;
@@ -243,7 +243,7 @@ class UnitOfWork implements PropertyListenerInterface
     {
         $hash = spl_object_hash($entity);
         if (isset($this->entitiesShouldBePersisted[$hash]) === true
-            && $this->entitiesShouldBePersisted[$hash] === self::STATE_DELETE
+            && $this->entitiesShouldBePersisted[$hash]['state'] === self::STATE_DELETE
         ) {
             return true;
         }
@@ -262,18 +262,18 @@ class UnitOfWork implements PropertyListenerInterface
      */
     public function process(): void
     {
-        foreach ($this->entitiesShouldBePersisted as $hash => $state) {
-            switch ($state) {
+        foreach ($this->entitiesShouldBePersisted as $details) {
+            switch ($details['state']) {
                 case self::STATE_MANAGED:
-                    $this->processManaged($hash);
+                    $this->processManaged($details['entity']);
                     break;
 
                 case self::STATE_NEW:
-                    $this->processNew($hash);
+                    $this->processNew($details['entity']);
                     break;
 
                 case self::STATE_DELETE:
-                    $this->processDelete($hash);
+                    $this->processDelete($details['entity']);
                     break;
             }
         }
@@ -288,13 +288,12 @@ class UnitOfWork implements PropertyListenerInterface
     /**
      * Update all applicable entities in database
      *
-     * @param string $hash
+     * @param NotifyPropertyInterface $entity
      * @throws Exception
      * @throws QueryException
      */
-    protected function processManaged($hash): void
+    protected function processManaged(NotifyPropertyInterface $entity): void
     {
-        $entity = $this->getEntity($hash);
         if (isset($this->entitiesChanged[$entity]) === false) {
             return;
         }
@@ -312,7 +311,7 @@ class UnitOfWork implements PropertyListenerInterface
 
         $this->metadataRepository->findMetadataForEntity(
             $entity,
-            function (Metadata $metadata) use ($entity, $properties, $hash) {
+            function (Metadata $metadata) use ($entity, $properties) {
                 $connection = $metadata->getConnection($this->connectionPool);
                 $query = $metadata->generateQueryForUpdate(
                     $connection,
@@ -325,7 +324,7 @@ class UnitOfWork implements PropertyListenerInterface
                 $query->prepareExecute()->execute();
 
                 $this->entitiesChanged->offsetUnset($entity);
-                unset($this->entitiesShouldBePersisted[$hash]);
+                unset($this->entitiesShouldBePersisted[spl_object_hash($entity)]);
             },
             function () use ($entity) {
                 throw new QueryException('Could not find repository matching entity "' . \get_class($entity) . '"');
@@ -339,16 +338,11 @@ class UnitOfWork implements PropertyListenerInterface
      * @throws Exception
      * @throws QueryException
      */
-    protected function processNew($hash): void
+    protected function processNew(NotifyPropertyInterface $entity): void
     {
-        $entity = $this->getEntity($hash);
-        if ($entity === null) {
-            return;
-        }
-
         $this->metadataRepository->findMetadataForEntity(
             $entity,
-            function (Metadata $metadata) use ($entity, $hash) {
+            function (Metadata $metadata) use ($entity) {
                 $connection = $metadata->getConnection($this->connectionPool);
                 $query = $metadata->generateQueryForInsert(
                     $connection,
@@ -361,7 +355,7 @@ class UnitOfWork implements PropertyListenerInterface
                 $metadata->setEntityPropertyForAutoIncrement($entity, $connection->master());
 
                 $this->entitiesChanged->offsetUnset($entity);
-                unset($this->entitiesShouldBePersisted[$hash]);
+                unset($this->entitiesShouldBePersisted[spl_object_hash($entity)]);
 
                 $this->manage($entity);
             },
@@ -374,13 +368,12 @@ class UnitOfWork implements PropertyListenerInterface
     /**
      * Delete all flagged entities from database
      *
-     * @param string $hash
+     * @param NotifyPropertyInterface $entity
      * @throws Exception
      * @throws QueryException
      */
-    protected function processDelete($hash): void
+    protected function processDelete(NotifyPropertyInterface $entity): void
     {
-        $entity = $this->getEntity($hash);
         $properties = [];
         if (isset($this->entitiesChanged[$entity])) {
             foreach ($this->entitiesChanged[$entity] as $property => $values) {
@@ -419,15 +412,4 @@ class UnitOfWork implements PropertyListenerInterface
             $this->statements[$statementName][spl_object_hash($connection)] = $connection;
         }
     }
-
-    private function getEntity(string $hash): ?NotifyPropertyInterface
-    {
-        foreach ($this->entities->getIterator() as $entity) {
-            if (spl_object_hash($entity) === $hash) {
-                return $entity;
-            }
-        }
-
-        return null;
-    }
 }