Ensure memory-efficient high-level document symbol cache #944

* Introduce SymbolBody representation (compact), from which the
    body can be reconstructed, rendering body computation lazy
  * Introduce SymbolBodyFactory for the creation of such instances,
    ensuring that instances created from the same factory share
    the same file content, allowing the in-memory and on-disk caches
    to have a significantly smaller memory footprint
  * Change type of "body" key in UnifiedSymbolRepresentation to use
    SymbolBody instead of str
  * Adjust all usages of the "body" key

src/serena/symbol.py

@@ -300,7 +300,11 @@ class LanguageServerSymbol(Symbol, ToStringMixin):
 
     @property
     def body(self) -> str | None:
-        return self.symbol_root.get("body")
+        body = self.symbol_root.get("body")
+        if body is None:
+            return None
+        else:
+            return body.get_text()
 
     def get_name_path(self) -> str:
         """
@@ -418,9 +422,10 @@ class LanguageServerSymbol(Symbol, ToStringMixin):
             result["body_location"] = {"start_line": body_start_line, "end_line": body_end_line}
 
         if include_body:
-            if self.body is None:
+            body = self.body
+            if body is None:
                 log.warning("Requested body for symbol, but it is not present. The symbol might have been loaded with include_body=False.")
-            result["body"] = self.body
+            result["body"] = body
 
         if child_inclusion_predicate is None:
             child_inclusion_predicate = lambda s: True

src/solidlsp/ls.py

@@ -84,6 +84,58 @@ class LSPFileBuffer:
         return self.contents.split("\n")
 
 
+class SymbolBody:
+    """
+    Representation of the body of a symbol, which allows the extraction of the symbol's text
+    from the lines of the file it is defined in.
+
+    Instances that share the same lines buffer are memory-efficient,
+    using only 4 integers and a reference to the lines buffer from which the text can be extracted,
+    i.e. a core representation of only about 40 bytes per body.
+    """
+
+    def __init__(self, lines: list[str], start_line: int, start_col: int, end_line: int, end_col: int) -> None:
+        self._lines = lines
+        self._start_line = start_line
+        self._start_col = start_col
+        self._end_line = end_line
+        self._end_col = end_col
+
+    def get_text(self) -> str:
+        # extract relevant lines
+        symbol_body = "\n".join(self._lines[self._start_line : self._end_line + 1])
+
+        # remove leading indentation
+        symbol_body = symbol_body[self._start_col :]
+
+        # TODO: handle end_col properly (this was never implemented)
+
+        return symbol_body
+
+
+class SymbolBodyFactory:
+    """
+    A factory for the creation of SymbolBody instances from symbols dictionaries.
+    Instances created from the same factory instance are memory-efficient, as they share
+    the same lines buffer.
+    """
+
+    def __init__(self, file_buffer: LSPFileBuffer):
+        self._lines = file_buffer.split_lines()
+
+    def create_symbol_body(self, symbol: GenericDocumentSymbol) -> SymbolBody:
+        existing_body = symbol.get("body", None)
+        if existing_body and isinstance(existing_body, SymbolBody):
+            return existing_body
+
+        assert "location" in symbol
+        start_line = symbol["location"]["range"]["start"]["line"]  # type: ignore
+        end_line = symbol["location"]["range"]["end"]["line"]  # type: ignore
+        start_col = symbol["location"]["range"]["start"]["character"]  # type: ignore
+        end_col = symbol["location"]["range"]["end"]["character"]  # type: ignore
+        return SymbolBody(self._lines, start_line, start_col, end_line, end_col)
+
+
 class DocumentSymbols:
     # IMPORTANT: Instances of this class are persisted in the high-level document symbol cache
 
@@ -206,7 +258,7 @@ class SolidLanguageServer(ABC):
     """
     RAW_DOCUMENT_SYMBOL_CACHE_FILENAME = "raw_document_symbols.pkl"
     RAW_DOCUMENT_SYMBOL_CACHE_FILENAME_LEGACY_FALLBACK = "document_symbols_cache_v23-06-25.pkl"
-    DOCUMENT_SYMBOL_CACHE_VERSION = 3
+    DOCUMENT_SYMBOL_CACHE_VERSION = 4
     DOCUMENT_SYMBOL_CACHE_FILENAME = "document_symbols.pkl"
 
     # To be overridden and extended by subclasses
@@ -1106,7 +1158,7 @@ class SolidLanguageServer(ABC):
             assert isinstance(root_symbols, list), f"Unexpected response from Language Server: {root_symbols}"
             log.debug("Received %d root symbols for %s from the language server", len(root_symbols), relative_file_path)
 
