From b1f700fc275bf1d7e9461b61a9ebfdb1fba19620 Mon Sep 17 00:00:00 2001 From: Ashish Date: Mon, 15 Jun 2026 08:24:21 -0700 Subject: [PATCH] fix(compression): convert tree-sitter byte offsets to char offsets (#892) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description tree-sitter reports node positions as byte offsets into the UTF-8 encoding, but `CodeStructureHandler` builds a character-indexed mask. Any multi-byte character (accents, emoji, CJK in docstrings/comments/strings) shifted every subsequent span, preserving the wrong characters and leaking signature bytes into bodies. Stacked on #890. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/compression/handlers/code_handler.py`: remap spans through a byte->char table before masking; pure-ASCII content (byte == char) skips the conversion. - `tests/test_compression/test_code_handler.py`: regression test with `café münü 🎉` in a comment, asserting the following signature and body are correctly aligned. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_compression/ -q 93 passed, 8 skipped ``` ## Real Behavior Proof - Environment: local macOS, Python 3.11, tree-sitter-language-pack 1.8.1, branch `fix/code-byte-char-offsets` (stacked on #890). - Exact command / steps: `pytest tests/test_compression/ -q`. - Observed result: With 9 extra UTF-8 bytes ahead of it, a function signature is exactly preserved and its body stays compressible; before, the offsets were shifted. - Not tested: End-to-end through the live proxy pipeline. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — library change. See Test Output. ## Additional Notes Stacked on #890 — review the top commit until that merges. PR 4 of 7. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Fable 5 --- headroom/compression/handlers/code_handler.py | 41 +++++++++++++++++++ tests/test_compression/test_code_handler.py | 27 ++++++++++++ 2 files changed, 68 insertions(+) diff --git a/headroom/compression/handlers/code_handler.py b/headroom/compression/handlers/code_handler.py index b25d94da3..309876316 100644 --- a/headroom/compression/handlers/code_handler.py +++ b/headroom/compression/handlers/code_handler.py @@ -477,6 +477,13 @@ class CodeStructureHandler(BaseStructureHandler): visit_node(_ts_root(tree)) + # tree-sitter spans are BYTE offsets into the UTF-8 encoding; + # the mask is indexed by CHARACTER. Any non-ASCII character + # (docstrings, comments, string literals) shifts every later + # span, so convert before masking. Skipped for pure-ASCII + # content where the offsets coincide. + spans = self._byte_spans_to_char_spans(spans, content) + # Build mask from spans mask = self._spans_to_mask(spans, len(content)) @@ -553,6 +560,40 @@ class CodeStructureHandler(BaseStructureHandler): }, ) + @staticmethod + def _byte_spans_to_char_spans(spans: list[CodeSpan], content: str) -> list[CodeSpan]: + """Convert byte-offset spans to character-offset spans. + + tree-sitter reports node positions as byte offsets in the UTF-8 + encoding. For pure-ASCII content byte == char and the spans are + returned unchanged. Otherwise a byte->char table is built once + and every span endpoint is remapped. + """ + n_bytes = len(content.encode("utf-8")) + if n_bytes == len(content): + return spans + + # byte_to_char[b] = index of the character containing byte b; + # byte_to_char[n_bytes] = len(content) so exclusive ends map. + byte_to_char = [0] * (n_bytes + 1) + byte_pos = 0 + for char_idx, ch in enumerate(content): + ch_width = len(ch.encode("utf-8")) + for b in range(byte_pos, byte_pos + ch_width): + byte_to_char[b] = char_idx + byte_pos += ch_width + byte_to_char[n_bytes] = len(content) + + return [ + CodeSpan( + start=byte_to_char[min(span.start, n_bytes)], + end=byte_to_char[min(span.end, n_bytes)], + role=span.role, + is_structural=span.is_structural, + ) + for span in spans + ] + def _spans_to_mask(self, spans: list[CodeSpan], length: int) -> list[bool]: """Convert spans to character-level mask. diff --git a/tests/test_compression/test_code_handler.py b/tests/test_compression/test_code_handler.py index 556317e6a..2cd3f0f28 100644 --- a/tests/test_compression/test_code_handler.py +++ b/tests/test_compression/test_code_handler.py @@ -139,6 +139,33 @@ class TestTreeSitterContainers: result.mask.mask[i] for i in range(start, start + len("let body_line = 5;")) ), "impl method body must be compressible" + def test_non_ascii_content_mask_alignment(self, handler): + """Byte offsets must be converted to char offsets. + + Regression: tree-sitter reports byte offsets into the UTF-8 + encoding, but the mask is char-indexed. Multi-byte characters + (here: accents + an emoji, 9 extra bytes) shifted every later + span, preserving the wrong characters. + """ + code = ( + "# café münü 🎉 comment\n" + "def target(x: int) -> int:\n" + " body_value = 9\n" + " return body_value\n" + ) + result = handler.get_mask(code, language="python") + + sig = "def target(x: int) -> int:" + start = code.index(sig) + assert all(result.mask.mask[i] for i in range(start, start + len(sig))), ( + "signature after non-ASCII content must be exactly preserved" + ) + + bstart = code.index("body_value = 9") + assert not any( + result.mask.mask[i] for i in range(bstart, bstart + len("body_value = 9")) + ), "body after non-ASCII content must stay compressible" + def test_preservation_ratio_sane_for_class_code(self, handler): """A class with substantial method bodies should NOT preserve everything — the whole point of the handler."""