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_common::instruction::Instruction,
10    sbpf_ir::{Cfg, CfgRodata, InputNode, control_flow_graph},
11    std::collections::{HashMap, HashSet},
12};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum CfgDumpStage {
16    BeforeDfe,
17    AfterDfe,
18}
19
20impl CfgDumpStage {
21    pub fn file_name(self) -> &'static str {
22        match self {
23            Self::BeforeDfe => "dfe-before.dot",
24            Self::AfterDfe => "dfe-after.dot",
25        }
26    }
27}
28
29/// Runs every CFG optimization on a single CFG built from the AST, then writes
30/// the result back once: dead-function elimination, then the ALU32 extension
31/// fix-up.
32pub fn optimize(ast: &mut AST) {
33    optimize_with_observer(ast, |_, _| {});
34}
35
36/// Like [`optimize`], and exposes the CFG before and after dead-function
37/// elimination. The observer owns any optional diagnostics or I/O, keeping the
38/// passes themselves pure.
39pub fn optimize_with_observer<F>(ast: &mut AST, mut observe: F)
40where
41    F: FnMut(CfgDumpStage, &Cfg),
42{
43    let mut cfg = cfg_for_ast(ast);
44    observe(CfgDumpStage::BeforeDfe, &cfg);
45    remove_dead_functions(&mut cfg);
46    observe(CfgDumpStage::AfterDfe, &cfg);
47    sbpf_analyze::fix_alu32_extension(&mut cfg);
48    write_cfg_to_ast(ast, &cfg);
49}
50
51/// Removes functions not reachable from the entry via `call imm`.
52pub fn eliminate_unreachable_functions(ast: &mut AST) {
53    let mut cfg = cfg_for_ast(ast);
54    remove_dead_functions(&mut cfg);
55    write_cfg_to_ast(ast, &cfg);
56}
57
58/// Makes eBPF's ALU32 zero-extension hold under SBPF, where `add32`, `sub32` and
59/// `mul32` sign-extend. Each such instruction is kept, promoted to its 64-bit form,
60/// or zero-extended. Assumes the AST follows eBPF semantics (LLVM output).
61pub fn fix_alu32_extension(ast: &mut AST) {
62    let mut cfg = cfg_for_ast(ast);
63    sbpf_analyze::fix_alu32_extension(&mut cfg);
64    write_cfg_to_ast(ast, &cfg);
65}
66
67/// Builds the CFG for the AST's text section. Each CFG node records the index of
68/// the AST node it came from, which [`write_cfg_to_ast`] uses to map back.
69pub(crate) fn cfg_for_ast(ast: &AST) -> Cfg {
70    let function_entries = function_entries(ast);
71    let entry_label = ast.nodes.iter().find_map(|node| {
72        if let ASTNode::GlobalDecl { global_decl } = node {
73            Some(global_decl.entry_label.as_str())
74        } else {
75            None
76        }
77    });
78    let nodes = ast.nodes.iter().map(|node| match node {
79        ASTNode::Label { label, .. } => InputNode::Label(label.name.as_str()),
80        ASTNode::Instruction { instruction, .. } => InputNode::Instruction(instruction),
81        _ => InputNode::Other,
82    });
83    let mut cfg = control_flow_graph(nodes, &function_entries, entry_label);
84    if !ast.rodata_nodes.is_empty() {
85        let rodata = ast.rodata_nodes.iter().filter_map(|node| {
86            let ASTNode::ROData { rodata, offset } = node else {
87                return None;
88            };
89            Some(CfgRodata::new(&rodata.name, *offset, rodata.get_size()))
90        });
91        let references = ast
92            .rodata_relocations()
93            .iter()
94            .map(|relocation| (relocation.offset, relocation.target.as_str()));
95        cfg.set_rodata(rodata, references);
96    }
97    cfg
98}
99
100/// Writes a CFG built by [`cfg_for_ast`] back into the AST, then recomputes
101/// offsets:
102/// - labels and instructions no longer in the CFG (e.g. removed functions) are
103///   dropped;
104/// - each instruction takes its CFG node's current form, followed by any
105///   instructions a pass inserted after it;
106/// - every other node (e.g. `GlobalDecl`) is kept as is.
107pub(crate) fn write_cfg_to_ast(ast: &mut AST, cfg: &Cfg) {
108    let live: HashSet<usize> = cfg
109        .all_blocks()
110        .flat_map(|(_, block)| block.node_ids().iter().copied())
111        .collect();
112
113    let mut rewritten: HashMap<usize, Vec<Instruction>> = HashMap::new();
114    for (_, block) in cfg.all_blocks() {
115        let mut last_source = None;
116        for node in block.instructions() {
117            let Some(instruction) = node.instruction() else {
118                continue;
119            };
120            match node.source_node_id() {
121                Some(id) => {
122                    rewritten.insert(id, vec![instruction.clone()]);
123                    last_source = Some(id);
124                }
125                None => {
126                    let id = last_source
127                        .expect("inserted instructions follow a source instruction in their block");
128                    rewritten
129                        .get_mut(&id)
130                        .expect("source instruction was recorded")
131                        .push(instruction.clone());
132                }
133            }
134        }
135    }
136
137    ast.nodes = std::mem::take(&mut ast.nodes)
138        .into_iter()
139        .enumerate()
140        .flat_map(|(idx, node)| match node {
141            ASTNode::Label { .. } if !live.contains(&idx) => vec![],
142            ASTNode::Instruction { .. } => rewritten
143                .remove(&idx)
144                .unwrap_or_default()
145                .into_iter()
146                .map(|instruction| ASTNode::Instruction {
147                    instruction,
148                    offset: 0,
149                })
150                .collect(),
151            node => vec![node],
152        })
153        .collect();
154
155    assign_offsets(ast);
156}
157
158/// Recomputes byte offsets for all labels and instructions in the AST from
159/// scratch, in node order. Called after any pass that alters the node list so
160/// there is a single authoritative place that assigns offsets.
161pub fn assign_offsets(ast: &mut AST) {
162    let mut text_offset = 0u64;
163    let mut text_size = 0u64;
164    for node in &mut ast.nodes {
165        match node {
166            ASTNode::Label { offset, .. } => *offset = text_offset,
167            ASTNode::Instruction {
168                instruction,
169                offset,
170            } => {
171                *offset = text_offset;
172                let size = instruction.get_size();
173                text_offset += size;
174                text_size += size;
175            }
176            _ => {}
177        }
178    }
179    ast.set_text_size(text_size);
180}
181
182fn function_entries(ast: &AST) -> HashSet<String> {
183    ast.function_entries().clone()
184}
185
186#[cfg(test)]
187mod tests {
188    use {
189        super::*,
190        crate::{
191            SbpfArch,
192            ast::{OptimizationConfig, build_program},
193            astnode::{GlobalDecl, Label, ROData},
194            parser::Token,
195        },
196        either::Either,
197        sbpf_common::{
198            inst_param::{Number, Register},
199            instruction::Instruction,
200            opcode::Opcode,
201        },
202    };
203
204    #[test]
205    fn test_optimizer_preserves_unreachable_blocks_in_reachable_function() {
206        let mut ast = AST::new();
207        ast.add_function_entry("entrypoint".to_string());
208        ast.nodes = vec![
209            label_node("entrypoint", 0),
210            instruction_node(
211                Opcode::Ja,
212                None,
213                0,
214                Some(Either::Left("target".to_string())),
215            ),
216            instruction_node(Opcode::Mov64Imm, Some(0), 8, None),
217            label_node("target", 16),
218            instruction_node(Opcode::Exit, None, 16, None),
219        ];
220        ast.set_text_size(24);
221
222        eliminate_unreachable_functions(&mut ast);
223
224        assert_eq!(ast.nodes.len(), 5);
225        assert!(matches!(
226            &ast.nodes[1],
227            ASTNode::Instruction { instruction, offset }
228                if instruction.opcode == Opcode::Ja
229                    && instruction.off == Some(Either::Left("target".to_string()))
230                    && *offset == 0
231        ));
232        assert!(matches!(
233            &ast.nodes[2],
234            ASTNode::Instruction { instruction, offset }
235                if instruction.opcode == Opcode::Mov64Imm && *offset == 8
236        ));
237        assert!(matches!(
238            &ast.nodes[3],
239            ASTNode::Label { label, offset } if label.name == "target" && *offset == 16
240        ));
241        assert!(matches!(
242            &ast.nodes[4],
243            ASTNode::Instruction { instruction, offset }
244                if instruction.opcode == Opcode::Exit && *offset == 16
245        ));
246    }
247
248    #[test]
249    fn test_optimizer_removes_uncalled_function_only() {
250        let mut ast = AST::new();
251        ast.add_function_entry("entrypoint".to_string());
252        ast.add_function_entry("dead".to_string());
253        ast.add_function_entry("callee".to_string());
254        ast.nodes = vec![
255            label_node("entrypoint", 0),
256            call_node("callee", 0),
257            instruction_node(Opcode::Exit, None, 8, None),
258            label_node("dead", 16),
259            instruction_node(Opcode::Exit, None, 16, None),
260            label_node("callee", 24),
261            instruction_node(Opcode::Exit, None, 24, None),
262        ];
263        ast.set_text_size(32);
264
265        eliminate_unreachable_functions(&mut ast);
266
267        assert!(
268            !ast.nodes
269                .iter()
270                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
271        );
272        assert!(
273            ast.nodes
274                .iter()
275                .any(|node| matches!(node, ASTNode::Label { label, offset }
276                if label.name == "callee" && *offset == 16))
277        );
278        assert_eq!(
279            ast.nodes
280                .iter()
281                .filter(|node| matches!(node, ASTNode::Instruction { .. }))
282                .count(),
283            3
284        );
285    }
286
287    #[test]
288    fn test_optimizer_uses_declared_entry_when_it_is_not_first() {
289        let mut ast = AST::new();
290        ast.add_function_entry("helper".to_string());
291        ast.add_function_entry("dead".to_string());
292        ast.add_function_entry("entrypoint".to_string());
293        ast.nodes = vec![
294            ASTNode::GlobalDecl {
295                global_decl: GlobalDecl {
296                    entry_label: "entrypoint".to_string(),
297                    span: 0..0,
298                },
299            },
300            label_node("helper", 0),
301            instruction_node(Opcode::Exit, None, 0, None),
302            label_node("dead", 8),
303            instruction_node(Opcode::Exit, None, 8, None),
304            label_node("entrypoint", 16),
305            call_node("helper", 16),
306            instruction_node(Opcode::Exit, None, 24, None),
307        ];
308        ast.set_text_size(32);
309
310        let cfg = cfg_for_ast(&ast);
311        assert_eq!(cfg.functions()[0].name(), "entrypoint");
312        assert_eq!(cfg.functions()[0].entry_block_id(), Some(2)); // entrypoint is block 2 in source order
313
314        eliminate_unreachable_functions(&mut ast);
315
316        assert!(ast.nodes.iter().any(
317            |node| matches!(node, ASTNode::Label { label, .. } if label.name == "entrypoint")
318        ));
319        assert!(
320            ast.nodes
321                .iter()
322                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "helper"))
323        );
324        assert!(
325            !ast.nodes
326                .iter()
327                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
328        );
329    }
330
331    #[test]
332    fn test_optimizer_preserves_function_targeted_by_rodata_relocation() {
333        let mut ast = AST::new();
334        ast.add_function_entry("entrypoint".to_string());
335        ast.add_function_entry("target".to_string());
336        ast.add_function_entry("dead".to_string());
337        ast.nodes = vec![
338            ASTNode::GlobalDecl {
339                global_decl: GlobalDecl {
340                    entry_label: "entrypoint".to_string(),
341                    span: 0..0,
342                },
343            },
344            label_node("entrypoint", 0),
345            instruction_node(Opcode::Exit, None, 0, None),
346            label_node("target", 8),
347            instruction_node(Opcode::Exit, None, 8, None),
348            label_node("dead", 16),
349            instruction_node(Opcode::Exit, None, 16, None),
350        ];
351        ast.rodata_nodes.push(ASTNode::ROData {
352            rodata: ROData {
353                name: "pointer".to_string(),
354                args: vec![
355                    Token::Directive("quad".to_string(), 0..0),
356                    Token::VectorLiteral(vec![Number::Int(0)], 0..0),
357                ],
358                span: 0..0,
359            },
360            offset: 0,
361        });
362        ast.set_text_size(24);
363        ast.set_rodata_size(8);
364        ast.add_rodata_relocation(0, "target".to_string());
365
366        eliminate_unreachable_functions(&mut ast);
367
368        assert!(ast.nodes.iter().any(
369            |node| matches!(node, ASTNode::Label { label, .. } if label.name == "entrypoint")
370        ));
371        assert!(
372            ast.nodes
373                .iter()
374                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "target"))
375        );
376        assert!(
377            !ast.nodes
378                .iter()
379                .any(|node| matches!(node, ASTNode::Label { label, .. } if label.name == "dead"))
380        );
381    }
382
383    #[test]
384    fn test_write_back_drops_dead_functions_and_recomputes_offsets_cumulatively() {
385        let mut ast = AST::new();
386        for name in ["entrypoint", "dead_a", "live", "dead_b", "target"] {
387            ast.add_function_entry(name.to_string());
388        }
389        ast.nodes = vec![
390            label_node("entrypoint", 0),
391            call_node("live", 0),
392            call_node("target", 8),
393            instruction_node(Opcode::Exit, None, 16, None),
394            label_node("dead_a", 24),
395            instruction_node(Opcode::Exit, None, 24, None),
396            label_node("live", 32),
397            instruction_node(Opcode::Exit, None, 32, None),
398            label_node("dead_b", 40),
399            instruction_node(Opcode::Exit, None, 40, None),
400            label_node("target", 48),
401            instruction_node(Opcode::Exit, None, 48, None),
402        ];
403        ast.set_text_size(56);
404
405        eliminate_unreachable_functions(&mut ast);
406
407        assert_eq!(
408            listing(&ast.nodes),
409            [
410                ("entrypoint:".to_string(), 0),
411                ("call".to_string(), 0),
412                ("call".to_string(), 8),
413                ("exit".to_string(), 16),
414                ("live:".to_string(), 24),
415                ("exit".to_string(), 24),
416                ("target:".to_string(), 32),
417                ("exit".to_string(), 32),
418            ]
419        );
420    }
421
422    #[test]
423    fn test_assign_offsets_recomputes_from_scratch() {
424        let mut ast = AST::new();
425        ast.nodes = vec![
426            label_node("entrypoint", 999),
427            instruction_node(Opcode::Exit, None, 999, None),
428            label_node("next", 999),
429            instruction_node(Opcode::Exit, None, 999, None),
430        ];
431
432        assign_offsets(&mut ast);
433
434        assert!(matches!(
435            &ast.nodes[0],
436            ASTNode::Label { label, offset } if label.name == "entrypoint" && *offset == 0
437        ));
438        assert!(matches!(
439            &ast.nodes[1],
440            ASTNode::Instruction { offset, .. } if *offset == 0
441        ));
442        assert!(matches!(
443            &ast.nodes[2],
444            ASTNode::Label { label, offset } if label.name == "next" && *offset == 8
445        ));
446    }
447
448    fn alu_node(opcode: Opcode, dst: u8, src: Option<u8>, imm: Option<i64>) -> ASTNode {
449        ASTNode::Instruction {
450            instruction: Instruction {
451                opcode,
452                dst: Some(Register { n: dst }),
453                src: src.map(|n| Register { n }),
454                off: None,
455                imm: imm.map(|v| Either::Right(Number::Int(v))),
456                span: 0..0,
457            },
458            offset: 0,
459        }
460    }
461
462    /// `add32 r1, r2` whose full 64-bit result is shifted and returned, then a
463    /// later label. The shift reads r1 directly, so the def itself is fixed up.
464    fn alu32_program() -> AST {
465        let mut ast = AST::new();
466        ast.add_function_entry("entrypoint".to_string());
467        ast.nodes = vec![
468            label_node("entrypoint", 0),
469            alu_node(Opcode::Add32Reg, 1, Some(2), None),
470            instruction_node(Opcode::Ja, None, 8, Some(Either::Left("done".to_string()))),
471            label_node("done", 16),
472            alu_node(Opcode::Lsh64Imm, 1, None, Some(1)),
473            alu_node(Opcode::Mov64Reg, 0, Some(1), None),
474            instruction_node(Opcode::Exit, None, 32, None),
475        ];
476        ast.set_text_size(40);
477        ast
478    }
479
480    fn listing(nodes: &[ASTNode]) -> Vec<(String, u64)> {
481        nodes
482            .iter()
483            .filter_map(|node| match node {
484                ASTNode::Label { label, offset } => Some((format!("{}:", label.name), *offset)),
485                ASTNode::Instruction {
486                    instruction,
487                    offset,
488                } => Some((instruction.opcode.to_string(), *offset)),
489                _ => None,
490            })
491            .collect()
492    }
493
494    #[test]
495    fn test_fix_alu32_extension_inserts_fixup_and_reassigns_offsets() {
496        let mut ast = alu32_program();
497
498        fix_alu32_extension(&mut ast);
499
500        assert_eq!(
501            listing(&ast.nodes),
502            [
503                ("entrypoint:".to_string(), 0),
504                ("add32".to_string(), 0),
505                ("mov32".to_string(), 8),
506                ("ja".to_string(), 16),
507                ("done:".to_string(), 24),
508                ("lsh64".to_string(), 24),
509                ("mov64".to_string(), 32),
510                ("exit".to_string(), 40),
511            ]
512        );
513        let ASTNode::Instruction { instruction, .. } = &ast.nodes[2] else {
514            panic!("expected the inserted instruction");
515        };
516        assert_eq!(instruction.opcode, Opcode::Mov32Reg);
517        assert_eq!(instruction.dst, Some(Register { n: 1 }));
518        assert_eq!(instruction.src, Some(Register { n: 1 }));
519    }
520
521    #[test]
522    fn test_fix_alu32_extension_narrows_copy_in_place() {
523        let mut ast = AST::new();
524        ast.add_function_entry("entrypoint".to_string());
525        ast.nodes = vec![
526            label_node("entrypoint", 0),
527            alu_node(Opcode::Add32Reg, 1, Some(2), None),
528            alu_node(Opcode::Mov64Reg, 0, Some(1), None),
529            instruction_node(Opcode::Exit, None, 16, None),
530        ];
531        ast.set_text_size(24);
532
533        fix_alu32_extension(&mut ast);
534
535        assert_eq!(
536            listing(&ast.nodes),
537            [
538                ("entrypoint:".to_string(), 0),
539                ("add32".to_string(), 0),
540                ("mov32".to_string(), 8),
541                ("exit".to_string(), 16),
542            ]
543        );
544        let ASTNode::Instruction { instruction, .. } = &ast.nodes[2] else {
545            panic!("expected the narrowed copy");
546        };
547        assert_eq!(instruction.opcode, Opcode::Mov32Reg);
548        assert_eq!(instruction.dst, Some(Register { n: 0 }));
549        assert_eq!(instruction.src, Some(Register { n: 1 }));
550    }
551
552    #[test]
553    fn test_fix_alu32_extension_promotes_in_place() {
554        let mut ast = AST::new();
555        ast.add_function_entry("entrypoint".to_string());
556        ast.nodes = vec![
557            label_node("entrypoint", 0),
558            alu_node(Opcode::Lddw, 1, None, Some(0x9000_0000)),
559            alu_node(Opcode::Sub32Imm, 1, None, Some(1)),
560            alu_node(Opcode::Mov64Reg, 0, Some(1), None),
561            instruction_node(Opcode::Exit, None, 32, None),
562        ];
563        ast.set_text_size(40);
564
565        fix_alu32_extension(&mut ast);
566
567        assert_eq!(
568            listing(&ast.nodes),
569            [
570                ("entrypoint:".to_string(), 0),
571                ("lddw".to_string(), 0),
572                ("sub64".to_string(), 16),
573                ("mov64".to_string(), 24),
574                ("exit".to_string(), 32),
575            ]
576        );
577    }
578
579    #[test]
580    fn test_fix_alu32_extension_leaves_unaffected_program_untouched() {
581        let mut ast = AST::new();
582        ast.add_function_entry("entrypoint".to_string());
583        ast.nodes = vec![
584            label_node("entrypoint", 0),
585            alu_node(Opcode::Add32Reg, 1, Some(2), None),
586            alu_node(Opcode::Mov32Reg, 0, Some(1), None),
587            instruction_node(Opcode::Exit, None, 16, None),
588        ];
589        assign_offsets(&mut ast);
590        let before = listing(&ast.nodes);
591
592        fix_alu32_extension(&mut ast);
593
594        assert_eq!(listing(&ast.nodes), before);
595    }
596
597    #[test]
598    fn test_build_program_runs_fixup_only_when_optimizing() {
599        for arch in [SbpfArch::V0, SbpfArch::V3] {
600            let disabled =
601                build_program(alu32_program(), arch, OptimizationConfig::disabled()).unwrap();
602            assert!(
603                !listing(disabled.code_section.get_nodes())
604                    .iter()
605                    .any(|(op, _)| op == "mov32"),
606                "{arch:?}: disabled must not rewrite"
607            );
608
609            let enabled =
610                build_program(alu32_program(), arch, OptimizationConfig::enabled()).unwrap();
611            let ops: Vec<String> = listing(enabled.code_section.get_nodes())
612                .into_iter()
613                .map(|(op, _)| op)
614                .collect();
615            assert_eq!(
616                ops,
617                [
618                    "entrypoint:",
619                    "add32",
620                    "mov32",
621                    "ja",
622                    "done:",
623                    "lsh64",
624                    "mov64",
625                    "exit"
626                ],
627                "{arch:?}"
628            );
629        }
630    }
631
632    fn label_node(name: &str, offset: u64) -> ASTNode {
633        ASTNode::Label {
634            label: Label {
635                name: name.to_string(),
636                span: 0..0,
637            },
638            offset,
639        }
640    }
641
642    fn instruction_node(
643        opcode: Opcode,
644        dst: Option<u8>,
645        offset: u64,
646        off: Option<Either<String, i16>>,
647    ) -> ASTNode {
648        ASTNode::Instruction {
649            instruction: Instruction {
650                opcode,
651                dst: dst.map(|n| Register { n }),
652                src: None,
653                off,
654                imm: None,
655                span: 0..0,
656            },
657            offset,
658        }
659    }
660
661    fn call_node(target: &str, offset: u64) -> ASTNode {
662        ASTNode::Instruction {
663            instruction: Instruction {
664                opcode: Opcode::Call,
665                dst: None,
666                src: None,
667                off: None,
668                imm: Some(Either::Left(target.to_string())),
669                span: 0..0,
670            },
671            offset,
672        }
673    }
674}
sbpf_assembler/optimizer/mod.rs

sbpf_assembler/optimizer/
mod.rs