-            file_lines = file_data.split_lines()
+            body_factory = SymbolBodyFactory(file_data)
 
             def convert_to_unified_symbol(original_symbol_dict: GenericDocumentSymbol) -> ls_types.UnifiedSymbolInformation:
                 """
@@ -1137,8 +1189,7 @@ class SolidLanguageServer(ABC):
                 if "relativePath" not in location:
                     location["relativePath"] = relative_file_path  # type: ignore
 
-                if "body" not in item:
-                    item["body"] = self.retrieve_symbol_body(item, file_lines=file_lines)
+                item["body"] = self.create_symbol_body(item, factory=body_factory)
 
                 # handle missing selectionRange
                 if "selectionRange" not in item:
@@ -1475,32 +1526,17 @@ class SolidLanguageServer(ABC):
 
         return ls_types.SignatureHelp(**response)  # type: ignore
 
-    def retrieve_symbol_body(
+    def create_symbol_body(
         self,
         symbol: ls_types.UnifiedSymbolInformation | LSPTypes.SymbolInformation,
-        file_lines: list[str] | None = None,
-        file_buffer: LSPFileBuffer | None = None,
-    ) -> str:
-        """
-        Load the body of the given symbol. If the body is already contained in the symbol, just return it.
-        """
-        existing_body = symbol.get("body", None)
-        if existing_body:
-            return str(existing_body)
+        factory: SymbolBodyFactory | None = None,
+    ) -> SymbolBody:
+        if factory is None:
+            assert "relativePath" in symbol["location"]
+            with self._open_file_context(symbol["location"]["relativePath"]) as f:  # type: ignore
+                factory = SymbolBodyFactory(f)
 
-        assert "location" in symbol
-        symbol_start_line = symbol["location"]["range"]["start"]["line"]
-        symbol_end_line = symbol["location"]["range"]["end"]["line"]
-        assert "relativePath" in symbol["location"]
-        if file_lines is None:
-            with self._open_file_context(symbol["location"]["relativePath"], file_buffer) as f:  # type: ignore
-                file_lines = f.split_lines()
-        symbol_body = "\n".join(file_lines[symbol_start_line : symbol_end_line + 1])
-
-        # remove leading indentation
-        symbol_start_column = symbol["location"]["range"]["start"]["character"]  # type: ignore
-        symbol_body = symbol_body[symbol_start_column:]
-        return symbol_body
+        return factory.create_symbol_body(symbol)
 
     def request_referencing_symbols(
         self,
@@ -1549,8 +1585,12 @@ class SolidLanguageServer(ABC):
             ref_col = ref["range"]["start"]["character"]
 
             with self.open_file(ref_path) as file_data:
+                body_factory = SymbolBodyFactory(file_data)
+
                 # Get the containing symbol for this reference
-                containing_symbol = self.request_containing_symbol(ref_path, ref_line, ref_col, include_body=include_body)
+                containing_symbol = self.request_containing_symbol(
+                    ref_path, ref_line, ref_col, include_body=include_body, body_factory=body_factory
+                )
                 if containing_symbol is None:
                     # TODO: HORRIBLE HACK! I don't know how to do it better for now...
                     # THIS IS BOUND TO BREAK IN MANY CASES! IT IS ALSO SPECIFIC TO PYTHON!
@@ -1588,11 +1628,6 @@ class SolidLanguageServer(ABC):
                     )
                     name = os.path.splitext(os.path.basename(ref_path))[0]
 
-                    if include_body:
-                        body = self.retrieve_full_file_content(ref_path)
-                    else:
-                        body = ""
-
                     containing_symbol = ls_types.UnifiedSymbolInformation(
                         kind=ls_types.SymbolKind.File,
                         range=fileRange,
@@ -1600,8 +1635,11 @@ class SolidLanguageServer(ABC):
                         location=location,
                         name=name,
                         children=[],
-                        body=body,
                     )
+
+                    if include_body:
+                        containing_symbol["body"] = self.create_symbol_body(containing_symbol, factory=body_factory)
+
                 if containing_symbol is None or (not include_file_symbols and containing_symbol["kind"] == ls_types.SymbolKind.File):
                     continue
 
@@ -1648,6 +1686,7 @@ class SolidLanguageServer(ABC):
         column: int | None = None,
         strict: bool = False,
         include_body: bool = False,
+        body_factory: SymbolBodyFactory | None = None,
     ) -> ls_types.UnifiedSymbolInformation | None:
         """
         Finds the first symbol containing the position for the given file.
@@ -1747,7 +1786,7 @@ class SolidLanguageServer(ABC):
             # Return the one with the greatest starting position (i.e. the innermost container).
             containing_symbol = max(containing_symbols, key=lambda s: s["location"]["range"]["start"]["line"])
             if include_body:
-                containing_symbol["body"] = self.retrieve_symbol_body(containing_symbol)
+                containing_symbol["body"] = self.create_symbol_body(containing_symbol, factory=body_factory)
             return containing_symbol
         else:
             return None

src/solidlsp/ls_types.py

