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."""