From 11ee00442d6ac17af457270aedeae1c07ac577ad Mon Sep 17 00:00:00 2001
From: Brown <github@muglug.com>
Date: Wed, 26 Aug 2020 11:00:38 -0400
Subject: [PATCH] Refine iterable key after is_array check

Ref #4038
---
 .../Type/SimpleAssertionReconciler.php        | 23 ++++++++++++++++---
 tests/Template/FunctionTemplateTest.php       | 19 +++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/Psalm/Internal/Type/SimpleAssertionReconciler.php b/src/Psalm/Internal/Type/SimpleAssertionReconciler.php
index ee98dc1c5..d7771ff37 100644
--- a/src/Psalm/Internal/Type/SimpleAssertionReconciler.php
+++ b/src/Psalm/Internal/Type/SimpleAssertionReconciler.php
@@ -1499,9 +1499,9 @@ class SimpleAssertionReconciler extends \Psalm\Type\Reconciler
                 $did_remove_type = true;
             } elseif ($type instanceof Atomic\TIterable) {
                 $clone_type = clone $type;
-                if ($clone_type->type_params[0]->isMixed()) {
-                    $clone_type->type_params[0] = Type::getArrayKey();
-                }
+
+                self::refineArrayKey($clone_type->type_params[0]);
+
                 $array_types[] = new TArray($clone_type->type_params);
 
                 $did_remove_type = true;
@@ -1539,6 +1539,23 @@ class SimpleAssertionReconciler extends \Psalm\Type\Reconciler
             : Type::getEmpty();
     }
 
+    private static function refineArrayKey(Union $key_type) : void
+    {
+        foreach ($key_type->getAtomicTypes() as $key => $cat) {
+            if ($cat instanceof TTemplateParam) {
+                self::refineArrayKey($cat->as);
+                $key_type->bustCache();
+            } elseif ($cat instanceof TScalar || $cat instanceof TMixed) {
+                $key_type->removeType($key);
+                $key_type->addType(new Type\Atomic\TArrayKey());
+            } elseif (!$cat instanceof TString && !$cat instanceof TInt) {
+                // this should ideally prompt some sort of error
+                $key_type->removeType($key);
+                $key_type->addType(new Type\Atomic\TArrayKey());
+            }
+        }
+    }
+
     /**
      * @param   string[]  $suppressed_issues
      * @param   0|1|2    $failed_reconciliation
diff --git a/tests/Template/FunctionTemplateTest.php b/tests/Template/FunctionTemplateTest.php
index 5085b6587..8e39ba404 100644
--- a/tests/Template/FunctionTemplateTest.php
+++ b/tests/Template/FunctionTemplateTest.php
@@ -1383,6 +1383,25 @@ class FunctionTemplateTest extends TestCase
                         return null;
                     }'
             ],
+            'inferIterableArrayKeyAfterIsArrayCheck' => [
+                '<?php
+                    /**
+                     * @template Key
+                     * @template Element
+                     * @psalm-param iterable<Key, Element> $input
+                     * @psalm-return Iterator<Key, Element>
+                     */
+                    function to_iterator(iterable $input): Iterator
+                    {
+                        if (\is_array($input)) {
+                            return new \ArrayIterator($input);
+                        } elseif ($input instanceof Iterator) {
+                            return $input;
+                        } else {
+                            return new \IteratorIterator($input);
+                        }
+                    }'
+            ],
         ];
     }