Refine iterable key after is_array check

Ref #4038

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

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);
+                        }
+                    }'
+            ],
         ];
     }