@@ -5,12 +5,16 @@ Defines wrapper objects around the types returned by LSP to ensure decoupling be
 from __future__ import annotations
 
 from enum import Enum, IntEnum
-from typing import NotRequired, Union
+from typing import TYPE_CHECKING, NotRequired, Union
 
 from typing_extensions import TypedDict
 
 from solidlsp.lsp_protocol_handler.lsp_types import DiagnosticSeverity
 
+if TYPE_CHECKING:
+    from .ls import SymbolBody
+
+
 URI = str
 DocumentUri = str
 Uint = int
@@ -187,8 +191,12 @@ class SymbolTag(IntEnum):
 
 
 class UnifiedSymbolInformation(TypedDict):
-    """Represents information about programming constructs like variables, classes,
+    """
+    Represents information about programming constructs like variables, classes,
     interfaces etc.
+
+    This is a unifying extension of `lsp_types.SymbolInformation` and `lsp_types.DocumentSymbol`,
+    with added fields for SolidLSP/Serena use.
     """
 
     deprecated: NotRequired[bool]
@@ -234,7 +242,7 @@ class UnifiedSymbolInformation(TypedDict):
     """ The range that should be selected and revealed when this symbol is being picked, e.g the name of a function.
     Must be contained by the `range`. """
 
-    body: NotRequired[str]
+    body: NotRequired["SymbolBody"]
     """ The body of the symbol. """
 
     children: list[UnifiedSymbolInformation]

test/solidlsp/bash/test_bash_basic.py

@@ -50,7 +50,7 @@ class TestBashLanguageServerBasics:
         assert greet_user_symbol is not None, "Should find greet_user function"
 
         if "body" in greet_user_symbol:
-            body = greet_user_symbol["body"]
+            body = greet_user_symbol["body"].get_text()
             assert "function greet_user()" in body, "Function body should contain function definition"
             assert "case" in body.lower(), "Function body should contain case statement"
 

test/solidlsp/dart/test_dart_basic.py

@@ -212,7 +212,8 @@ class TestDartLanguageServer:
             assert containing_symbol["name"] == "add"
             assert containing_symbol["kind"] == SymbolKind.Method
             if "body" in containing_symbol:
-                assert "add" in containing_symbol["body"] or "final result" in containing_symbol["body"]
+                body = containing_symbol["body"].get_text()
+                assert "add" in body or "final result" in body
 
     @pytest.mark.parametrize("language_server", [Language.DART], indirect=True)
     def test_request_containing_symbol_class(self, language_server: SolidLanguageServer) -> None:

test/solidlsp/elixir/test_elixir_symbol_retrieval.py

@@ -53,7 +53,7 @@ class TestElixirLanguageServerSymbols:
             assert containing_symbol["name"] == "def create_user(pid, id, name, email, roles \\\\ [])"
             assert containing_symbol["kind"] == SymbolKind.Method or containing_symbol["kind"] == SymbolKind.Function
             if "body" in containing_symbol:
-                assert "def create_user" in containing_symbol["body"]
+                assert "def create_user" in containing_symbol["body"].get_text()
 
     @pytest.mark.parametrize("language_server", [Language.ELIXIR], indirect=True)
     def test_request_containing_symbol_module(self, language_server: SolidLanguageServer) -> None:

test/solidlsp/erlang/test_erlang_symbol_retrieval.py

@@ -52,7 +52,7 @@ class TestErlangLanguageServerSymbols:
             assert "create_user" in containing_symbol["name"]
             assert containing_symbol["kind"] == SymbolKind.Method or containing_symbol["kind"] == SymbolKind.Function
             if "body" in containing_symbol:
-                assert "create_user" in containing_symbol["body"]
+                assert "create_user" in containing_symbol["body"].get_text()
 
     @pytest.mark.parametrize("language_server", [Language.ERLANG], indirect=True)
     def test_request_containing_symbol_module(self, language_server: SolidLanguageServer) -> None:

test/solidlsp/python/test_symbol_retrieval.py

@@ -34,7 +34,7 @@ class TestLanguageServerSymbols:
         assert containing_symbol["name"] == "create_user"
         assert containing_symbol["kind"] == SymbolKind.Method
         if "body" in containing_symbol:
-            assert containing_symbol["body"].strip().startswith("def create_user(self")
+            assert containing_symbol["body"].get_text().strip().startswith("def create_user(self")
 
     @pytest.mark.parametrize("language_server", [Language.PYTHON], indirect=True)
     def test_references_to_variables(self, language_server: SolidLanguageServer) -> None:
@@ -463,7 +463,6 @@ class TestLanguageServerSymbols:
             for ref in language_server.request_referencing_symbols(file_path, 4, 6, include_imports=False, include_file_symbols=True)
         ]
         assert {ref["kind"] for ref in references_to_typing} == {SymbolKind.File}
-        assert {ref["body"] for ref in references_to_typing} == {""}
 
         # now include bodies
         references_to_typing = [

test/solidlsp/ruby/test_ruby_symbol_retrieval.py

@@ -53,7 +53,7 @@ class TestRubyLanguageServerSymbols:
 
         # Verify body content if available
         if "body" in containing_symbol:
-            body = containing_symbol["body"]
+            body = containing_symbol["body"].get_text()
             assert "def create_user" in body, "Method body should contain method definition"
             assert len(body.strip()) > 0, "Method body should not be empty"
 

test/solidlsp/toml/test_toml_basic.py

@@ -137,7 +137,7 @@ class TestTomlLanguageServerBasics:
         # Check that body exists and contains expected content
         # Note: Taplo includes the section header in the body
         assert "body" in package_symbol, "'package' symbol should have body"
-        package_body = package_symbol["body"]
+        package_body = package_symbol["body"].get_text()
         assert 'name = "test_project"' in package_body, "Body should contain 'name' field"
         assert 'version = "0.1.0"' in package_body, "Body should contain 'version' field"
         assert 'edition = "2021"' in package_body, "Body should contain 'edition' field"
@@ -146,7 +146,7 @@ class TestTomlLanguageServerBasics:
         deps_symbol = next((s for s in all_symbols if s.get("name") == "dependencies"), None)
         assert deps_symbol is not None, "Should find 'dependencies' symbol"
         assert "body" in deps_symbol, "'dependencies' symbol should have body"
-        deps_body = deps_symbol["body"]
+        deps_body = deps_symbol["body"].get_text()
         assert "serde" in deps_body, "Body should contain serde dependency"
         assert "tokio" in deps_body, "Body should contain tokio dependency"
 
@@ -160,7 +160,7 @@ class TestTomlLanguageServerBasics:
         )
         assert features_symbol is not None, "Should find top-level 'features' table symbol"
         assert "body" in features_symbol, "'features' symbol should have body"
-        features_body = features_symbol["body"]
+        features_body = features_symbol["body"].get_text()
         assert "default" in features_body, "Body should contain 'default' feature"
 
     @pytest.mark.parametrize("language_server", [Language.TOML], indirect=True)

test/solidlsp/toml/test_toml_edge_cases.py

@@ -152,7 +152,7 @@ class TestTomlEdgeCases:
         assert endpoint_symbol is not None
 
         if "body" in endpoint_symbol:
-            body = endpoint_symbol["body"]
+            body = endpoint_symbol["body"].get_text()
             # Body should contain the inline table syntax
             assert "url" in body or "version" in body, f"Body should contain inline table contents, got: {body}"
 

test/solidlsp/vue/test_vue_symbol_retrieval.py

@@ -43,7 +43,7 @@ class TestVueSymbolRetrieval:
 
         # Verify the body is included if available
         if "body" in containing_symbol:
-            assert "handleDigit" in containing_symbol["body"], "Function body should contain function name"
+            assert "handleDigit" in containing_symbol["body"].get_text(), "Function body should contain function name"
 
     @pytest.mark.parametrize("language_server", [Language.VUE], indirect=True)
     def test_request_containing_symbol_computed_property(self, language_server: SolidLanguageServer) -> None:

test/solidlsp/yaml_ls/test_yaml_basic.py

@@ -124,7 +124,7 @@ class TestYAMLLanguageServerBasics:
 
         # Check that body exists and contains expected content
         assert "body" in app_symbol, "'app' symbol should have body"
-        app_body = app_symbol["body"]
+        app_body = app_symbol["body"].get_text()
         assert "app:" in app_body, "Body should start with 'app:'"
         assert "name: test-application" in app_body, "Body should contain 'name' field"
         assert "version: 1.0.0" in app_body, "Body should contain 'version' field"
@@ -135,13 +135,13 @@ class TestYAMLLanguageServerBasics:
         name_symbols = [s for s in all_symbols if s.get("name") == "name" and "body" in s]
         assert len(name_symbols) > 0, "Should find 'name' symbols with bodies"
         # At least one should contain "test-application"
-        assert any("test-application" in s["body"] for s in name_symbols), "Should find name with test-application"
+        assert any("test-application" in s["body"].get_text() for s in name_symbols), "Should find name with test-application"
 
         # Find the database symbol and check its body
         database_symbol = next((s for s in all_symbols if s.get("name") == "database"), None)
         assert database_symbol is not None, "Should find 'database' symbol"
         assert "body" in database_symbol, "'database' symbol should have body"
-        db_body = database_symbol["body"]
+        db_body = database_symbol["body"].get_text()
         assert "database:" in db_body, "Body should start with 'database:'"
         assert "host: localhost" in db_body, "Body should contain host configuration"
         assert "port: 5432" in db_body, "Body should contain port configuration"