sbpf_assembler/optimizer/
mod.rs

1mod canonicalize;
2
3pub(crate) use canonicalize::{
4    canonicalize_control_flow_targets, remove_temp_control_flow_target_labels,
5};
6use {
7    crate::{ast::AST, astnode::ASTNode},
8    sbpf_analyze::remove_dead_functions,
9    sbpf_ir::{Cfg, CfgRodata, InputNode, control_flow_graph},
10    std::collections::HashSet,
11};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum CfgDumpStage {
15    BeforeDfe,
16    AfterDfe,
17}
18
19impl CfgDumpStage {
20    pub fn file_name(self) -> &'static str {
21        match self {
22            Self::BeforeDfe => "dfe-before.dot",
23            Self::AfterDfe => "dfe-after.dot",
24        }
25    }
26}
27
28/// Removes functions not reachable from the entry via `call imm`.
29pub fn eliminate_unreachable_functions(ast: &mut AST) {
30    eliminate_unreachable_functions_with_observer(ast, |_, _| {});
31}
32
33/// Removes unreachable functions and exposes the CFG before and after the pass.
34/// The observer owns any optional diagnostics or I/O, keeping the pass itself pure.
35pub fn eliminate_unreachable_functions_with_observer<F>(ast: &mut AST, mut observe: F)
36where
37    F: FnMut(CfgDumpStage, &Cfg),
38{
39    let mut cfg = cfg_for_ast(ast);
40    observe(CfgDumpStage::BeforeDfe, &cfg);
41
42    let removed_functions = remove_dead_functions(&mut cfg);
43
44    if !removed_functions.is_empty() {
45        let dead_node_ids: HashSet<usize> = removed_functions
46            .into_iter()
47            .flat_map(|f| f.node_ids)
48            .collect();
49        strip_dead_nodes(ast, &dead_node_ids);
50    }
51
52    assign_offsets(ast);
53
54    let cfg = cfg_for_ast(ast);
55    observe(CfgDumpStage::AfterDfe, &cfg);
56}
57
58/// Removes AST nodes belonging to dead functions, identified by their index in
59/// `ast.nodes`. Non-label/instruction nodes (e.g. `GlobalDecl`) are always kept.
60fn strip_dead_nodes(ast: &mut AST, dead_node_ids: &HashSet<usize>) {
61    ast.nodes = std::mem::take(&mut ast.nodes)
62        .into_iter()
63        .enumerate()
64        .filter(|(idx, node)| {
65            !matches!(node, ASTNode::Label { .. } | ASTNode::Instruction { .. })
66                || !dead_node_ids.contains(idx)
67        })
68        .map(|(_, node)| node)
69        .collect();
70}
71
72/// Recomputes byte offsets for all labels and instructions in the AST from
73/// scratch, in node order. Called after any pass that alters the node list so
74/// there is a single authoritative place that assigns offsets.
75pub fn assign_offsets(ast: &mut AST) {
76    let mut text_offset = 0u64;
77    let mut text_size = 0u64;
78    for node in &mut ast.nodes {
79        match node {
80            ASTNode::Label { offset, .. } => *offset = text_offset,
81            ASTNode::Instruction {
82                instruction,
83                offset,
84            } => {
85                *offset = text_offset;
86                let size = instruction.get_size();
87                text_offset += size;
88                text_size += size;
89            }
90            _ => {}
91        }
92    }
93    ast.set_text_size(text_size);
94}
95
96fn cfg_for_ast(ast: &AST) -> Cfg {
97    let function_entries = function_entries(ast);
98    let entry_label = ast.nodes.iter().find_map(|node| {
99        if let ASTNode::GlobalDecl { global_decl } = node {
100            Some(global_decl.entry_label.as_str())
101        } else {
102            None
103        }
104    });
105    let nodes = ast.nodes.iter().map(|node| match node {
106        ASTNode::Label { label, .. } => InputNode::Label(label.name.as_str()),
107        ASTNode::Instruction { instruction, .. } => InputNode::Instruction(instruction),
108        _ => InputNode::Other,
109    });
110    let mut cfg = control_flow_graph(nodes, &function_entries, entry_label);
111    if !ast.rodata_nodes.is_empty() {
112        let rodata = ast.rodata_nodes.iter().filter_map(|node| {
113            let ASTNode::ROData { rodata, offset } = node else {
114                return None;
115            };
116            Some(CfgRodata::new(&rodata.name, *offset, rodata.get_size()))
117        });
118        let references = ast
119            .rodata_relocations()
120            .iter()
121            .map(|relocation| (relocation.offset, relocation.target.as_str()));
122        cfg.set_rodata(rodata, references);
123    }
124    cfg
125}
126
127fn function_entries(ast: &AST) -> HashSet<String> {
128    ast.function_entries().clone()
129}
130
131#[cfg(test)]
132mod tests {
133    use {
134        super::*,
135        crate::{
136            astnode::{GlobalDecl, Label, ROData},
137            parser::Token,
138        },
139        either::Either,
140        sbpf_common::{
141            inst_param::{Number, Register},
142            instruction::Instruction,
143            opcode::Opcode,
144        },
145    };
146
147    #[test]
148    fn test_optimizer_preserves_unreachable_blocks_in_reachable_function() {
149        let mut ast = AST::new();
150        ast.add_function_entry("entrypoint".to_string());
151        ast.nodes = vec![
152            label_node("entrypoint", 0),
153            instruction_node(
154                Opcode::Ja,
155                None,
156                0,
157                Some(Either::Left("target".to_string())),
158            ),
159            instruction_node(Opcode::Mov64Imm, Some(0), 8, None),
160            label_node("target", 16),
161            instruction_node(Opcode::Exit, None, 16, None),
162        ];
163        ast.set_text_size(24);
164
165        eliminate_unreachable_functions(&mut ast);
166
167        assert_eq!(ast.nodes.len(), 5);
168        assert!(matches!(
169            &ast.nodes[1],
170            ASTNode::Instruction { instruction, offset }
171                if instruction.opcode == Opcode::Ja
172                    && instruction.off == Some(Either::Left("target".to_string()))
173                    && *offset == 0
174        ));
175        assert!(matches!(
176            &ast.nodes[2],
177            ASTNode::Instruction { instruction, offset }
178                if instruction.opcode == Opcode::Mov64Imm && *offset == 8
179        ));
180        assert!(matches!(
181            &ast.nodes[3],
182            ASTNode::Label { label, offset } if label.name == "target" && *offset == 16
183        ));
184        assert!(matches!(
185            &ast.nodes[4],
186            ASTNode::Instruction { instruction, offset }
187                if instruction.opcode == Opcode::Exit && *offset == 16
188        ));
189    }
190
191    #[test]
192    fn test_optimizer_removes_uncalled_function_only() {
193        let mut ast = AST::new();
194        ast.add_function_entry("entrypoint".to_string());
195        ast.add_function_entry("dead".to_string());
196        ast.add_function_entry("callee".to_string());
197        ast.nodes = vec![
198            label_node("entrypoint", 0),
199            call_node("callee", 0),
200            instruction_node(Opcode::Exit, None, 8, None),
201            label_node("dead", 16),
202            instruction_node(Opcode::Exit, None, 16, None),
203            label_node("callee", 24),
204            instruction_node(Opcode::Exit, None, 24, None),
205        ];
206        ast.set_text_size(32);
207
208        eliminate_unreachable_functions(&mut ast);
209
210        assert!(
211            !ast.nodes
212                .iter()
213                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
214        );
215        assert!(
216            ast.nodes
217                .iter()
218                .any(|node| matches!(node, ASTNode::Label { label, offset }
219                if label.name == "callee" && *offset == 16))
220        );
221        assert_eq!(
222            ast.nodes
223                .iter()
224                .filter(|node| matches!(node, ASTNode::Instruction { .. }))
225                .count(),
226            3
227        );
228    }
229
230    #[test]
231    fn test_optimizer_uses_declared_entry_when_it_is_not_first() {
232        let mut ast = AST::new();
233        ast.add_function_entry("helper".to_string());
234        ast.add_function_entry("dead".to_string());
235        ast.add_function_entry("entrypoint".to_string());
236        ast.nodes = vec![
237            ASTNode::GlobalDecl {
238                global_decl: GlobalDecl {
239                    entry_label: "entrypoint".to_string(),
240                    span: 0..0,
241                },
242            },
243            label_node("helper", 0),
244            instruction_node(Opcode::Exit, None, 0, None),
245            label_node("dead", 8),
246            instruction_node(Opcode::Exit, None, 8, None),
247            label_node("entrypoint", 16),
248            call_node("helper", 16),
249            instruction_node(Opcode::Exit, None, 24, None),
250        ];
251        ast.set_text_size(32);
252
253        let cfg = cfg_for_ast(&ast);
254        assert_eq!(cfg.functions()[0].name(), "entrypoint");
255        assert_eq!(cfg.functions()[0].entry_block_id(), Some(2)); // entrypoint is block 2 in source order
256
257        eliminate_unreachable_functions(&mut ast);
258
259        assert!(ast.nodes.iter().any(
260            |node| matches!(node, ASTNode::Label { label, .. } if label.name == "entrypoint")
261        ));
262        assert!(
263            ast.nodes
264                .iter()
265                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "helper"))
266        );
267        assert!(
268            !ast.nodes
269                .iter()
270                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
271        );
272    }
273
274    #[test]
275    fn test_optimizer_preserves_function_targeted_by_rodata_relocation() {
276        let mut ast = AST::new();
277        ast.add_function_entry("entrypoint".to_string());
278        ast.add_function_entry("target".to_string());
279        ast.add_function_entry("dead".to_string());
280        ast.nodes = vec![
281            ASTNode::GlobalDecl {
282                global_decl: GlobalDecl {
283                    entry_label: "entrypoint".to_string(),
284                    span: 0..0,
285                },
286            },
287            label_node("entrypoint", 0),
288            instruction_node(Opcode::Exit, None, 0, None),
289            label_node("target", 8),
290            instruction_node(Opcode::Exit, None, 8, None),
291            label_node("dead", 16),
292            instruction_node(Opcode::Exit, None, 16, None),
293        ];
294        ast.rodata_nodes.push(ASTNode::ROData {
295            rodata: ROData {
296                name: "pointer".to_string(),
297                args: vec![
298                    Token::Directive("quad".to_string(), 0..0),
299                    Token::VectorLiteral(vec![Number::Int(0)], 0..0),
300                ],
301                span: 0..0,
302            },
303            offset: 0,
304        });
305        ast.set_text_size(24);
306        ast.set_rodata_size(8);
307        ast.add_rodata_relocation(0, "target".to_string());
308
309        eliminate_unreachable_functions(&mut ast);
310
311        assert!(ast.nodes.iter().any(
312            |node| matches!(node, ASTNode::Label { label, .. } if label.name == "entrypoint")
313        ));
314        assert!(
315            ast.nodes
316                .iter()
317                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "target"))
318        );
319        assert!(
320            !ast.nodes
321                .iter()
322                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
323        );
324    }
325
326    #[test]
327    fn test_strip_dead_nodes_and_assign_offsets_recomputes_cumulatively() {
328        let mut ast = AST::new();
329        // node indices: 0=entrypoint, 1=exit, 2=dead_a, 3=exit, 4=live, 5=exit, 6=dead_b, 7=exit, 8=target, 9=exit
330        ast.nodes = vec![
331            label_node("entrypoint", 0),
332            instruction_node(Opcode::Exit, None, 0, None),
333            label_node("dead_a", 8),
334            instruction_node(Opcode::Exit, None, 8, None),
335            label_node("live", 16),
336            instruction_node(Opcode::Exit, None, 16, None),
337            label_node("dead_b", 24),
338            instruction_node(Opcode::Exit, None, 24, None),
339            label_node("target", 32),
340            instruction_node(Opcode::Exit, None, 32, None),
341        ];
342        ast.set_text_size(40);
343
344        strip_dead_nodes(&mut ast, &HashSet::from([2usize, 3, 6, 7]));
345        assign_offsets(&mut ast);
346
347        assert!(matches!(
348            &ast.nodes[2],
349            ASTNode::Label { label, offset } if label.name == "live" && *offset == 8
350        ));
351        assert!(matches!(
352            &ast.nodes[4],
353            ASTNode::Label { label, offset } if label.name == "target" && *offset == 16
354        ));
355    }
356
357    #[test]
358    fn test_assign_offsets_recomputes_from_scratch() {
359        let mut ast = AST::new();
360        ast.nodes = vec![
361            label_node("entrypoint", 999),
362            instruction_node(Opcode::Exit, None, 999, None),
363            label_node("next", 999),
364            instruction_node(Opcode::Exit, None, 999, None),
365        ];
366
367        assign_offsets(&mut ast);
368
369        assert!(matches!(
370            &ast.nodes[0],
371            ASTNode::Label { label, offset } if label.name == "entrypoint" && *offset == 0
372        ));
373        assert!(matches!(
374            &ast.nodes[1],
375            ASTNode::Instruction { offset, .. } if *offset == 0
376        ));
377        assert!(matches!(
378            &ast.nodes[2],
379            ASTNode::Label { label, offset } if label.name == "next" && *offset == 8
380        ));
381    }
382
383    fn label_node(name: &str, offset: u64) -> ASTNode {
384        ASTNode::Label {
385            label: Label {
386                name: name.to_string(),
387                span: 0..0,
388            },
389            offset,
390        }
391    }
392
393    fn instruction_node(
394        opcode: Opcode,
395        dst: Option<u8>,
396        offset: u64,
397        off: Option<Either<String, i16>>,
398    ) -> ASTNode {
399        ASTNode::Instruction {
400            instruction: Instruction {
401                opcode,
402                dst: dst.map(|n| Register { n }),
403                src: None,
404                off,
405                imm: None,
406                span: 0..0,
407            },
408            offset,
409        }
410    }
411
412    fn call_node(target: &str, offset: u64) -> ASTNode {
413        ASTNode::Instruction {
414            instruction: Instruction {
415                opcode: Opcode::Call,
416                dst: None,
417                src: None,
418                off: None,
419                imm: Some(Either::Left(target.to_string())),
420                span: 0..0,
421            },
422            offset,
423        }
424    }
425}
sbpf_assembler/optimizer/mod.rs

sbpf_assembler/optimizer/
mod.rs
