perf(cache): reduce cache load transient memory

Author: orenlab

codeclone/cache/integrity.py

@@ -11,7 +11,8 @@
 from collections.abc import Mapping
 from pathlib import Path
 
-from ..utils.json_io import json_text as _json_text
+import orjson
+
 from ..utils.json_io import read_json_document as _read_json_document
 from ..utils.json_io import (
     write_json_document_atomically as _write_json_document_atomically,
@@ -39,12 +40,11 @@ def as_str_dict(value: object) -> dict[str, object] | None:
 
 
 def canonical_json(data: object) -> str:
-    return _json_text(data, sort_keys=True)
+    return _canonical_json_bytes(data).decode("utf-8")
 
 
 def sign_cache_payload(data: Mapping[str, object]) -> str:
-    canonical = canonical_json(data)
-    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+    return hashlib.sha256(_canonical_json_bytes(data)).hexdigest()
 
 
 def verify_cache_payload_signature(
@@ -54,6 +54,10 @@ def verify_cache_payload_signature(
     return hmac.compare_digest(signature, sign_cache_payload(payload))
 
 
+def _canonical_json_bytes(data: object) -> bytes:
+    return orjson.dumps(data, option=orjson.OPT_SORT_KEYS)
+
+
 def read_json_document(path: Path, *, max_bytes: int | None = None) -> object:
     if max_bytes is None:
         return _read_json_document(path)

codeclone/cache/store.py

@@ -406,11 +406,17 @@ def _load_and_validate(self, raw_obj: object) -> CacheData | None:
             files_dict = _as_str_dict(payload.get("files"))
             if files_dict is None:
                 return self._reject_invalid_cache_format(schema_version=version)
+            segment_projection_obj = payload.get("sr")
+            payload.pop("files", None)
+            payload.clear()
+            raw.clear()
 
         parsed_files: dict[str, CacheEntry] = {}
         with span(name="cache.decode_entries") as decode_span:
             decode_span.set_counter("cache_entries", len(files_dict))
-            for wire_path, file_entry_obj in files_dict.items():
+            while files_dict:
+                wire_path = next(iter(files_dict))
+                file_entry_obj = files_dict.pop(wire_path)
                 runtime_path = runtime_filepath_from_wire(wire_path, root=self.root)
                 parsed_entry = self._decode_entry(file_entry_obj, runtime_path)
                 if parsed_entry is None:
@@ -419,7 +425,7 @@ def _load_and_validate(self, raw_obj: object) -> CacheData | None:
             decode_span.set_counter("decoded_entries", len(parsed_files))
         with span(name="cache.segment_projection"):
             self.segment_report_projection = decode_segment_report_projection(
-                payload.get("sr"),
+                segment_projection_obj,
                 root=self.root,
             )
 

tests/test_cache.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import hashlib
 import json
 import os
 from collections.abc import Callable
@@ -81,7 +82,7 @@
     _unit_dict_from_model,
 )
 from codeclone.cache.integrity import as_str_dict as _as_str_dict
-from codeclone.cache.integrity import sign_cache_payload
+from codeclone.cache.integrity import canonical_json, sign_cache_payload
 from codeclone.cache.projection import (
     runtime_filepath_from_wire,
     wire_filepath_from_runtime,
@@ -1452,6 +1453,42 @@ def test_cache_signature_validation_ignores_json_whitespace(tmp_path: Path) -> N
     assert loaded.get_file_entry("x.py") is not None
 
 
+def test_cache_signature_matches_legacy_string_digest_for_unicode_payload() -> None:
+    cache = Cache(Path("cache.json"))
+    payload = _analysis_payload(
+        cache,
+        files={
+            "unicodé.py": {
+                "st": [1, 10],
+                "rn": ["Ω", "é"],
+            }
+        },
+    )
+    legacy_digest = hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest()
+
+    assert sign_cache_payload(payload) == legacy_digest
+
+
+def test_cache_load_accepts_legacy_string_signed_unicode_payload(
+    tmp_path: Path,
+) -> None:
+    cache_path = tmp_path / "cache.json"
+    cache = Cache(cache_path)
+    cache.put_file_entry("unicodé.py", {"mtime_ns": 1, "size": 10}, [], [], [])
+    cache.save()
+
+    raw = json.loads(cache_path.read_text("utf-8"))
+    payload = cast(dict[str, object], raw["payload"])
+    raw["sig"] = hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest()
+    cache_path.write_text(json.dumps(raw, ensure_ascii=False, indent=2), "utf-8")
+
+    loaded = Cache(cache_path)
+    loaded.load()
+
+    assert loaded.load_warning is None
+    assert loaded.get_file_entry("unicodé.py") is not None
+
+
 def test_decode_wire_file_and_name_section_helpers_cover_valid_and_invalid() -> None:
     encoded = _encode_wire_file_entry(
         {