bpf_script/
optimizer.rs

1use bpf_ins::{Instruction, Opcode};
2
3/// An [`Optimizer`] takes a set of input instructions and pushes optimized
4/// instructions to the output (the second argument) if it succeeeds. On success
5/// it should return `true`.
6type Optimizer = fn(&mut &[Instruction], &mut Vec<Instruction>) -> bool;
7
8/// Makes the following optimization:
9///
10///   r2 = r1   | r2 = *(r1 + N)
11///   r2 += N   |
12///   r2 = *r2  |
13///
14fn optimize_mov_add_load(inp: &mut &[Instruction], out: &mut Vec<Instruction>) -> bool {
15    const NEEDED: usize = 3;
16    if inp.len() < NEEDED {
17        return false;
18    }
19    let (ins, rem) = inp.split_at(NEEDED);
20    let load_size = if let Opcode::Memory(memory) = ins[2].get_opcode() {
21        *memory.get_size()
22    } else {
23        return false;
24    };
25
26    let offset1: i16 = match ins[1].get_imm().try_into() {
27        Ok(offset) => offset,
28        Err(_) => return false,
29    };
30
31    let check0 = Instruction::movx64(ins[0].get_dst_reg(), ins[0].get_src_reg());
32    let check1 = Instruction::add64(ins[1].get_dst_reg(), offset1.into());
33    let check2 = Instruction::loadx(ins[2].get_dst_reg(), ins[2].get_src_reg(), 0, load_size);
34
35    if check0 != ins[0] || check1 != ins[1] || check2 != ins[2] {
36        return false;
37    }
38
39    *inp = rem;
40    out.push(Instruction::loadx(
41        ins[0].get_dst_reg(),
42        ins[0].get_src_reg(),
43        offset1,
44        load_size,
45    ));
46    true
47}
48
49///
50/// Makes the following optimization:
51///
52///   r2 += N   | r2 = *(r2 + N)
53///   r2 = *r2  |
54///
55// I would think these could return Option<&'static [Instruction]>
56fn optimize_add_load(inp: &mut &[Instruction], out: &mut Vec<Instruction>) -> bool {
57    const NEEDED: usize = 2;
58    if inp.len() < NEEDED {
59        return false;
60    }
61    let (ins, rem) = inp.split_at(NEEDED);
62    let load_size = if let Opcode::Memory(memory) = ins[1].get_opcode() {
63        *memory.get_size()
64    } else {
65        return false;
66    };
67
68    let offset0: i16 = match ins[0].get_imm().try_into() {
69        Ok(offset) => offset,
70        Err(_) => return false,
71    };
72
73    let check0 = Instruction::add64(ins[0].get_dst_reg(), offset0.into());
74    let check1 = Instruction::loadx(ins[1].get_dst_reg(), ins[1].get_src_reg(), 0, load_size);
75
76    if check0 != ins[0] || check1 != ins[1] {
77        return false;
78    }
79
80    *inp = rem;
81    out.push(Instruction::loadx(
82        ins[0].get_dst_reg(),
83        ins[0].get_dst_reg(),
84        offset0,
85        load_size,
86    ));
87    true
88}
89
90fn no_optimization(inp: &mut &[Instruction], out: &mut Vec<Instruction>) -> bool {
91    let (ins, rem) = match inp.split_first() {
92        Some((ins, rem)) => (ins, rem),
93        None => return false,
94    };
95    out.push(*ins);
96    *inp = rem;
97    true
98}
99
100/// List of optimizers used by the `optimize` function.
101static OPTIMIZERS: [Optimizer; 3] = [optimize_mov_add_load, optimize_add_load, no_optimization];
102
103/// Applies various optimizations to the given list of instructions.
104///
105/// # Arguments
106///
107/// * `instructions` - The program, as a list of instructions, to optimize.
108pub fn optimize(mut instructions: &[Instruction]) -> Vec<Instruction> {
109    let mut optimized = vec![];
110    let instructions = &mut instructions;
111    while !instructions.is_empty() {
112        for optimizer in OPTIMIZERS {
113            optimizer(instructions, &mut optimized);
114        }
115    }
116
117    optimized
118}
bpf_script/optimizer.rs

bpf_script/
optimizer.rs
