Implement symbol definitions (#9)

src/spinasm_lsp/server.py

@@ -153,19 +153,12 @@ async def completions(
     symbol_completions = [
         lsp.CompletionItem(
             label=symbol,
-            kind=lsp.CompletionItemKind.Constant,
+            kind=lsp.CompletionItemKind.Constant
+            if symbol in parser.symtbl
+            else lsp.CompletionItemKind.Module,
             detail=_get_defined_hover(symbol, parser=parser),
         )
-        for symbol in parser.symtbl
-    ]
-
-    label_completions = [
-        lsp.CompletionItem(
-            label=label,
-            kind=lsp.CompletionItemKind.Reference,
-            detail=_get_defined_hover(label, parser=parser),
-        )
-        for label in parser.jmptbl
+        for symbol in {**parser.symtbl, **parser.jmptbl}
     ]
 
     opcode_completions = [
@@ -182,7 +175,7 @@ async def completions(
 
     return lsp.CompletionList(
         is_incomplete=False,
-        items=symbol_completions + label_completions + opcode_completions,
+        items=symbol_completions + opcode_completions,
     )
 
 
@@ -211,6 +204,26 @@ async def definition(
     )
 
 
+@server.feature(lsp.TEXT_DOCUMENT_DOCUMENT_SYMBOL)
+async def document_symbol_definitions(
+    ls: SPINAsmLanguageServer, params: lsp.DocumentSymbolParams
+) -> lsp.DocumentSymbol | None:
+    """Returns the definitions of all symbols in the document."""
+    parser = await ls.get_parser(params.text_document.uri)
+
+    return [
+        lsp.DocumentSymbol(
+            name=symbol,
+            kind=lsp.SymbolKind.Module
+            if symbol in parser.jmptbl
+            else lsp.SymbolKind.Constant,
+            range=definition,
+            selection_range=definition,
+        )
+        for symbol, definition in parser.definitions.items()
+    ]
+
+
 @server.feature(lsp.TEXT_DOCUMENT_PREPARE_RENAME)
 async def prepare_rename(ls: SPINAsmLanguageServer, params: lsp.PrepareRenameParams):
     """Called by the client to determine if renaming the symbol at the given location

tests/conftest.py

@@ -10,14 +10,22 @@
 assert TEST_PATCHES, "No test patches found in the patches directory."
 
 
-class AssignmentDict(TypedDict):
+class DefinitionDict(TypedDict):
     """A dictionary track where a symbol is referenced and defined."""
 
     symbol: str
     referenced: lsp.Position
     defined: lsp.Location
 
 
+class SymbolDefinitionDict(TypedDict):
+    """A dictionary to record definition locations for a symbol."""
+
+    symbol: str
+    range: lsp.Range
+    kind: lsp.SymbolKind
+
+
 class HoverDict(TypedDict):
     """A dictionary to record hover information for a symbol."""
 
@@ -44,8 +52,39 @@ class RenameDict(TypedDict):
     changes: list[lsp.TextEdit]
 
 
+SYMBOL_DEFINITIONS: list[SymbolDefinitionDict] = [
+    {
+        # Variable
+        "symbol": "apout",
+        "kind": lsp.SymbolKind.Constant,
+        "range": lsp.Range(
+            start=lsp.Position(line=23, character=4),
+            end=lsp.Position(line=23, character=9),
+        ),
+    },
+    {
+        # Memory
+        "symbol": "lap2a",
+        "kind": lsp.SymbolKind.Constant,
+        "range": lsp.Range(
+            start=lsp.Position(line=16, character=4),
+            end=lsp.Position(line=16, character=9),
+        ),
+    },
+    {
+        # Label
+        "symbol": "endclr",
+        "kind": lsp.SymbolKind.Module,
+        "range": lsp.Range(
+            start=lsp.Position(line=41, character=0),
+            end=lsp.Position(line=41, character=6),
+        ),
+    },
+]
+
+
 # Example assignments from the "Basic.spn" patch, for testing definition locations
-DEFINITIONS: list[AssignmentDict] = [
+DEFINITIONS: list[DefinitionDict] = [
     {
         # Variable
         "symbol": "apout",

tests/test_server.py

@@ -9,8 +9,11 @@
     PATCH_DIR,
     PREPARE_RENAMES,
     RENAMES,
+    SYMBOL_DEFINITIONS,
+    DefinitionDict,
     PrepareRenameDict,
     RenameDict,
+    SymbolDefinitionDict,
 )
 
 
@@ -32,18 +35,18 @@ async def client(request, lsp_client: LanguageClient):
 
 
 @pytest.mark.asyncio()
-@pytest.mark.parametrize("assignment", DEFINITIONS, ids=lambda x: x["symbol"])
-async def test_definition(assignment: dict, client: LanguageClient):
+@pytest.mark.parametrize("definition", DEFINITIONS, ids=lambda x: x["symbol"])
+async def test_definition(definition: DefinitionDict, client: LanguageClient):
     """Test that the definition location of different assignments is correct."""
-    uri = assignment["defined"].uri
+    uri = definition["defined"].uri
     result = await client.text_document_definition_async(
         params=lsp.DefinitionParams(
-            position=assignment["referenced"],
+            position=definition["referenced"],
             text_document=lsp.TextDocumentIdentifier(uri=uri),
         )
     )
 
-    assert result == assignment["defined"]
+    assert result == definition["defined"]
 
 
 @pytest.mark.asyncio()
@@ -216,3 +219,23 @@ async def test_rename(rename: RenameDict, client: LanguageClient):
     )
 
     assert result.changes[uri] == rename["changes"]
+
+
+@pytest.mark.parametrize("symbol", SYMBOL_DEFINITIONS, ids=lambda x: x["symbol"])
+@pytest.mark.asyncio()
+async def test_symbol_definitions(symbol: SymbolDefinitionDict, client: LanguageClient):
+    """Test that the definitions of all symbols in the document are returned."""
+    patch = PATCH_DIR / "Basic.spn"
+
+    result = await client.text_document_document_symbol_async(
+        params=lsp.DocumentSymbolParams(
+            text_document=lsp.TextDocumentIdentifier(uri=f"file:///{patch.absolute()}"),
+        )
+    )
+
+    matching = [item for item in result if item.name == symbol["symbol"].upper()]
+    assert matching, f"Symbol {symbol['symbol'].upper()} not in document symbols"
+
+    item = matching[0]
+    assert item.kind == symbol["kind"]
+    assert item.range == symbol["range"]