From 9b8e8ab9649c8a85709ba5ba482858b821651ef2 Mon Sep 17 00:00:00 2001
From: Valentin Udaltsov <udaltsov.valentin@gmail.com>
Date: Fri, 17 Apr 2020 22:07:41 +0300
Subject: [PATCH] Support non-empty-arrays in array_keys (#3168)

Closes #3160
---
 .../Stubs/CoreGenericFunctions.phpstub        |  9 ++++----
 tests/Template/FunctionTemplateTest.php       | 21 +++++++++++++++++--
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub b/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub
index 311156741..a8024e541 100644
--- a/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub
+++ b/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub
@@ -1,12 +1,13 @@
 <?php
 /**
  * @psalm-template T as array-key
+ * @psalm-template A as non-empty-array<T, mixed>|array<T, mixed>
  *
- * @param array<T, mixed> $arr
- * @param mixed           $search_value
- * @param bool            $strict
+ * @param A     $arr
+ * @param mixed $search_value
+ * @param bool  $strict
  *
- * @return list<T>
+ * @return (A is non-empty-array ? non-empty-list<T> : list<T>)
  * @psalm-pure
  */
 function array_keys(array $arr, $search_value = null, bool $strict = false)
diff --git a/tests/Template/FunctionTemplateTest.php b/tests/Template/FunctionTemplateTest.php
index e44a12af7..faa5e09a7 100644
--- a/tests/Template/FunctionTemplateTest.php
+++ b/tests/Template/FunctionTemplateTest.php
@@ -93,7 +93,7 @@ class FunctionTemplateTest extends TestCase
                      * @template T as array-key
                      *
                      * @param array<T, mixed> $arr
-                     * @return array<int, T>
+                     * @return list<T>
                      */
                     function my_array_keys($arr) {
                         return array_keys($arr);
@@ -101,7 +101,24 @@ class FunctionTemplateTest extends TestCase
 
                     $a = my_array_keys(["hello" => 5, "goodbye" => new \Exception()]);',
                 'assertions' => [
-                    '$a' => 'array<int, string>',
+                    '$a' => 'list<string>',
+                ],
+            ],
+            'genericNonEmptyArrayKeys' => [
+                '<?php
+                    /**
+                     * @template T as array-key
+                     *
+                     * @param non-empty-array<T, mixed> $arr
+                     * @return non-empty-list<T>
+                     */
+                    function my_array_keys($arr) {
+                        return array_keys($arr);
+                    }
+
+                    $a = my_array_keys(["hello" => 5, "goodbye" => new \Exception()]);',
+                'assertions' => [
+                    '$a' => 'non-empty-list<string>',
                 ],
             ],
             'genericArrayFlip' => [