1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use bpf_ins::{Instruction, Opcode};

/// An [`Optimizer`] takes a set of input instructions and pushes optimized
/// instructions to the output (the second argument) if it succeeeds. On success
/// it should return `true`.
type Optimizer = fn(&mut &[Instruction], &mut Vec<Instruction>) -> bool;

/// Makes the following optimization:
///
///   r2 = r1   | r2 = *(r1 + N)
///   r2 += N   |
///   r2 = *r2  |
///
fn optimize_mov_add_load(inp: &mut &[Instruction], out: &mut Vec<Instruction>) -> bool {
    const NEEDED: usize = 3;
    if inp.len() < NEEDED {
        return false;
    }
    let (ins, rem) = inp.split_at(NEEDED);
    let load_size = if let Opcode::Memory(memory) = ins[2].get_opcode() {
        *memory.get_size()
    } else {
        return false;
    };

    let offset1: i16 = match ins[1].get_imm().try_into() {
        Ok(offset) => offset,
        Err(_) => return false,
    };

    let check0 = Instruction::movx64(ins[0].get_dst_reg(), ins[0].get_src_reg());
    let check1 = Instruction::add64(ins[1].get_dst_reg(), offset1.into());
    let check2 = Instruction::loadx(ins[2].get_dst_reg(), ins[2].get_src_reg(), 0, load_size);

    if check0 != ins[0] || check1 != ins[1] || check2 != ins[2] {
        return false;
    }

    *inp = rem;
    out.push(Instruction::loadx(
        ins[0].get_dst_reg(),
        ins[0].get_src_reg(),
        offset1,
        load_size,
    ));
    true
}

///
/// Makes the following optimization:
///
///   r2 += N   | r2 = *(r2 + N)
///   r2 = *r2  |
///
// I would think these could return Option<&'static [Instruction]>
fn optimize_add_load(inp: &mut &[Instruction], out: &mut Vec<Instruction>) -> bool {
    const NEEDED: usize = 2;
    if inp.len() < NEEDED {
        return false;
    }
    let (ins, rem) = inp.split_at(NEEDED);
    let load_size = if let Opcode::Memory(memory) = ins[1].get_opcode() {
        *memory.get_size()
    } else {
        return false;
    };

    let offset0: i16 = match ins[0].get_imm().try_into() {
        Ok(offset) => offset,
        Err(_) => return false,
    };

    let check0 = Instruction::add64(ins[0].get_dst_reg(), offset0.into());
    let check1 = Instruction::loadx(ins[1].get_dst_reg(), ins[1].get_src_reg(), 0, load_size);

    if check0 != ins[0] || check1 != ins[1] {
        return false;
    }

    *inp = rem;
    out.push(Instruction::loadx(
        ins[0].get_dst_reg(),
        ins[0].get_dst_reg(),
        offset0,
        load_size,
    ));
    true
}

fn no_optimization(inp: &mut &[Instruction], out: &mut Vec<Instruction>) -> bool {
    let (ins, rem) = match inp.split_first() {
        Some((ins, rem)) => (ins, rem),
        None => return false,
    };
    out.push(*ins);
    *inp = rem;
    true
}

/// List of optimizers used by the `optimize` function.
static OPTIMIZERS: [Optimizer; 3] = [optimize_mov_add_load, optimize_add_load, no_optimization];

/// Applies various optimizations to the given list of instructions.
///
/// # Arguments
///
/// * `instructions` - The program, as a list of instructions, to optimize.
pub fn optimize(mut instructions: &[Instruction]) -> Vec<Instruction> {
    let mut optimized = vec![];
    let instructions = &mut instructions;
    while !instructions.is_empty() {
        for optimizer in OPTIMIZERS {
            optimizer(instructions, &mut optimized);
        }
    }

    optimized
}