diff --git a/EntityDependencyIterator.inc b/EntityDependencyIterator.inc
index ad7bed3..5867c93 100644
--- a/EntityDependencyIterator.inc
+++ b/EntityDependencyIterator.inc
@@ -126,9 +126,9 @@ class EntityDependencyIterator implements RecursiveIterator {
 
       // Don't add dependencies that already were checked.
       foreach ($this->dependencies as $entity_type => $dependencies) {
-        foreach ($dependencies as $entity_id) {
+        foreach ($dependencies as $key => $entity_id) {
           if (isset($this->checked[$entity_type][$entity_id])) {
-            unset($this->dependencies[$entity_type][$entity_id]);
+            unset($this->dependencies[$entity_type][$key]);
           }
           else {
             $this->causes[$entity_type][$entity_id] = array(
@@ -230,7 +230,14 @@ class EntityDependencyIterator implements RecursiveIterator {
    * Moves the current position to the next element.
    */
   public function next() {
-    return next($this->entities);
+    do {
+      $current = next($this->entities);
+      // If reached the end of array, then exit.
+      if (!$current || is_array($current) || !isset($this->entityType)) {
+        break;
+      }
+      // We need to skip all already traversed elements.
+    } while (isset($this->traversed[$this->entityType][$current]));
   }
 
   /**