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, 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    control_flow_graph(nodes, &function_entries, entry_label)
111}
112
113fn function_entries(ast: &AST) -> HashSet<String> {
114    ast.function_entries().clone()
115}
116
117#[cfg(test)]
118mod tests {
119    use {
120        super::*,
121        crate::astnode::{GlobalDecl, Label},
122        either::Either,
123        sbpf_common::{inst_param::Register, instruction::Instruction, opcode::Opcode},
124    };
125
126    #[test]
127    fn test_optimizer_preserves_unreachable_blocks_in_reachable_function() {
128        let mut ast = AST::new();
129        ast.add_function_entry("entrypoint".to_string());
130        ast.nodes = vec![
131            label_node("entrypoint", 0),
132            instruction_node(
133                Opcode::Ja,
134                None,
135                0,
136                Some(Either::Left("target".to_string())),
137            ),
138            instruction_node(Opcode::Mov64Imm, Some(0), 8, None),
139            label_node("target", 16),
140            instruction_node(Opcode::Exit, None, 16, None),
141        ];
142        ast.set_text_size(24);
143
144        eliminate_unreachable_functions(&mut ast);
145
146        assert_eq!(ast.nodes.len(), 5);
147        assert!(matches!(
148            &ast.nodes[1],
149            ASTNode::Instruction { instruction, offset }
150                if instruction.opcode == Opcode::Ja
151                    && instruction.off == Some(Either::Left("target".to_string()))
152                    && *offset == 0
153        ));
154        assert!(matches!(
155            &ast.nodes[2],
156            ASTNode::Instruction { instruction, offset }
157                if instruction.opcode == Opcode::Mov64Imm && *offset == 8
158        ));
159        assert!(matches!(
160            &ast.nodes[3],
161            ASTNode::Label { label, offset } if label.name == "target" && *offset == 16
162        ));
163        assert!(matches!(
164            &ast.nodes[4],
165            ASTNode::Instruction { instruction, offset }
166                if instruction.opcode == Opcode::Exit && *offset == 16
167        ));
168    }
169
170    #[test]
171    fn test_optimizer_removes_uncalled_function_only() {
172        let mut ast = AST::new();
173        ast.add_function_entry("entrypoint".to_string());
174        ast.add_function_entry("dead".to_string());
175        ast.add_function_entry("callee".to_string());
176        ast.nodes = vec![
177            label_node("entrypoint", 0),
178            call_node("callee", 0),
179            instruction_node(Opcode::Exit, None, 8, None),
180            label_node("dead", 16),
181            instruction_node(Opcode::Exit, None, 16, None),
182            label_node("callee", 24),
183            instruction_node(Opcode::Exit, None, 24, None),
184        ];
185        ast.set_text_size(32);
186
187        eliminate_unreachable_functions(&mut ast);
188
189        assert!(
190            !ast.nodes
191                .iter()
192                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
193        );
194        assert!(
195            ast.nodes
196                .iter()
197                .any(|node| matches!(node, ASTNode::Label { label, offset }
198                if label.name == "callee" && *offset == 16))
199        );
200        assert_eq!(
201            ast.nodes
202                .iter()
203                .filter(|node| matches!(node, ASTNode::Instruction { .. }))
204                .count(),
205            3
206        );
207    }
208
209    #[test]
210    fn test_optimizer_uses_declared_entry_when_it_is_not_first() {
211        let mut ast = AST::new();
212        ast.add_function_entry("helper".to_string());
213        ast.add_function_entry("dead".to_string());
214        ast.add_function_entry("entrypoint".to_string());
215        ast.nodes = vec![
216            ASTNode::GlobalDecl {
217                global_decl: GlobalDecl {
218                    entry_label: "entrypoint".to_string(),
219                    span: 0..0,
220                },
221            },
222            label_node("helper", 0),
223            instruction_node(Opcode::Exit, None, 0, None),
224            label_node("dead", 8),
225            instruction_node(Opcode::Exit, None, 8, None),
226            label_node("entrypoint", 16),
227            call_node("helper", 16),
228            instruction_node(Opcode::Exit, None, 24, None),
229        ];
230        ast.set_text_size(32);
231
232        let cfg = cfg_for_ast(&ast);
233        assert_eq!(cfg.functions()[0].name(), "entrypoint");
234        assert_eq!(cfg.functions()[0].entry_block_id(), Some(2)); // entrypoint is block 2 in source order
235
236        eliminate_unreachable_functions(&mut ast);
237
238        assert!(ast.nodes.iter().any(
239            |node| matches!(node, ASTNode::Label { label, .. } if label.name == "entrypoint")
240        ));
241        assert!(
242            ast.nodes
243                .iter()
244                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "helper"))
245        );
246        assert!(
247            !ast.nodes
248                .iter()
249                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
250        );
251    }
252
253    #[test]
254    fn test_strip_dead_nodes_and_assign_offsets_recomputes_cumulatively() {
255        let mut ast = AST::new();
256        // node indices: 0=entrypoint, 1=exit, 2=dead_a, 3=exit, 4=live, 5=exit, 6=dead_b, 7=exit, 8=target, 9=exit
257        ast.nodes = vec![
258            label_node("entrypoint", 0),
259            instruction_node(Opcode::Exit, None, 0, None),
260            label_node("dead_a", 8),
261            instruction_node(Opcode::Exit, None, 8, None),
262            label_node("live", 16),
263            instruction_node(Opcode::Exit, None, 16, None),
264            label_node("dead_b", 24),
265            instruction_node(Opcode::Exit, None, 24, None),
266            label_node("target", 32),
267            instruction_node(Opcode::Exit, None, 32, None),
268        ];
269        ast.set_text_size(40);
270
271        strip_dead_nodes(&mut ast, &HashSet::from([2usize, 3, 6, 7]));
272        assign_offsets(&mut ast);
273
274        assert!(matches!(
275            &ast.nodes[2],
276            ASTNode::Label { label, offset } if label.name == "live" && *offset == 8
277        ));
278        assert!(matches!(
279            &ast.nodes[4],
280            ASTNode::Label { label, offset } if label.name == "target" && *offset == 16
281        ));
282    }
283
284    #[test]
285    fn test_assign_offsets_recomputes_from_scratch() {
286        let mut ast = AST::new();
287        ast.nodes = vec![
288            label_node("entrypoint", 999),
289            instruction_node(Opcode::Exit, None, 999, None),
290            label_node("next", 999),
291            instruction_node(Opcode::Exit, None, 999, None),
292        ];
293
294        assign_offsets(&mut ast);
295
296        assert!(matches!(
297            &ast.nodes[0],
298            ASTNode::Label { label, offset } if label.name == "entrypoint" && *offset == 0
299        ));
300        assert!(matches!(
301            &ast.nodes[1],
302            ASTNode::Instruction { offset, .. } if *offset == 0
303        ));
304        assert!(matches!(
305            &ast.nodes[2],
306            ASTNode::Label { label, offset } if label.name == "next" && *offset == 8
307        ));
308    }
309
310    fn label_node(name: &str, offset: u64) -> ASTNode {
311        ASTNode::Label {
312            label: Label {
313                name: name.to_string(),
314                span: 0..0,
315            },
316            offset,
317        }
318    }
319
320    fn instruction_node(
321        opcode: Opcode,
322        dst: Option<u8>,
323        offset: u64,
324        off: Option<Either<String, i16>>,
325    ) -> ASTNode {
326        ASTNode::Instruction {
327            instruction: Instruction {
328                opcode,
329                dst: dst.map(|n| Register { n }),
330                src: None,
331                off,
332                imm: None,
333                span: 0..0,
334            },
335            offset,
336        }
337    }
338
339    fn call_node(target: &str, offset: u64) -> ASTNode {
340        ASTNode::Instruction {
341            instruction: Instruction {
342                opcode: Opcode::Call,
343                dst: None,
344                src: None,
345                off: None,
346                imm: Some(Either::Left(target.to_string())),
347                span: 0..0,
348            },
349            offset,
350        }
351    }
352}
sbpf_assembler/optimizer/mod.rs

sbpf_assembler/optimizer/
mod.rs
