Skip to main content

runar_compiler_rust/codegen/
optimizer.rs

1//! Peephole optimizer -- runs on Stack IR before emission.
2//!
3//! Scans for short sequences of stack operations that can be replaced with
4//! fewer or cheaper opcodes. Applies rules iteratively until a fixed point
5//! is reached. Mirrors the TypeScript peephole optimizer.
6
7use super::stack::{PushValue, StackOp};
8
9const MAX_ITERATIONS: usize = 100;
10
11/// Apply peephole optimization to a list of stack ops.
12pub fn optimize_stack_ops(ops: &[StackOp]) -> Vec<StackOp> {
13    // First, recursively optimize nested if-blocks
14    let mut current: Vec<StackOp> = ops.iter().map(optimize_nested_if).collect();
15
16    for _ in 0..MAX_ITERATIONS {
17        let (result, changed) = apply_one_pass(&current);
18        if !changed {
19            break;
20        }
21        current = result;
22    }
23
24    current
25}
26
27fn optimize_nested_if(op: &StackOp) -> StackOp {
28    match op {
29        StackOp::If {
30            then_ops,
31            else_ops,
32        } => {
33            let optimized_then = optimize_stack_ops(then_ops);
34            let optimized_else = if else_ops.is_empty() {
35                Vec::new()
36            } else {
37                optimize_stack_ops(else_ops)
38            };
39            StackOp::If {
40                then_ops: optimized_then,
41                else_ops: optimized_else,
42            }
43        }
44        other => other.clone(),
45    }
46}
47
48fn apply_one_pass(ops: &[StackOp]) -> (Vec<StackOp>, bool) {
49    let mut result = Vec::new();
50    let mut changed = false;
51    let mut i = 0;
52
53    while i < ops.len() {
54        // Try 4-op window
55        if i + 3 < ops.len() {
56            if let Some(replacement) =
57                match_window_4(&ops[i], &ops[i + 1], &ops[i + 2], &ops[i + 3])
58            {
59                result.extend(replacement);
60                i += 4;
61                changed = true;
62                continue;
63            }
64        }
65        // Try 3-op window
66        if i + 2 < ops.len() {
67            if let Some(replacement) = match_window_3(&ops[i], &ops[i + 1], &ops[i + 2]) {
68                result.extend(replacement);
69                i += 3;
70                changed = true;
71                continue;
72            }
73        }
74        // Existing 2-op window
75        if i + 1 < ops.len() {
76            if let Some(replacement) = match_window_2(&ops[i], &ops[i + 1]) {
77                result.extend(replacement);
78                i += 2;
79                changed = true;
80                continue;
81            }
82        }
83
84        result.push(ops[i].clone());
85        i += 1;
86    }
87
88    (result, changed)
89}
90
91fn match_window_3(a: &StackOp, b: &StackOp, c: &StackOp) -> Option<Vec<StackOp>> {
92    // Constant folding: PUSH(a) PUSH(b) ADD → PUSH(a+b)
93    if let (StackOp::Push(PushValue::Int(va)), StackOp::Push(PushValue::Int(vb))) = (a, b) {
94        if is_opcode(c, "OP_ADD") {
95            return Some(vec![StackOp::Push(PushValue::Int(va + vb))]);
96        }
97        if is_opcode(c, "OP_SUB") {
98            return Some(vec![StackOp::Push(PushValue::Int(va - vb))]);
99        }
100        if is_opcode(c, "OP_MUL") {
101            return Some(vec![StackOp::Push(PushValue::Int(va * vb))]);
102        }
103    }
104    None
105}
106
107fn match_window_4(
108    a: &StackOp,
109    b: &StackOp,
110    c: &StackOp,
111    d: &StackOp,
112) -> Option<Vec<StackOp>> {
113    // Chain folding: PUSH(a) ADD PUSH(b) ADD → PUSH(a+b) ADD
114    if let StackOp::Push(PushValue::Int(va)) = a {
115        if is_opcode(b, "OP_ADD") {
116            if let StackOp::Push(PushValue::Int(vb)) = c {
117                if is_opcode(d, "OP_ADD") {
118                    return Some(vec![
119                        StackOp::Push(PushValue::Int(va + vb)),
120                        StackOp::Opcode("OP_ADD".to_string()),
121                    ]);
122                }
123            }
124        }
125        if is_opcode(b, "OP_SUB") {
126            if let StackOp::Push(PushValue::Int(vb)) = c {
127                if is_opcode(d, "OP_SUB") {
128                    return Some(vec![
129                        StackOp::Push(PushValue::Int(va + vb)),
130                        StackOp::Opcode("OP_SUB".to_string()),
131                    ]);
132                }
133            }
134        }
135    }
136    None
137}
138
139fn match_window_2(a: &StackOp, b: &StackOp) -> Option<Vec<StackOp>> {
140    // PUSH x, DROP -> remove both
141    if matches!(a, StackOp::Push(_)) && matches!(b, StackOp::Drop) {
142        return Some(vec![]);
143    }
144
145    // DUP, DROP -> remove both
146    if matches!(a, StackOp::Dup) && matches!(b, StackOp::Drop) {
147        return Some(vec![]);
148    }
149
150    // SWAP, SWAP -> remove both
151    if matches!(a, StackOp::Swap) && matches!(b, StackOp::Swap) {
152        return Some(vec![]);
153    }
154
155    // PUSH 1, OP_ADD -> OP_1ADD
156    if is_push_int(a, 1) && is_opcode(b, "OP_ADD") {
157        return Some(vec![StackOp::Opcode("OP_1ADD".to_string())]);
158    }
159
160    // PUSH 1, OP_SUB -> OP_1SUB
161    if is_push_int(a, 1) && is_opcode(b, "OP_SUB") {
162        return Some(vec![StackOp::Opcode("OP_1SUB".to_string())]);
163    }
164
165    // PUSH 0, OP_ADD -> remove both
166    if is_push_int(a, 0) && is_opcode(b, "OP_ADD") {
167        return Some(vec![]);
168    }
169
170    // PUSH 0, OP_SUB -> remove both
171    if is_push_int(a, 0) && is_opcode(b, "OP_SUB") {
172        return Some(vec![]);
173    }
174
175    // OP_NOT, OP_NOT -> remove both
176    if is_opcode(a, "OP_NOT") && is_opcode(b, "OP_NOT") {
177        return Some(vec![]);
178    }
179
180    // OP_NEGATE, OP_NEGATE -> remove both
181    if is_opcode(a, "OP_NEGATE") && is_opcode(b, "OP_NEGATE") {
182        return Some(vec![]);
183    }
184
185    // OP_EQUAL, OP_VERIFY -> OP_EQUALVERIFY
186    if is_opcode(a, "OP_EQUAL") && is_opcode(b, "OP_VERIFY") {
187        return Some(vec![StackOp::Opcode("OP_EQUALVERIFY".to_string())]);
188    }
189
190    // OP_CHECKSIG, OP_VERIFY -> OP_CHECKSIGVERIFY
191    if is_opcode(a, "OP_CHECKSIG") && is_opcode(b, "OP_VERIFY") {
192        return Some(vec![StackOp::Opcode("OP_CHECKSIGVERIFY".to_string())]);
193    }
194
195    // OP_NUMEQUAL, OP_VERIFY -> OP_NUMEQUALVERIFY
196    if is_opcode(a, "OP_NUMEQUAL") && is_opcode(b, "OP_VERIFY") {
197        return Some(vec![StackOp::Opcode("OP_NUMEQUALVERIFY".to_string())]);
198    }
199
200    // OP_CHECKMULTISIG, OP_VERIFY -> OP_CHECKMULTISIGVERIFY
201    if is_opcode(a, "OP_CHECKMULTISIG") && is_opcode(b, "OP_VERIFY") {
202        return Some(vec![StackOp::Opcode(
203            "OP_CHECKMULTISIGVERIFY".to_string(),
204        )]);
205    }
206
207    // OP_DUP, OP_DROP -> remove both
208    if is_opcode(a, "OP_DUP") && is_opcode(b, "OP_DROP") {
209        return Some(vec![]);
210    }
211
212    // OP_OVER, OP_OVER -> OP_2DUP
213    if matches!(a, StackOp::Over) && matches!(b, StackOp::Over) {
214        return Some(vec![StackOp::Opcode("OP_2DUP".to_string())]);
215    }
216
217    // OP_DROP, OP_DROP -> OP_2DROP
218    if matches!(a, StackOp::Drop) && matches!(b, StackOp::Drop) {
219        return Some(vec![StackOp::Opcode("OP_2DROP".to_string())]);
220    }
221
222    // PUSH(0) + Roll{depth:0} → remove both
223    if is_push_int(a, 0) && matches!(b, StackOp::Roll { depth: 0 }) {
224        return Some(vec![]);
225    }
226
227    // PUSH(1) + Roll{depth:1} → Swap
228    if is_push_int(a, 1) && matches!(b, StackOp::Roll { depth: 1 }) {
229        return Some(vec![StackOp::Swap]);
230    }
231
232    // PUSH(2) + Roll{depth:2} → Rot
233    if is_push_int(a, 2) && matches!(b, StackOp::Roll { depth: 2 }) {
234        return Some(vec![StackOp::Rot]);
235    }
236
237    // PUSH(0) + Pick{depth:0} → Dup
238    if is_push_int(a, 0) && matches!(b, StackOp::Pick { depth: 0 }) {
239        return Some(vec![StackOp::Dup]);
240    }
241
242    // PUSH(1) + Pick{depth:1} → Over
243    if is_push_int(a, 1) && matches!(b, StackOp::Pick { depth: 1 }) {
244        return Some(vec![StackOp::Over]);
245    }
246
247    // Also match Push + Opcode("OP_ROLL"/"OP_PICK") (used by SLH-DSA codegen)
248    if is_push_int(a, 0) && is_opcode(b, "OP_ROLL") {
249        return Some(vec![]);
250    }
251    if is_push_int(a, 1) && is_opcode(b, "OP_ROLL") {
252        return Some(vec![StackOp::Swap]);
253    }
254    if is_push_int(a, 2) && is_opcode(b, "OP_ROLL") {
255        return Some(vec![StackOp::Rot]);
256    }
257    if is_push_int(a, 0) && is_opcode(b, "OP_PICK") {
258        return Some(vec![StackOp::Dup]);
259    }
260    if is_push_int(a, 1) && is_opcode(b, "OP_PICK") {
261        return Some(vec![StackOp::Over]);
262    }
263
264    // SHA256 + SHA256 → HASH256
265    if is_opcode(a, "OP_SHA256") && is_opcode(b, "OP_SHA256") {
266        return Some(vec![StackOp::Opcode("OP_HASH256".to_string())]);
267    }
268
269    // PUSH 0 + NUMEQUAL → NOT
270    if is_push_int(a, 0) && is_opcode(b, "OP_NUMEQUAL") {
271        return Some(vec![StackOp::Opcode("OP_NOT".to_string())]);
272    }
273
274    None
275}
276
277fn is_push_int(op: &StackOp, n: i128) -> bool {
278    matches!(op, StackOp::Push(PushValue::Int(v)) if *v == n)
279}
280
281fn is_opcode(op: &StackOp, code: &str) -> bool {
282    matches!(op, StackOp::Opcode(c) if c == code)
283}