From 45807e27b89676244cb0c7541b723a2f318f1f95 Mon Sep 17 00:00:00 2001
From: ging-dev <thanh1101dev@gmail.com>
Date: Tue, 26 Sep 2023 01:52:23 +0700
Subject: [PATCH] feat(LSP): Improve the autocomplete for nested array shape

---
 .../Expression/Fetch/ArrayFetchAnalyzer.php   | 11 +++++++
 tests/LanguageServer/CompletionTest.php       | 33 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/ArrayFetchAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/ArrayFetchAnalyzer.php
index 9c6d85d7639..9830012c3a6 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/ArrayFetchAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/ArrayFetchAnalyzer.php
@@ -480,6 +480,17 @@ public static function getArrayAccessTypeGivenOffset(
 
         $key_values = [];
 
+        if ($codebase->store_node_types
+            && !$context->collect_initializations
+            && !$context->collect_mutations
+        ) {
+            $codebase->analyzer->addNodeType(
+                $statements_analyzer->getFilePath(),
+                $stmt->var,
+                $array_type->getId(),
+            );
+        }
+
         if ($stmt->dim instanceof PhpParser\Node\Scalar\String_) {
             $value_type = Type::getAtomicStringFromLiteral($stmt->dim->value);
             if ($value_type instanceof TLiteralString) {
diff --git a/tests/LanguageServer/CompletionTest.php b/tests/LanguageServer/CompletionTest.php
index 290562509bb..942556f56cd 100644
--- a/tests/LanguageServer/CompletionTest.php
+++ b/tests/LanguageServer/CompletionTest.php
@@ -1380,6 +1380,39 @@ public function testCompletionOnArrayKey(): void
         $this->assertCount(2, $completion_items);
     }
 
+    public function testCompletionOnNestedArrayKey(): void
+    {
+        $codebase = $this->codebase;
+        $config = $codebase->config;
+        $config->throw_exception = false;
+
+        $this->addFile(
+            'somefile.php',
+            '<?php
+                $my_array = ["foo" => ["bar" => 1]];
+                $my_array["foo"][]
+                ',
+        );
+
+        $codebase->file_provider->openFile('somefile.php');
+        $codebase->scanFiles();
+        $this->analyzeFile('somefile.php', new Context());
+
+        $completion_data = $codebase->getCompletionDataAtPosition('somefile.php', new Position(2, 33));
+        $this->assertSame(
+            [
+                'array{bar: 1}',
+                '[',
+                92,
+            ],
+            $completion_data,
+        );
+
+        $completion_items = $codebase->getCompletionItemsForArrayKeys($completion_data[0]);
+
+        $this->assertCount(1, $completion_items);
+    }
+
     public function testTypeContextForFunctionArgument(): void
     {