From e86c6390cec4fc0f932b006b36d5b924511a5b0b Mon Sep 17 00:00:00 2001 From: Zhenjia ZHOU Date: Thu, 30 Jul 2026 00:14:29 +0800 Subject: [PATCH] fix(rust): port CJK-aware relevance-query matching to CodeCompressor (#2634) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The Rust port of `CodeCompressor` (#1154, parity-only) did not carry over the CJK-aware relevance-query matching from `headroom/transforms/code_compressor.py` (`_CONTEXT_DELIMS` / `_CJK_CHARS` / `_query_context_tokens()` / `_symbol_in_context()`, lines 2353-2387, called from lines 987/1009): - Rust tokenized the context with an ASCII-only delimiter class `[\s,;:.()\[\]{}"']+`, so a CJK query (no spaces, CJK punctuation) collapses into a single blob and never isolates an ASCII symbol name. - The substring-fallback guard `chars().count() > 3` had no CJK relaxation, so a short ASCII name glued to CJK text (e.g. `run` in `修复run函数的报错`, `db` in `请保留db相关的逻辑`) could never receive the +3.0 context boost — while Python does boost it. Same `(code, context)` input, different `symbol_scores`. This PR ports the two Python helpers with identical semantics: - `query_context_tokens()` — delimiter class extended with the CJK/full-width punctuation and ideographic space from Python's `_CONTEXT_DELIMS`; returns `(words, lowered, has_cjk)` with CJK detection over U+3000-U+9FFF, U+AC00-U+D7AF, U+FF00-U+FFEF (Python's `_CJK_CHARS`). - `symbol_in_context()` — exact token match, plus the substring fallback gated by `> 3` **characters** (Python `len()`, not bytes), relaxed when the query contains CJK. The call site in `analyze_symbols` now uses these helpers; no other behavior changed. Pure-ASCII query behavior is identical to before (exact token match, `>3`-gated substring fallback), which the tests pin down. Closes #2630 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `crates/headroom-core/src/transforms/code_compressor.rs`: extract `query_context_tokens()` / `symbol_in_context()` free functions mirroring the Python helpers (CJK/full-width delimiter class, CJK detection, CJK-relaxed `>3`-character guard); replace the inline ASCII-only tokenization + guard in `analyze_symbols` with calls to them. - Unit tests mirroring `tests/test_transforms/test_code_compressor_cjk.py` case-for-case, plus a character-vs-byte guard test and an end-to-end `compress_with` test asserting `symbol_scores`. ## Testing - [x] Unit tests pass (`cargo test -p headroom-core` — Rust-only change; Python untouched) - [x] Linting passes (`cargo fmt --check`, `cargo clippy -p headroom-core --all-targets` — no new warnings) - [ ] Type checking passes (`mypy headroom`) — N/A, no Python changes - [x] New tests added for new functionality - [x] Manual testing performed New tests (all in `code_compressor.rs` `mod tests`): - `cjk_query_isolates_wrapped_ascii_symbol` — full-width parens isolate `parse_config` - `cjk_query_matches_short_ascii_name_glued_to_cjk` — `db` (len 2) glued to CJK matches via the relaxed guard - `english_short_name_substring_still_gated` — `db` vs "keep the database helper" must NOT match (ASCII guard unchanged) - `english_exact_token_match_unchanged`, `english_long_name_substring_fallback_unchanged`, `empty_context_matches_nothing` - `guard_counts_chars_not_bytes` — the guard is a character count, matching Python `len()` - `cjk_context_boosts_named_symbol_end_to_end` — full `compress_with` run asserting `symbol_scores` (red on main, green here — see proof) ### Test Output ```text $ cargo test -p headroom-core --lib -- code_compressor::tests test transforms::code_compressor::tests::empty_and_short_passthrough ... ok test transforms::code_compressor::tests::empty_context_matches_nothing ... ok test transforms::code_compressor::tests::estimate_tokens_uses_chars_div_4_min_1 ... ok test transforms::code_compressor::tests::py_round3_matches_cpython ... ok test transforms::code_compressor::tests::py_round_int_is_half_to_even ... ok test transforms::code_compressor::tests::cjk_query_isolates_wrapped_ascii_symbol ... ok test transforms::code_compressor::tests::english_short_name_substring_still_gated ... ok test transforms::code_compressor::tests::cjk_query_matches_short_ascii_name_glued_to_cjk ... ok test transforms::code_compressor::tests::english_exact_token_match_unchanged ... ok test transforms::code_compressor::tests::english_long_name_substring_fallback_unchanged ... ok test transforms::code_compressor::tests::guard_counts_chars_not_bytes ... ok test transforms::code_compressor::tests::cjk_context_boosts_named_symbol_end_to_end ... ok test transforms::code_compressor::tests::detect_language_basic ... ok test result: ok. 13 passed; 0 failed; 0 ignored; 0 measured; 899 filtered out; finished in 0.05s $ cargo test -p headroom-core # per-binary summaries lib .......................... ok. 911 passed; 0 failed; 1 ignored auth_mode .................... ok. 16 passed; 0 failed cache_control ................ ok. 14 passed; 0 failed ccr_backends ................. ok. 7 passed; 0 failed ccr_roundtrip ................ ok. 15 passed; 0 failed code_compressor_parity ....... ok. 1 passed; 0 failed (recorded byte-parity fixtures) live_zone_ccr ................ ok. 3 passed; 0 failed live_zone_dispatch ........... ok. 6 passed; 0 failed live_zone_thresholds ......... ok. 2 passed; 0 failed live_zone_token_validation ... ok. 3 passed; 0 failed recommendations_loader ....... ok. 4 passed; 0 failed tokenizer_proptest ........... ok. 5 passed; 0 failed doc-tests .................... ok. 1 passed; 0 failed; 2 ignored $ cargo fmt --check # clean $ cargo clippy -p headroom-core --all-targets # no new warnings ``` ## Real Behavior Proof - Environment: macOS (Darwin 25.5.0), repo-pinned Rust toolchain (`rust-toolchain.toml`), branch based on current `main`. - Exact command / steps: the end-to-end test was written first and run against unmodified `main` (red), then after the fix (green). Input: Python source with two signal-symmetric functions `run` and `keep`; context `修复run函数的报错`. The Python reference gives `run` the boost (`_symbol_in_context('run', ...) == True`, `_symbol_in_context('keep', ...) == False`, verified against the live Python implementation), so expected normalized scores are `run = 1.0`, `keep = 0.0`. - Observed result: on unmodified main the end-to-end test fails (`left: 0.5, right: 1.0` — the CJK query `修复run函数的报错` gives `run` no boost, both symbols collapse to 0.5, while Python scores `run=1.0, keep=0.0`); on this branch all 8 new tests pass and the same query boosts `run` to 1.0, matching Python. Full output: Before (unmodified `main` + new test only — Rust gives no boost, both symbols collapse to 0.5): ```text ---- transforms::code_compressor::tests::cjk_context_boosts_named_symbol_end_to_end stdout ---- thread '...cjk_context_boosts_named_symbol_end_to_end' panicked at crates/headroom-core/src/transforms/code_compressor.rs:1890:9: assertion `left == right` failed: run must get the context boost left: 0.5 right: 1.0 test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 904 filtered out ``` After (this branch): the same test passes, including its ASCII control case (`fix the runner` must NOT boost `run` — scores stay 0.5/0.5), proving pure-ASCII behavior is unchanged. The recorded byte-parity fixture suite (`code_compressor_parity`) also still passes. - Not tested: real proxy traffic end-to-end (change is confined to the symbol-scoring context boost inside the Rust compressor; the Python implementation is the behavioral reference and is untouched). `kompress_parity` was not run locally — it is model-gated and my sandbox blocks the model fetch; it is unrelated to this change and CI covers its skip path. ## 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 — N/A (internal behavior fix) - [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 - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Additional Notes For background, the Python-side CJK handling comes from the merged CJK sweep work (#2220 and follow-ups); #1154 predates part of it, which is likely how the port missed it. Longer names wrapped in full-width punctuation happened to still match in Rust via the substring fallback, but the token set itself was wrong; this PR restores exact-token semantics for those too. --- .../src/transforms/code_compressor.rs | 183 ++++++++++++++++-- 1 file changed, 164 insertions(+), 19 deletions(-) diff --git a/crates/headroom-core/src/transforms/code_compressor.rs b/crates/headroom-core/src/transforms/code_compressor.rs index c06447b52..11b6a55ad 100644 --- a/crates/headroom-core/src/transforms/code_compressor.rs +++ b/crates/headroom-core/src/transforms/code_compressor.rs @@ -413,6 +413,59 @@ fn get_definition_name(node: Node, code: &str) -> Option { None } +/// Tokenize a relevance query for symbol-name matching (CJK-aware). +/// Mirrors `_query_context_tokens`: returns (word set, lowercased query, +/// has_cjk). Symbol names are ASCII identifiers; CJK relevance queries have +/// no spaces and use CJK/full-width punctuation, so an ASCII-only delimiter +/// class would collapse the whole query into one blob and never isolate an +/// ASCII name the user asked to keep. CJK/full-width punctuation and the +/// ideographic space are therefore delimiters too. +fn query_context_tokens(context: &str) -> (BTreeSet, String, bool) { + if context.is_empty() { + return (BTreeSet::new(), String::new(), false); + } + static DELIMS: std::sync::OnceLock = std::sync::OnceLock::new(); + let delims = DELIMS.get_or_init(|| { + // Same class as Python `_CONTEXT_DELIMS`. + regex::Regex::new(r#"[\s,;:.()\[\]{}"'，、；：。．！？（）【】「」『』《》〈〉·…—　]+"#) + .unwrap() + }); + static CJK: std::sync::OnceLock = std::sync::OnceLock::new(); + let cjk = CJK.get_or_init(|| { + // Same class as Python `_CJK_CHARS`: + // U+3000-U+9FFF, U+AC00-U+D7AF (Hangul), U+FF00-U+FFEF (full-width). + regex::Regex::new(r"[\u{3000}-\u{9FFF}\u{AC00}-\u{D7AF}\u{FF00}-\u{FFEF}]").unwrap() + }); + let lowered = context.to_lowercase(); + let words: BTreeSet = delims + .split(&lowered) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect(); + let has_cjk = cjk.is_match(&lowered); + (words, lowered, has_cjk) +} + +/// Whether the relevance query names this symbol. Mirrors `_symbol_in_context`: +/// exact token match, or a substring fallback gated by len>3 (in characters, +/// like Python's `len`) for ASCII queries but relaxed for CJK queries — a +/// short ASCII name glued to CJK has no delimiter to isolate it, so the +/// exact match can't fire and the guard would wrongly drop it. +fn symbol_in_context( + name_lower: &str, + words: &BTreeSet, + context_lower: &str, + has_cjk: bool, +) -> bool { + if words.is_empty() || name_lower.is_empty() { + return false; + } + if words.contains(name_lower) { + return true; + } + context_lower.contains(name_lower) && (name_lower.chars().count() > 3 || has_cjk) +} + fn is_public_symbol(name: &str, language: CodeLanguage) -> bool { if name.is_empty() { return false; @@ -1066,18 +1119,8 @@ impl CodeAwareCompressor { ref_counts.insert(qname.clone(), (count - def_count).max(0)); } - // Context words (empty when context is ""). - let context_lower = context.to_lowercase(); - let context_words: BTreeSet = if context.is_empty() { - BTreeSet::new() - } else { - static SPLIT: std::sync::OnceLock = std::sync::OnceLock::new(); - let re = SPLIT.get_or_init(|| regex::Regex::new(r#"[\s,;:.()\[\]{}"']+"#).unwrap()); - re.split(&context_lower) - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .collect() - }; + // Context words (empty when context is ""). Mirrors `_query_context_tokens`. + let (context_words, context_lower, context_has_cjk) = query_context_tokens(context); // Raw importance signals per symbol. let mut raw_signals: Vec<(String, f64)> = Vec::new(); @@ -1106,13 +1149,14 @@ impl CodeAwareCompressor { raw += 1.0; } - if !context_words.is_empty() { - let name_lower = short.to_lowercase(); - if context_words.contains(&name_lower) - || (name_lower.chars().count() > 3 && context_lower.contains(&name_lower)) - { - raw += 3.0; - } + // Context boost: the relevance query named this symbol. + if symbol_in_context( + &short.to_lowercase(), + &context_words, + &context_lower, + context_has_cjk, + ) { + raw += 3.0; } raw_signals.push((qname.clone(), raw)); } @@ -1866,6 +1910,107 @@ mod tests { assert_eq!(lang, CodeLanguage::Unknown); } + // CJK-aware relevance-query matching. Mirrors + // tests/test_transforms/test_code_compressor_cjk.py (Python reference). + + #[test] + fn cjk_query_isolates_wrapped_ascii_symbol() { + // Full-width parens around the name must still tokenize parse_config out. + let (words, lowered, has_cjk) = + query_context_tokens("请重点保留（parse_config）的解析配置"); + assert!(has_cjk); + assert!(words.contains("parse_config")); + assert!(symbol_in_context("parse_config", &words, &lowered, has_cjk)); + } + + #[test] + fn cjk_query_matches_short_ascii_name_glued_to_cjk() { + // 'db' (len 2) glued to CJK has no delimiter to isolate it; the len>3 + // guard is relaxed for CJK so the substring fallback still matches. + let (words, lowered, has_cjk) = query_context_tokens("请保留db相关的逻辑"); + assert!(has_cjk); + assert!(symbol_in_context("db", &words, &lowered, has_cjk)); + } + + #[test] + fn english_short_name_substring_still_gated() { + // ASCII query unchanged: a short name that is only a substring (not a + // token) of an English query must NOT match (avoids spurious boosts). + let (words, lowered, has_cjk) = query_context_tokens("keep the database helper"); + assert!(!has_cjk); + assert!(!symbol_in_context("db", &words, &lowered, has_cjk)); + } + + #[test] + fn english_exact_token_match_unchanged() { + let (words, lowered, has_cjk) = query_context_tokens("keep parse_config and helper"); + assert!(!has_cjk); + assert!(symbol_in_context("parse_config", &words, &lowered, has_cjk)); + assert!(symbol_in_context("helper", &words, &lowered, has_cjk)); + } + + #[test] + fn english_long_name_substring_fallback_unchanged() { + // ASCII path, len>3 substring fallback: 'parse_config' is not a + // standalone token but is a substring of 'parse_configs' -> must match. + let (words, lowered, has_cjk) = query_context_tokens("parse_configs and related helpers"); + assert!(!has_cjk); + assert!(!words.contains("parse_config")); + assert!(symbol_in_context("parse_config", &words, &lowered, has_cjk)); + } + + #[test] + fn empty_context_matches_nothing() { + let (words, lowered, has_cjk) = query_context_tokens(""); + assert!(words.is_empty()); + assert_eq!(lowered, ""); + assert!(!has_cjk); + assert!(!symbol_in_context("foo", &words, &lowered, has_cjk)); + } + + #[test] + fn guard_counts_chars_not_bytes() { + // Python's len() counts characters. A 4-char name that is >3 in chars + // must take the substring fallback on an ASCII query even though a + // byte-length comparison would agree here; conversely a 3-char name + // must not, even when it is many bytes away from any CJK. + let (words, lowered, has_cjk) = query_context_tokens("prefer the runs_fast variant"); + assert!(!has_cjk); + assert!(symbol_in_context("runs", &words, &lowered, has_cjk)); + assert!(!symbol_in_context("run", &words, &lowered, has_cjk)); + } + + /// Symmetric pair of Python functions: identical raw importance signals, + /// so any score difference comes only from the context boost. + const CJK_BOOST_CODE: &str = "import os\n\n\ +def run(config):\n value = config.get(\"alpha\")\n result = value + 1\n total = result * 2\n scaled = total - value\n merged = scaled + result\n print(merged)\n print(scaled)\n print(total)\n return merged\n\n\ +def keep(config):\n value = config.get(\"beta\")\n result = value + 2\n total = result * 3\n scaled = total - value\n merged = scaled + result\n print(merged)\n print(scaled)\n print(total)\n return merged\n"; + + fn score_of(result: &CodeCompressionResult, name: &str) -> f64 { + result + .symbol_scores + .iter() + .find(|(k, _)| k == name) + .map(|(_, v)| *v) + .unwrap_or_else(|| panic!("no score for {name}: {:?}", result.symbol_scores)) + } + + #[test] + fn cjk_context_boosts_named_symbol_end_to_end() { + // Python reference: a CJK query with no spaces still boosts the ASCII + // symbol it names ("run" glued to CJK, len 3 <= guard, has_cjk relaxes it). + let c = CodeAwareCompressor::new(CodeCompressorConfig::default()); + let r = c.compress_with(CJK_BOOST_CODE, Some("python"), "修复run函数的报错"); + assert_eq!(score_of(&r, "run"), 1.0, "run must get the context boost"); + assert_eq!(score_of(&r, "keep"), 0.0); + + // ASCII query unchanged: "run" is only a substring of "runner" and the + // len>3 guard is NOT relaxed without CJK -> no boost, symmetric scores. + let r = c.compress_with(CJK_BOOST_CODE, Some("python"), "fix the runner"); + assert_eq!(score_of(&r, "run"), 0.5); + assert_eq!(score_of(&r, "keep"), 0.5); + } + #[test] fn empty_and_short_passthrough() { let c = CodeAwareCompressor::new(CodeCompressorConfig::default());