Fix WorkspaceIndex not reindexing when document changes. (#12)

* Fix WorkspaceIndex not reindexing when document changes.

* Fix coding standards.

Co-authored-by: Vladislav Kosev <[email protected]>

Author: BladeMF

lib/LanguageServerWorseReflection/Workspace/WorkspaceIndex.php

@@ -20,6 +20,14 @@ class WorkspaceIndex
      * @var array<string, TextDocument>
      */
     private $byName = [];
+    /**
+     * @var array<string, TextDocument>
+     */
+    private $documents = [];
+    /**
+     * @var array<string, array<string>>
+     */
+    private $documentToNameMap = [];
 
     public function __construct(SourceCodeReflector $reflector)
     {
@@ -37,41 +45,64 @@ public function documentForName(Name $name): ?TextDocument
 
     public function index(TextDocument $textDocument): void
     {
+        $this->documents[(string)$textDocument->uri()] = $textDocument;
+        $this->updateDocument($textDocument);
+    }
+
+    private function updateDocument(TextDocument $textDocument): void
+    {
+        $newNames = [];
         foreach ($this->reflector->reflectClassesIn($textDocument) as $reflectionClass) {
-            $this->byName[$reflectionClass->name()->full()] = $textDocument;
+            $newNames[] = $reflectionClass->name()->full();
         }
-
+        
         foreach ($this->reflector->reflectFunctionsIn($textDocument) as $reflectionFunction) {
-            $this->byName[$reflectionFunction->name()->full()] = $textDocument;
+            $newNames[] = $reflectionFunction->name()->full();
+        }
+
+        $this->updateNames($textDocument, $newNames, $this->documentToNameMap[(string)$textDocument->uri()] ?? []);
+    }
+
+    private function updateNames(TextDocument $textDocument, array $newNames, array $currentNames): void
+    {
+        $namesToRemove = array_diff($currentNames, $newNames);
+        
+        foreach ($newNames as $name) {
+            $this->byName[$name] = $textDocument;
+        }
+        foreach ($namesToRemove as $name) {
+            unset($this->byName[$name]);
+        }
+
+        if (!empty($newNames)) {
+            $this->documentToNameMap[(string)$textDocument->uri()] = $newNames;
+        } else {
+            unset($this->documentToNameMap[(string)$textDocument->uri()]);
         }
     }
 
     public function update(TextDocumentUri $textDocumentUri, string $updatedText): void
     {
-        foreach ($this->byName as $className => $textDocument) {
-            if ($textDocumentUri != $textDocument->uri()) {
-                continue;
-            }
-
-            $this->byName[$className] = TextDocumentBuilder::fromTextDocument($textDocument)->text($updatedText)->build();
-            return;
+        $textDocument = $this->documents[(string)$textDocumentUri] ?? null;
+        if ($textDocument === null) {
+            throw new RuntimeException(sprintf(
+                'Could not find document "%s"',
+                $textDocumentUri->__toString()
+            ));
         }
+        $this->updateDocument(TextDocumentBuilder::fromTextDocument($textDocument)->text($updatedText)->build());
     }
 
     public function remove(TextDocumentUri $textDocumentUri): void
     {
-        foreach ($this->byName as $className => $textDocument) {
-            if ($textDocumentUri != $textDocument->uri()) {
-                continue;
-            }
-
-            unset($this->byName[$className]);
-            return;
+        $textDocument = $this->documents[(string)$textDocumentUri] ?? null;
+        if ($textDocument === null) {
+            throw new RuntimeException(sprintf(
+                'Could not find document "%s"',
+                $textDocumentUri->__toString()
+            ));
         }
-
-        throw new RuntimeException(sprintf(
-            'Could not find document "%s"',
-            $textDocumentUri->__toString()
-        ));
+        $this->updateNames($textDocument, [], $this->documentToNameMap[(string)$textDocument->uri()] ?? []);
+        unset($this->documents[(string)$textDocumentUri]);
     }
 }

tests/LanguageServerWorseReflection/Unit/Workspace/WorkspaceIndexTest.php

@@ -23,8 +23,8 @@ protected function setUp(): void
 
     public function testIndexesClassesAndReturnsTextDocuments(): void
     {
-        $document1 = TextDocumentBuilder::create('<?php namespace Test {class Foobar {}}')->build();
-        $document2 = TextDocumentBuilder::create('<?php class Barfoo {}')->build();
+        $document1 = TextDocumentBuilder::create('<?php namespace Test {class Foobar {}}')->uri('/test/Foobar')->build();
+        $document2 = TextDocumentBuilder::create('<?php class Barfoo {}')->uri('/test/Barfoo')->build();
 
         $this->index->index($document1);
         $this->index->index($document2);
@@ -36,8 +36,8 @@ public function testIndexesClassesAndReturnsTextDocuments(): void
 
     public function testIndexesFunctionsAndReturnsTextDocuments(): void
     {
-        $document1 = TextDocumentBuilder::create('<?php namespace Test {function barbar() {}}')->build();
-        $document2 = TextDocumentBuilder::create('<?php function barfoo{}')->build();
+        $document1 = TextDocumentBuilder::create('<?php namespace Test {function barbar() {}}')->uri('/Test/funcs')->build();
+        $document2 = TextDocumentBuilder::create('<?php function barfoo{}')->uri('/globals')->build();
 
         $this->index->index($document1);
         $this->index->index($document2);
@@ -59,6 +59,26 @@ public function testUpdatesExistingTextDocument(): void
         self::assertSame($updatedText, $this->index->documentForName(Name::fromString('Test\Foobar'))->__toString());
     }
 
+    public function testUpdateExistingEmptyDocument(): void
+    {
+        $document1 = TextDocumentBuilder::create('<?php')->uri('/test/foobar')->build();
+        $this->index->index($document1);
+
+        $updatedText = '<?php namespace Test {class Foobar {} function foofunc{}}';
+        $this->index->update($document1->uri(), $updatedText);
+
+        self::assertSame($updatedText, $this->index->documentForName(Name::fromString('Test\Foobar'))->__toString());
+        self::assertSame($updatedText, $this->index->documentForName(Name::fromString('Test\foofunc'))->__toString());
+
+        $updatedText = '<?php namespace Test2 {class Foobar {} function foofunc{}}';
+        $this->index->update($document1->uri(), $updatedText);
+
+        self::assertSame($updatedText, $this->index->documentForName(Name::fromString('Test2\Foobar'))->__toString());
+        self::assertSame($updatedText, $this->index->documentForName(Name::fromString('Test2\foofunc'))->__toString());
+        self::assertNull($this->index->documentForName(Name::fromString('Test\Foobar')));
+        self::assertNull($this->index->documentForName(Name::fromString('Test\foofunc')));
+    }
+
     public function testRemovesTextDocument(): void
     {
         $document1 = TextDocumentBuilder::create('<?php namespace Test {class Foobar {}}')->uri('/test/foobar')->build();