From 392f3194c005ae09095ae97c5e07c6abf97982ee Mon Sep 17 00:00:00 2001
From: Matthew Brown <github@muglug.com>
Date: Wed, 18 Apr 2018 23:47:21 -0400
Subject: [PATCH] Fix #684 - use getIterator return type for IteratorAggregate
 functions

---
 .../Statements/Block/ForeachChecker.php       | 65 +++++++++++++++++++
 tests/ForeachTest.php                         | 35 ++++++++++
 2 files changed, 100 insertions(+)

diff --git a/src/Psalm/Checker/Statements/Block/ForeachChecker.php b/src/Psalm/Checker/Statements/Block/ForeachChecker.php
index a14a9dcfe..5cc03f146 100644
--- a/src/Psalm/Checker/Statements/Block/ForeachChecker.php
+++ b/src/Psalm/Checker/Statements/Block/ForeachChecker.php
@@ -202,6 +202,71 @@ class ForeachChecker
                     }
 
                     if ($codebase->classImplements(
+                        $iterator_type->value,
+                        'IteratorAggregate'
+                    ) ||
+                        (
+                            $codebase->interfaceExists($iterator_type->value)
+                            && $codebase->interfaceExtends(
+                                $iterator_type->value,
+                                'IteratorAggregate'
+                            )
+                        )
+                    ) {
+                        $iterator_method = $iterator_type->value . '::getIterator';
+                        $self_class = $iterator_type->value;
+                        $iterator_class_type = $codebase->methods->getMethodReturnType(
+                            $iterator_method,
+                            $self_class
+                        );
+
+                        if ($iterator_class_type) {
+                            $array_type = ExpressionChecker::fleshOutType(
+                                $project_checker,
+                                $iterator_class_type,
+                                $self_class,
+                                $self_class
+                            );
+
+                            foreach ($array_type->getTypes() as $array_atomic_type) {
+                                if ($array_atomic_type instanceof Type\Atomic\TArray
+                                    || $array_atomic_type instanceof Type\Atomic\ObjectLike
+                                ) {
+                                    if ($array_atomic_type instanceof Type\Atomic\ObjectLike) {
+                                        $array_atomic_type = $array_atomic_type->getGenericValueType();
+                                    }
+
+                                    $key_type_part = $array_atomic_type->type_params[0];
+                                    $value_type_part = $array_atomic_type->type_params[1];
+                                } elseif ($array_atomic_type instanceof Type\Atomic\TGenericObject) {
+                                    $type_param_count = count($array_atomic_type->type_params);
+
+                                    $value_type_part = $array_atomic_type->type_params[$type_param_count - 1];
+                                    $key_type_part = $type_param_count > 1
+                                        ? $array_atomic_type->type_params[0]
+                                        : Type::getMixed();
+                                } else {
+                                    $key_type = Type::getMixed();
+                                    $value_type = Type::getMixed();
+                                    break;
+                                }
+
+                                if (!$key_type) {
+                                    $key_type = $key_type_part;
+                                } else {
+                                    $key_type = Type::combineUnionTypes($key_type, $key_type_part);
+                                }
+
+                                if (!$value_type) {
+                                    $value_type = $value_type_part;
+                                } else {
+                                    $value_type = Type::combineUnionTypes($value_type, $value_type_part);
+                                }
+                            }
+                        } else {
+                            $value_type = Type::getMixed();
+                        }
+                    } elseif ($codebase->classImplements(
                         $iterator_type->value,
                         'Iterator'
                     ) ||
diff --git a/tests/ForeachTest.php b/tests/ForeachTest.php
index 67a1096f6..4b5d42abe 100644
--- a/tests/ForeachTest.php
+++ b/tests/ForeachTest.php
@@ -49,6 +49,41 @@ class ForeachTest extends TestCase
                         foreach ($object as $foo) {}
                     }'
             ],
+            'arrayIteratorIteration' => [
+                '<?php
+                    class Item {
+                      /**
+                       * @var string
+                       */
+                      public $prop = "var";
+                    }
+
+                    class Collection implements IteratorAggregate {
+                      /**
+                       * @var Item[]
+                       */
+                      private $items = [];
+
+                      public function add(Item $item): void
+                      {
+                          $this->items[] = $item;
+                      }
+
+                      /**
+                        * @return \ArrayIterator<mixed, Item>
+                        */
+                      public function getIterator(): \ArrayIterator
+                      {
+                          return new \ArrayIterator($this->items);
+                      }
+                    }
+
+                    $collection = new Collection();
+                    $collection->add(new Item());
+                    foreach ($collection as $item) {
+                      echo $item->prop;
+                    }'
+            ],
         ];
     }