Prhub

#46507 [Rust Frontend] Make Granite4 string argument scanning incremental

原始 PR 作者 reidliu41 合并时间 2026-06-26 11:54 文件变更 4 提交数 1 评论 5 代码增减 +256 / -33

执行摘要

Granite4 字符串参数流式解析增量状态机,性能提升 650x

Granite4 工具调用中,arguments 可能以 JSON 字符串形式编码(如 "{\n "data":...}")。原有的 json_str 解析器无法记忆已扫描进度,每次新 chunk 到达时都会从整个缓冲区开头重新扫描,导致长字符串参数在多 chunk 场景下性能严重退化。PR 描述中的基准数据定量说明了问题:64KB 输入耗时 107.961ms,且大小翻倍时耗时增长约 4x。

值得精读。此 PR 展示了如何通过轻量级状态机将一个原本 O(n²) 的流式解析问题降为 O(n),设计模式通用性强。特别是 take_json_stringjson_str 的重构关系,体现了增量接口与完整接口的良好复用层次。此外,review 中的设计对话也值得关注——在实现新功能时同步重构已有接口来保持一致性,是提升代码质量的好榜样。

讨论亮点

审核人 BugenZhao 在 review 中提出一个重构建议:“I'm wondering if we can now simply let the full json_str to reuse the newly introduced streaming take_json_string by constructing a temporary state.” 作者采纳并实现了该建议,使 json_str 通过临时 JsonStringScanState 调用 take_json_string,再经由 decode_json_str 解码,避免了代码重复。该讨论体现了对接口一致性和代码复用的良好设计追求。

实现拆解

  1. 增量扫描原语rust/src/parser/src/utils.rs):新增 JsonStringScanState 结构体(含 scanned_lenescape 字段),以及 take_json_string 函数。该函数接受可变状态,在扫描过程中记录已处理字节数和转义状态,遇到不完整输入时保存状态并返回 Incomplete
  2. 重构一次性解析:原有的 json_str 函数被重写为内部构造临时 JsonStringScanState 并调用 take_json_string 完成扫描,然后通过新的 decode_json_str 函数进行解码。这样既保留了已有的完整字符串解析接口,又复用了增量扫描逻辑。
  3. Granite4 解析器集成rust/src/parser/src/tool/json/granite4.rs):引入 Granite4ArgsScan 枚举,包含 Object(JsonObjectScanState)String(JsonStringScanState) 两个变体。Granite4Mode::Args 中的 json_scan 字段替换为 args_scan: Option<Granite4ArgsScan>args_event 函数在遇到字符串起始引号时创建 String 状态并存储,后续 chunk 直接复用现有的 take_json_string 扫描。
  4. 基准测试与构建配置rust/src/parser/benches/granite4.rsrust/src/parser/Cargo.toml):新增 Criterion 基准,模拟 7 字节分块的流式输入,分别测试长字符串参数和长对象参数场景。在 Cargo.toml 中注册 granite4 bench 目标。
  5. 单元测试rust/src/parser/src/utils.rsmod tests):新增 4 个测试用例(take_json_string_consumes_complete_stringtake_json_string_resumes_after_incomplete_inputtake_json_string_tracks_escape_across_chunkstake_json_string_rejects_non_string_start),验证增量扫描在正常完成、跨 chunk 恢复、转义追踪和拒绝非字符串起始等情形的正确性。
文件 模块 状态 重要度
rust/src/parser/src/utils.rs 解析器 modified 8.63
rust/src/parser/src/tool/json/granite4.rs Granite4 解析器 modified 8.01
rust/src/parser/benches/granite4.rs 基准测试 added 8.28
rust/src/parser/Cargo.toml 构建配置 modified 2.91

关键符号

take_json_string json_str decode_json_str string_args_event args_event run_stream_group bench_granite4

关键源码片段

rust/src/parser/src/utils.rs core-logic

核心改动所在:定义了增量扫描状态 `JsonStringScanState`、增量扫描函数 `take_json_string`、一次性解析 `json_str` 的重构以及新的解码函数 `decode_json_str`。这是整个 PR 的基础设施。

