|
| 1 | +// omnimcode-core/src/bytecode_opt.rs — Peephole + constant-folding passes |
| 2 | +// over compiled OMNIcode bytecode. |
| 3 | +// |
| 4 | +// Design: every pass that removes an op replaces it with `Op::Nop` |
| 5 | +// instead of actually shrinking the Vec, so already-computed jump |
| 6 | +// offsets stay valid. The VM treats Nop as a free no-op. Worth ~3 |
| 7 | +// cycles per Nop in the hot loop, but simpler to maintain than a |
| 8 | +// full re-emit pass that would have to walk all jumps and recompute |
| 9 | +// offsets. For the kind of programs OMNIcode runs (small kernels + |
| 10 | +// recursion, not megaword loops), the simplicity wins. |
| 11 | + |
| 12 | +use crate::bytecode::*; |
| 13 | + |
| 14 | +#[derive(Debug, Default, Clone)] |
| 15 | +pub struct OptStats { |
| 16 | + pub constants_folded: usize, |
| 17 | + pub dead_loads_removed: usize, |
| 18 | + pub double_nots_collapsed: usize, |
| 19 | + pub double_negs_collapsed: usize, |
| 20 | +} |
| 21 | + |
| 22 | +impl OptStats { |
| 23 | + pub fn total(&self) -> usize { |
| 24 | + self.constants_folded |
| 25 | + + self.dead_loads_removed |
| 26 | + + self.double_nots_collapsed |
| 27 | + + self.double_negs_collapsed |
| 28 | + } |
| 29 | +} |
| 30 | + |
| 31 | +/// Optimize a single function in place. Returns the stats from this run. |
| 32 | +pub fn optimize_function(func: &mut CompiledFunction) -> OptStats { |
| 33 | + let mut stats = OptStats::default(); |
| 34 | + // Run passes until a fixpoint is reached. In practice 2-3 iterations. |
| 35 | + loop { |
| 36 | + let before = stats.total(); |
| 37 | + constant_fold_pass(func, &mut stats); |
| 38 | + dead_load_pass(func, &mut stats); |
| 39 | + double_unary_pass(func, &mut stats); |
| 40 | + if stats.total() == before { |
| 41 | + break; |
| 42 | + } |
| 43 | + } |
| 44 | + stats |
| 45 | +} |
| 46 | + |
| 47 | +/// Fold `LoadConst a; LoadConst b; <op>` into `Nop; Nop; LoadConst c`. |
| 48 | +/// The arithmetic and comparison ops are pure functions of the operand |
| 49 | +/// pair, so this is safe regardless of surrounding control flow as |
| 50 | +/// long as we don't disturb the jump-offset count (we don't — Nops |
| 51 | +/// preserve indices). |
| 52 | +fn constant_fold_pass(func: &mut CompiledFunction, stats: &mut OptStats) { |
| 53 | + let n = func.ops.len(); |
| 54 | + if n < 3 { |
| 55 | + return; |
| 56 | + } |
| 57 | + for i in 0..(n - 2) { |
| 58 | + let (a, b, op) = match (&func.ops[i], &func.ops[i + 1], &func.ops[i + 2]) { |
| 59 | + (Op::LoadConst(a_idx), Op::LoadConst(b_idx), op) => { |
| 60 | + (*a_idx, *b_idx, op.clone()) |
| 61 | + } |
| 62 | + _ => continue, |
| 63 | + }; |
| 64 | + let a_val = match func.constants.get(a) { |
| 65 | + Some(c) => c.clone(), |
| 66 | + None => continue, |
| 67 | + }; |
| 68 | + let b_val = match func.constants.get(b) { |
| 69 | + Some(c) => c.clone(), |
| 70 | + None => continue, |
| 71 | + }; |
| 72 | + let folded = match fold_binary(&a_val, &b_val, &op) { |
| 73 | + Some(v) => v, |
| 74 | + None => continue, |
| 75 | + }; |
| 76 | + let new_idx = func.constants.len(); |
| 77 | + func.constants.push(folded); |
| 78 | + func.ops[i] = Op::Nop; |
| 79 | + func.ops[i + 1] = Op::Nop; |
| 80 | + func.ops[i + 2] = Op::LoadConst(new_idx); |
| 81 | + stats.constants_folded += 1; |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +/// Remove `LoadConst N; Pop` pairs — the constant is loaded only to be |
| 86 | +/// discarded. Both become Nops. |
| 87 | +fn dead_load_pass(func: &mut CompiledFunction, stats: &mut OptStats) { |
| 88 | + let n = func.ops.len(); |
| 89 | + if n < 2 { |
| 90 | + return; |
| 91 | + } |
| 92 | + for i in 0..(n - 1) { |
| 93 | + if matches!(func.ops[i], Op::LoadConst(_)) && matches!(func.ops[i + 1], Op::Pop) { |
| 94 | + func.ops[i] = Op::Nop; |
| 95 | + func.ops[i + 1] = Op::Nop; |
| 96 | + stats.dead_loads_removed += 1; |
| 97 | + } |
| 98 | + } |
| 99 | +} |
| 100 | + |
| 101 | +/// Collapse `Not; Not` (and similar double-unary ops) to no-op. |
| 102 | +fn double_unary_pass(func: &mut CompiledFunction, stats: &mut OptStats) { |
| 103 | + let n = func.ops.len(); |
| 104 | + if n < 2 { |
| 105 | + return; |
| 106 | + } |
| 107 | + for i in 0..(n - 1) { |
| 108 | + match (&func.ops[i], &func.ops[i + 1]) { |
| 109 | + (Op::Not, Op::Not) => { |
| 110 | + func.ops[i] = Op::Nop; |
| 111 | + func.ops[i + 1] = Op::Nop; |
| 112 | + stats.double_nots_collapsed += 1; |
| 113 | + } |
| 114 | + (Op::Neg, Op::Neg) => { |
| 115 | + func.ops[i] = Op::Nop; |
| 116 | + func.ops[i + 1] = Op::Nop; |
| 117 | + stats.double_negs_collapsed += 1; |
| 118 | + } |
| 119 | + _ => {} |
| 120 | + } |
| 121 | + } |
| 122 | +} |
| 123 | + |
| 124 | +/// Apply a binary op to two constants. Returns None if the op isn't |
| 125 | +/// foldable (e.g. it's a control-flow op, or the constants are |
| 126 | +/// incompatible). |
| 127 | +fn fold_binary(a: &Const, b: &Const, op: &Op) -> Option<Const> { |
| 128 | + // Promote to float if either is float. |
| 129 | + let any_float = matches!(a, Const::Float(_)) || matches!(b, Const::Float(_)); |
| 130 | + if any_float { |
| 131 | + let af = const_to_float(a)?; |
| 132 | + let bf = const_to_float(b)?; |
| 133 | + return match op { |
| 134 | + Op::Add => Some(Const::Float(af + bf)), |
| 135 | + Op::Sub => Some(Const::Float(af - bf)), |
| 136 | + Op::Mul => Some(Const::Float(af * bf)), |
| 137 | + Op::Div => { |
| 138 | + if bf == 0.0 { |
| 139 | + None // can't fold div-by-zero (produces Singularity) |
| 140 | + } else { |
| 141 | + Some(Const::Float(af / bf)) |
| 142 | + } |
| 143 | + } |
| 144 | + Op::Eq => Some(Const::Bool(af == bf)), |
| 145 | + Op::Ne => Some(Const::Bool(af != bf)), |
| 146 | + Op::Lt => Some(Const::Bool(af < bf)), |
| 147 | + Op::Le => Some(Const::Bool(af <= bf)), |
| 148 | + Op::Gt => Some(Const::Bool(af > bf)), |
| 149 | + Op::Ge => Some(Const::Bool(af >= bf)), |
| 150 | + _ => None, |
| 151 | + }; |
| 152 | + } |
| 153 | + let ai = const_to_int(a)?; |
| 154 | + let bi = const_to_int(b)?; |
| 155 | + match op { |
| 156 | + Op::Add => Some(Const::Int(ai.wrapping_add(bi))), |
| 157 | + Op::Sub => Some(Const::Int(ai.wrapping_sub(bi))), |
| 158 | + Op::Mul => Some(Const::Int(ai.wrapping_mul(bi))), |
| 159 | + Op::Div => { |
| 160 | + if bi == 0 { |
| 161 | + None |
| 162 | + } else { |
| 163 | + Some(Const::Int(ai / bi)) |
| 164 | + } |
| 165 | + } |
| 166 | + Op::Mod => { |
| 167 | + if bi == 0 { |
| 168 | + None |
| 169 | + } else { |
| 170 | + Some(Const::Int(ai % bi)) |
| 171 | + } |
| 172 | + } |
| 173 | + Op::Eq => Some(Const::Bool(ai == bi)), |
| 174 | + Op::Ne => Some(Const::Bool(ai != bi)), |
| 175 | + Op::Lt => Some(Const::Bool(ai < bi)), |
| 176 | + Op::Le => Some(Const::Bool(ai <= bi)), |
| 177 | + Op::Gt => Some(Const::Bool(ai > bi)), |
| 178 | + Op::Ge => Some(Const::Bool(ai >= bi)), |
| 179 | + Op::BitAnd => Some(Const::Int(ai & bi)), |
| 180 | + Op::BitOr => Some(Const::Int(ai | bi)), |
| 181 | + Op::BitXor => Some(Const::Int(ai ^ bi)), |
| 182 | + Op::Shl => Some(Const::Int(ai.wrapping_shl((bi & 63) as u32))), |
| 183 | + Op::Shr => Some(Const::Int(ai.wrapping_shr((bi & 63) as u32))), |
| 184 | + _ => None, |
| 185 | + } |
| 186 | +} |
| 187 | + |
| 188 | +fn const_to_int(c: &Const) -> Option<i64> { |
| 189 | + match c { |
| 190 | + Const::Int(n) => Some(*n), |
| 191 | + Const::Bool(b) => Some(if *b { 1 } else { 0 }), |
| 192 | + _ => None, |
| 193 | + } |
| 194 | +} |
| 195 | + |
| 196 | +fn const_to_float(c: &Const) -> Option<f64> { |
| 197 | + match c { |
| 198 | + Const::Int(n) => Some(*n as f64), |
| 199 | + Const::Float(f) => Some(*f), |
| 200 | + Const::Bool(b) => Some(if *b { 1.0 } else { 0.0 }), |
| 201 | + _ => None, |
| 202 | + } |
| 203 | +} |
| 204 | + |
| 205 | +pub fn optimize_module(module: &mut Module) -> OptStats { |
| 206 | + let mut total = OptStats::default(); |
| 207 | + let stats_main = optimize_function(&mut module.main); |
| 208 | + total.constants_folded += stats_main.constants_folded; |
| 209 | + total.dead_loads_removed += stats_main.dead_loads_removed; |
| 210 | + total.double_nots_collapsed += stats_main.double_nots_collapsed; |
| 211 | + total.double_negs_collapsed += stats_main.double_negs_collapsed; |
| 212 | + for (_, func) in module.functions.iter_mut() { |
| 213 | + let s = optimize_function(func); |
| 214 | + total.constants_folded += s.constants_folded; |
| 215 | + total.dead_loads_removed += s.dead_loads_removed; |
| 216 | + total.double_nots_collapsed += s.double_nots_collapsed; |
| 217 | + total.double_negs_collapsed += s.double_negs_collapsed; |
| 218 | + } |
| 219 | + total |
| 220 | +} |
| 221 | + |
| 222 | +#[cfg(test)] |
| 223 | +mod tests { |
| 224 | + use super::*; |
| 225 | + use crate::compiler::compile_program; |
| 226 | + use crate::parser::Parser; |
| 227 | + |
| 228 | + fn compile_and_opt(src: &str) -> (Module, OptStats) { |
| 229 | + let mut parser = Parser::new(src); |
| 230 | + let stmts = parser.parse().unwrap(); |
| 231 | + let mut module = compile_program(&stmts).unwrap(); |
| 232 | + let stats = optimize_module(&mut module); |
| 233 | + (module, stats) |
| 234 | + } |
| 235 | + |
| 236 | + #[test] |
| 237 | + fn folds_simple_int_add() { |
| 238 | + let (_, stats) = compile_and_opt("h x = 2 + 3;"); |
| 239 | + assert!(stats.constants_folded >= 1); |
| 240 | + } |
| 241 | + |
| 242 | + #[test] |
| 243 | + fn chained_arithmetic_folds_to_one_constant() { |
| 244 | + let (m, stats) = compile_and_opt("h x = 1 + 2 + 3 + 4;"); |
| 245 | + assert!(stats.constants_folded >= 3, "expected >=3 folds, got {}", stats.constants_folded); |
| 246 | + // After folding, main should contain a single LoadConst(10) plus |
| 247 | + // StoreVar plus a return — at least one of the constants is 10. |
| 248 | + assert!(m |
| 249 | + .main |
| 250 | + .constants |
| 251 | + .iter() |
| 252 | + .any(|c| matches!(c, Const::Int(10))), |
| 253 | + ); |
| 254 | + } |
| 255 | + |
| 256 | + #[test] |
| 257 | + fn folds_bitwise() { |
| 258 | + let (m, stats) = compile_and_opt("h x = 255 & 15;"); |
| 259 | + assert!(stats.constants_folded >= 1); |
| 260 | + assert!(m |
| 261 | + .main |
| 262 | + .constants |
| 263 | + .iter() |
| 264 | + .any(|c| matches!(c, Const::Int(15))), |
| 265 | + ); |
| 266 | + } |
| 267 | + |
| 268 | + #[test] |
| 269 | + fn folds_shift() { |
| 270 | + let (m, stats) = compile_and_opt("h x = 1 << 8;"); |
| 271 | + assert!(stats.constants_folded >= 1); |
| 272 | + assert!(m |
| 273 | + .main |
| 274 | + .constants |
| 275 | + .iter() |
| 276 | + .any(|c| matches!(c, Const::Int(256))), |
| 277 | + ); |
| 278 | + } |
| 279 | + |
| 280 | + #[test] |
| 281 | + fn does_not_fold_div_by_zero() { |
| 282 | + // 10 / 0 must NOT be pre-folded — at runtime it produces a Singularity. |
| 283 | + let (_, stats) = compile_and_opt("h x = 10 / 0;"); |
| 284 | + assert_eq!(stats.constants_folded, 0, "must preserve div-by-zero semantics"); |
| 285 | + } |
| 286 | + |
| 287 | + #[test] |
| 288 | + fn folds_float_arithmetic() { |
| 289 | + let (m, stats) = compile_and_opt("h x = 1.5 + 2.5;"); |
| 290 | + assert!(stats.constants_folded >= 1); |
| 291 | + assert!(m |
| 292 | + .main |
| 293 | + .constants |
| 294 | + .iter() |
| 295 | + .any(|c| matches!(c, Const::Float(f) if (f - 4.0).abs() < 1e-9)), |
| 296 | + ); |
| 297 | + } |
| 298 | + |
| 299 | + #[test] |
| 300 | + fn folds_comparison() { |
| 301 | + let (m, stats) = compile_and_opt("h x = 10 < 20;"); |
| 302 | + assert!(stats.constants_folded >= 1); |
| 303 | + assert!(m |
| 304 | + .main |
| 305 | + .constants |
| 306 | + .iter() |
| 307 | + .any(|c| matches!(c, Const::Bool(true))), |
| 308 | + ); |
| 309 | + } |
| 310 | +} |
0 commit comments