如何用Rust更优雅地实现汇编指令窗口的窥孔优化模式匹配？

问题背景 #

基于抽象栈机模型的编译器需要实现窥孔优化，将Vec<LLStackInstruction>转换为优化后的指令序列。针对初始化/赋值操作，原流程是将右值入栈、左值地址入栈，再执行弹出和STR指令，目标是简化为直接将右值入栈并通过地址指定符存储，但当前的匹配代码嵌套层级极深，可读性差且维护困难。

当前实现代码 #

pub fn peephole_optimiser(instrs: Vec<LLStackInstruction>) -> Vec<LLStackInstruction> {
    let mut optimized_instructions: Vec<LLStackInstruction> = Vec::new();

    let mut iter = instrs.iter().peekable();

    while let Some(instruction) = iter.next() {
        let window_iter = iter.clone();
        let window: Vec<_> = [vec![instruction], window_iter.take(10).collect()].concat();
        
        if let Some(LLStackInstruction::Push(Register::FP)) = window.get(0) {
            if let Some(LLStackInstruction::Mov(Condition::None, r1, ArgType::Imm(n))) = window.get(1) {
                if let Some(LLStackInstruction::Push(r2)) = window.get(2) {
                    if let Some(LLStackInstruction::Pop(_)) = window.get(3) {
                        if let Some(LLStackInstruction::Pop(_)) = window.get(4) {
                            if let Some(LLStackInstruction::Sub(..)) = window.get(5) {
                                if let Some(LLStackInstruction::Branch(..)) = window.get(6) {
                                    if let Some(LLStackInstruction::Push(_)) = window.get(7) {
                                        if let Some(LLStackInstruction::Pop(_)) = window.get(8) {
                                            if let Some(LLStackInstruction::Pop(_)) = window.get(9) {
                                                if let Some(LLStackInstruction::Str(..)) = window.get(10) {
                                                    if r1 == r2 {
                                                        optimized_instructions.push(LLStackInstruction::Pop(Register::R4));
                                                        optimized_instructions.push(LLStackInstruction::Str(Register::R4, AddressSpecifier::ImmediateOffset(Register::FP, -n)));
                                                        iter.nth(9);
                                                        continue;
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        optimized_instructions.push(instruction.clone());
    }
    optimized_instructions
}

优化示例 #

// 紧凑初始化/赋值语句
// push {fp}
// mov r1, #88
// push {r1}
// pop {r0}
// pop {r1}
// subs r0, r1, r0
// bvs _overflow_err
// push {r0}
// pop {r3}
// pop {r4}
// str r4, [r3]

// 转换为
// pop {r4}
// str r4, [fp, #-88]

更优实现方案 #

1. 链式替代嵌套匹配 #

利用Rust的Option链式调用（and_then），把多层嵌套的if let拆成线性逻辑，大幅提升可读性：

pub fn peephole_optimiser(instrs: Vec<LLStackInstruction>) -> Vec<LLStackInstruction> {
    let mut optimized = Vec::new();
    let mut iter = instrs.iter().peekable();

    while let Some(first) = iter.next() {
        // 构建匹配窗口：当前指令+后续10条
        let mut window = vec![first];
        window.extend(iter.clone().take(10));

        let match_result = (window.get(0), window.get(1), window.get(2), window.get(3), window.get(4), window.get(5), window.get(6), window.get(7), window.get(8), window.get(9), window.get(10))
            .and_then(|(i0, i1, i2, i3, i4, i5, i6, i7, i8, i9, i10)| {
                let fp_push = matches!(i0, Some(LLStackInstruction::Push(Register::FP)));
                let mov_imm = if let Some(LLStackInstruction::Mov(Condition::None, r1, ArgType::Imm(n))) = i1 {
                    Some((r1.clone(), *n))
                } else {
                    None
                };
                let r_push = if let Some(LLStackInstruction::Push(r2)) = i2 {
                    Some(r2.clone())
                } else {
                    None
                };
                let pop1 = matches!(i3, Some(LLStackInstruction::Pop(_)));
                let pop2 = matches!(i4, Some(LLStackInstruction::Pop(_)));
                let sub = matches!(i5, Some(LLStackInstruction::Sub(..)));
                let branch = matches!(i6, Some(LLStackInstruction::Branch(..)));
                let push_r0 = matches!(i7, Some(LLStackInstruction::Push(_)));
                let pop3 = matches!(i8, Some(LLStackInstruction::Pop(_)));
                let pop4 = matches!(i9, Some(LLStackInstruction::Pop(_)));
                let str = matches!(i10, Some(LLStackInstruction::Str(..)));

                if fp_push && pop1 && pop2 && sub && branch && push_r0 && pop3 && pop4 && str {
                    mov_imm.zip(r_push).filter(|(r1, r2)| r1 == r2)
                } else {
                    None
                }
            });

        if let Some((_r, n)) = match_result {
            // 应用优化替换
            optimized.push(LLStackInstruction::Pop(Register::R4));
            optimized.push(LLStackInstruction::Str(Register::R4, AddressSpecifier::ImmediateOffset(Register::FP, -n)));
            // 跳过已匹配的后续10条指令
            iter.nth(9);
            continue;
        }

        // 无匹配则保留原指令
        optimized.push(first.clone());
    }

    optimized
}

2. 封装优化规则为结构体 #

如果后续要添加更多窥孔优化规则，可以把每个规则的匹配逻辑、条件判断、替换逻辑封装成独立结构体，提升扩展性：

// 定义优化规则 trait
trait PeepholeRule {
    // 尝试匹配窗口内的指令，返回替换后的指令序列（匹配成功时）和需要跳过的指令数量
    fn apply(&self, window: &[&LLStackInstruction]) -> Option<(Vec<LLStackInstruction>, usize)>;
}

// 针对初始化/赋值的具体规则
struct AssignOptimizationRule;

impl PeepholeRule for AssignOptimizationRule {
    fn apply(&self, window: &[&LLStackInstruction]) -> Option<(Vec<LLStackInstruction>, usize)> {
        if window.len() < 11 {
            return None;
        }

        let i0 = window[0];
        let i1 = window[1];
        let i2 = window[2];
        let i3 = window[3];
        let i4 = window[4];
        let i5 = window[5];
        let i6 = window[6];
        let i7 = window[7];
        let i8 = window[8];
        let i9 = window[9];
        let i10 = window[10];

        // 匹配模式校验
        let fp_push = matches!(i0, LLStackInstruction::Push(Register::FP));
        let (r1, n) = if let LLStackInstruction::Mov(Condition::None, r, ArgType::Imm(num)) = i1 {
            (r.clone(), num)
        } else {
            return None;
        };
        let r2 = if let LLStackInstruction::Push(r) = i2 { r.clone() } else { return None; };
        let pop1 = matches!(i3, LLStackInstruction::Pop(_));
        let pop2 = matches!(i4, LLStackInstruction::Pop(_));
        let sub = matches!(i5, LLStackInstruction::Sub(..));
        let branch = matches!(i6, LLStackInstruction::Branch(..));
        let push_r0 = matches!(i7, LLStackInstruction::Push(_));
        let pop3 = matches!(i8, LLStackInstruction::Pop(_));
        let pop4 = matches!(i9, LLStackInstruction::Pop(_));
        let str = matches!(i10, LLStackInstruction::Str(..));

        if fp_push && r1 == r2 && pop1 && pop2 && sub && branch && push_r0 && pop3 && pop4 && str {
            let replacement = vec![
                LLStackInstruction::Pop(Register::R4),
                LLStackInstruction::Str(Register::R4, AddressSpecifier::ImmediateOffset(Register::FP, -n))
            ];
            // 总共匹配11条指令，跳过后续10条
            Some((replacement, 10))
        } else {
            None
        }
    }
}

pub fn peephole_optimiser(instrs: Vec<LLStackInstruction>) -> Vec<LLStackInstruction> {
    let mut optimized = Vec::new();
    let mut iter = instrs.iter().peekable();
    let rules: Vec<Box<dyn PeepholeRule>> = vec![Box::new(AssignOptimizationRule)];

    while let Some(first) = iter.next() {
        let mut window = vec![first];
        window.extend(iter.clone().take(20)); // 预留足够大的窗口容纳所有规则的匹配长度

        let mut matched = false;
        for rule in &rules {
            if let Some((replacement, skip_count)) = rule.apply(&window) {
                optimized.extend(replacement);
                iter.nth(skip_count - 1); // 已取第一个指令，跳过剩余skip_count个元素
                matched = true;
                break;
            }
        }

        if !matched {
            optimized.push(first.clone());
        }
    }

    optimized
}

3. 多轮优化处理 #

如果优化后的指令序列可能产生新的窥孔优化机会，可以加入多轮优化逻辑，直到无优化可应用：

pub fn peephole_optimiser(mut instrs: Vec<LLStackInstruction>) -> Vec<LLStackInstruction> {
    loop {
        let optimized = single_pass_optimise(instrs);
        if optimized.len() == instrs.len() {
            // 长度不变，无更多优化空间，退出
            break optimized;
        }
        instrs = optimized;
    }
}

// 单轮优化函数（复用前面的链式匹配或规则封装实现）
fn single_pass_optimise(instrs: Vec<LLStackInstruction>) -> Vec<LLStackInstruction> {
    // ... 此处放入单轮优化代码 ...
}

方案优势 #

• 链式匹配让代码线性化，可读性和可维护性大幅提升
• 规则封装模式便于后续添加新的窥孔优化规则，扩展性强
• 多轮优化能处理优化后产生的新优化场景，优化更彻底

内容的提问来源于stack exchange，提问作者Ahmed Mahmud