/// 流式 JSON 字符串词法状态:记录已经扫描的长度和转义状态。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct JsonStringScanState {
    scanned_len: usize,
    escape: bool,
}/// 从上次扫描位置继续解析一个原始的 JSON 字符串,返回字符串长度(含引号)。
/// 当输入不完整时保存状态并返回 `Incomplete`。
pub fn take_json_string(
    input: &mut Partial<&str>,
    state: &mut JsonStringScanState,
) -> ModalResult<usize> {
    let text = **input;
    if text.is_empty() {
        return incomplete();
    }
    let bytes = text.as_bytes();
    if bytes[0] != b'"' {
        return Err(json_scan_error(
            "JSON string",
            StrContextValue::CharLiteral('"'),
        ));
    }
    // 从上游保存的位置开始扫描:初始化为引号后第一个字节,后续从 saved scanned_len 继续
    let mut index = if state.scanned_len == 0 {
        1 // 跳过起始引号
    } else if state.scanned_len <= bytes.len() {
        state.scanned_len
    } else {
        return incomplete();
    };    while index < bytes.len() {
        let byte = bytes[index];
        index += 1;        if state.escape {
            state.escape = false;
            continue;
        }
        match byte {
            b'\\' => state.escape = true,
            b'"' => {
                // 遇到结束引号,消耗输入并返回整个字符串长度
                input.next_slice(index);
                return Ok(index);
            }
            _ => {}
        }
    }
    // 输入耗尽,保存当前已扫描长度(即整个输入长度)
    state.scanned_len = text.len();
    incomplete()
}/// 一次性解析一个完整的 JSON 字符串(用于非流式场景)。
pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
    let text = **input;
    let checkpoint = input.checkpoint();
    let mut state = JsonStringScanState::default();
    let len = take_json_string(input, &mut state)?;
    decode_json_str(&text[..len]).inspect_err(|_| {
        input.reset(&checkpoint);
    })
}/// 解码一个完整的 JSON 字符串字面量(仅含引号和内容)。
pub fn decode_json_str(raw: &str) -> ModalResult<String> {
    serde_json::from_str::<String>(raw).map_err(|_| {
        json_scan_error(
            "JSON string",
            StrContextValue::Description("valid JSON string"),
        )
    })
}
rust/src/parser/src/tool/json/granite4.rs core-logic

消费侧集成:引入 `Granite4ArgsScan` 枚举以区分对象和字符串扫描状态,修改 `args_event` 函数使其在遇到字符串起始时创建字符串扫描状态,并跨 chunk 复用 `take_json_string`。

/// Granite4 解析器的参数扫描状态:可以是对象扫描或字符串扫描。
#[derive(Debug, Clone, PartialEq, Eq)]
enum Granite4ArgsScan {
    Object(JsonObjectScanState),
    String(JsonStringScanState),
}/// 解析 arguments 值:根据第一个字符决定是对象还是字符串,然后持续扫描。
fn args_event(
    input: &mut JsonToolInput<'_>,
    args_scan: &mut Option<Granite4ArgsScan>,
) -> ModalResult<Granite4Event> {
    // 如果已有扫描状态,直接继续
    if let Some(scan) = args_scan {
        return match scan {
            Granite4ArgsScan::Object(scan) => {
                let len = take_json_object(input, scan)?;
                Ok(Granite4Event::ObjectArgsDelta {
                    len,
                    complete: scan.complete(),
                })
            }
            Granite4ArgsScan::String(scan) => string_args_event(input, scan),
        };
    }    // 首次遇到 arguments 值,探测实际类型
    match peek(any).parse_next(input)? {
        '{' => {
            let mut scan = JsonObjectScanState::default();
            let len = take_json_object(input, &mut scan)?;
            let complete = scan.complete();
            *args_scan = Some(Granite4ArgsScan::Object(scan));
            Ok(Granite4Event::ObjectArgsDelta { len, complete })
        }
        '"' => {
            // 字符串起始:创建字符串扫描状态并立即扫描
            *args_scan = Some(Granite4ArgsScan::String(JsonStringScanState::default()));
            let Some(Granite4ArgsScan::String(scan)) = args_scan else {
                unreachable!("Granite4 string args state just set");
            };
            string_args_event(input, scan)
        }
        _ => Err(json_scan_error(
            "Granite4 tool arguments",
            StrContextValue::Description("object or string"),
        )),
    }
}/// 处理字符串参数:使用 `take_json_string` 扫描完整字符串后再解码。
fn string_args_event(
    input: &mut JsonToolInput<'_>,
    scan: &mut JsonStringScanState,
) -> ModalResult<Granite4Event> {
    let text = **input;
    let checkpoint = input.checkpoint();
    let len = take_json_string(input, scan)?;
    let decoded = decode_json_str(&text[..len])?;
    Ok(Granite4Event::StringArgs { decoded })
}

评论区精华

json_str 复用 take_json_string 以消除重复逻辑 设计

审核人 BugenZhao 建议让 `json_str` 通过构造临时状态来复用新增的流式 `take_json_string`,因为二者扫描逻辑基本相同。

结论:作者已采纳,将 `json_str` 改造为内部调用 `take_json_string` 再加 `decode_json_str`,消除了冗余代码。 · 已解决

风险与影响

本 PR 主要引入增量状态机来优化性能,核心风险是增量状态(scanned_lenescape)在多 chunk 场景下的正确性。但已有 4 个单元测试覆盖了完整字符串、跨 chunk 恢复、转义追踪和非法起始等边界情况。此外,json_str 的重构在逻辑上等价于原有实现,且保留了相同的行为。总体风险较低,但由于是解析器核心逻辑的变动,如果未来对 JsonStringScanState 修改或在其他解析器中复用,需警惕状态同步问题。没有引入新的依赖或 API 变化。

直接影响 Granite4 工具调用处理流程:在流式场景下,长字符串参数的解析时间从 O(n²) 降到 O(n),用户不易感知具体解析细节,但能体验到更低的端到端延迟和更高的吞吐量。不影响其他模型解析器(如 Llama3、Gemma4 等),因为这些解析器各自拥有独立的 args_event 逻辑。代码变更集中在 Rust 前端 parser 模块,对 Python 侧和 C++ 内核无影响。团队内该模块的维护者需要了解和认可增量扫描的设计模式。

增量状态机正确性 流式边界逃逸字符处理

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论