saf_core/
air_visitor.rs

1//! Visitor pattern for traversing AIR module structures.
2//!
3//! Provides a trait-based visitor with default no-op implementations,
4//! plus convenience `walk_module` and `walk_function` drivers.
5
6use crate::air::{AirBlock, AirFunction, AirModule, Instruction};
7
8/// Visitor trait for traversing AIR module structures.
9///
10/// Default implementations are no-ops, so visitors only need to
11/// implement the methods they care about. Return `false` from
12/// `visit_function` or `visit_block` to skip visiting children.
13pub trait AirVisitor {
14    /// Called for each function. Return `false` to skip blocks/instructions.
15    fn visit_function(&mut self, func: &AirFunction) -> bool {
16        let _ = func;
17        true
18    }
19
20    /// Called for each block. Return `false` to skip instructions.
21    fn visit_block(&mut self, func: &AirFunction, block: &AirBlock) -> bool {
22        let _ = (func, block);
23        true
24    }
25
26    /// Called for each instruction.
27    fn visit_instruction(&mut self, func: &AirFunction, block: &AirBlock, inst: &Instruction) {
28        let _ = (func, block, inst);
29    }
30}
31
32/// Walk an entire module, dispatching to the visitor.
33pub fn walk_module(module: &AirModule, visitor: &mut impl AirVisitor) {
34    for func in &module.functions {
35        if !visitor.visit_function(func) {
36            continue;
37        }
38        for block in &func.blocks {
39            if !visitor.visit_block(func, block) {
40                continue;
41            }
42            for inst in &block.instructions {
43                visitor.visit_instruction(func, block, inst);
44            }
45        }
46    }
47}
48
49/// Walk a single function.
50pub fn walk_function(func: &AirFunction, visitor: &mut impl AirVisitor) {
51    if !visitor.visit_function(func) {
52        return;
53    }
54    for block in &func.blocks {
55        if !visitor.visit_block(func, block) {
56            continue;
57        }
58        for inst in &block.instructions {
59            visitor.visit_instruction(func, block, inst);
60        }
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use crate::air::{AirBlock, AirFunction, AirModule, Instruction, Operation};
68    use crate::ids::{BlockId, FunctionId, InstId, ModuleId};
69    use std::collections::BTreeMap;
70
71    /// Helper to create a minimal test module.
72    fn test_module() -> AirModule {
73        let inst1 = Instruction::new(InstId::new(1), Operation::Ret);
74        let inst2 = Instruction::new(InstId::new(2), Operation::Ret);
75        let block = AirBlock {
76            id: BlockId::new(1),
77            label: None,
78            instructions: vec![inst1, inst2],
79        };
80        let func = AirFunction {
81            id: FunctionId::new(1),
82            name: "test_func".to_string(),
83            blocks: vec![block],
84            params: vec![],
85            entry_block: None,
86            is_declaration: false,
87            span: None,
88            symbol: None,
89            block_index: BTreeMap::new(),
90        };
91        AirModule {
92            id: ModuleId::new(1),
93            name: Some("test_mod".to_string()),
94            functions: vec![func],
95            globals: vec![],
96            source_files: vec![],
97            type_hierarchy: vec![],
98            constants: BTreeMap::new(),
99            types: BTreeMap::new(),
100            target_pointer_width: 8,
101            function_index: BTreeMap::new(),
102            name_index: BTreeMap::new(),
103        }
104    }
105
106    struct InstructionCounter {
107        count: usize,
108    }
109
110    impl AirVisitor for InstructionCounter {
111        fn visit_instruction(&mut self, _: &AirFunction, _: &AirBlock, _: &Instruction) {
112            self.count += 1;
113        }
114    }
115
116    #[test]
117    fn walk_module_counts_all_instructions() {
118        let module = test_module();
119        let mut counter = InstructionCounter { count: 0 };
120        walk_module(&module, &mut counter);
121        assert_eq!(counter.count, 2);
122    }
123
124    struct FunctionSkipper;
125
126    impl AirVisitor for FunctionSkipper {
127        fn visit_function(&mut self, _: &AirFunction) -> bool {
128            false // Skip all functions
129        }
130
131        fn visit_instruction(&mut self, _: &AirFunction, _: &AirBlock, _: &Instruction) {
132            panic!("Should not visit instructions when function is skipped");
133        }
134    }
135
136    #[test]
137    fn walk_module_skips_when_visit_function_returns_false() {
138        let module = test_module();
139        let mut skipper = FunctionSkipper;
140        walk_module(&module, &mut skipper);
141        // No panic means instructions were skipped correctly
142    }
143
144    struct BlockSkipper;
145
146    impl AirVisitor for BlockSkipper {
147        fn visit_block(&mut self, _: &AirFunction, _: &AirBlock) -> bool {
148            false // Skip all blocks
149        }
150
151        fn visit_instruction(&mut self, _: &AirFunction, _: &AirBlock, _: &Instruction) {
152            panic!("Should not visit instructions when block is skipped");
153        }
154    }
155
156    #[test]
157    fn walk_module_skips_when_visit_block_returns_false() {
158        let module = test_module();
159        let mut skipper = BlockSkipper;
160        walk_module(&module, &mut skipper);
161    }
162
163    #[test]
164    fn walk_function_counts_instructions() {
165        let module = test_module();
166        let mut counter = InstructionCounter { count: 0 };
167        walk_function(&module.functions[0], &mut counter);
168        assert_eq!(counter.count, 2);
169    }
170}
saf_core/air_visitor.rs

saf_core/
air_visitor.rs
