use crate::chunk::Chunk;
use crate::opcode::Opcode;
use crate::value::ConstValue;
pub fn peephole(chunk: &mut Chunk) {
let mut i: usize = 0;
while i < chunk.code.len() {
if let Some(rewrite) = try_match(chunk, i) {
let replacement_len = rewrite.replacement_bytes.len();
if apply_rewrite(chunk, i, rewrite) {
i += replacement_len;
continue;
}
}
match Opcode::from_u8(chunk.code[i]) {
Some(op) => i += 1 + op.operand_bytes() as usize,
None => i += 1,
}
}
}
struct Rewrite {
window_bytes_in: usize,
replacement_bytes: Vec<u8>,
surviving_line: u32,
new_constant: Option<ConstValue>,
}
fn try_match(chunk: &Chunk, i: usize) -> Option<Rewrite> {
let code = &chunk.code;
let const_byte = Opcode::Const as u8;
let pop_byte = Opcode::Pop as u8;
let dup_byte = Opcode::Dup as u8;
let not_byte = Opcode::Not as u8;
let jif_byte = Opcode::JumpIfFalse as u8;
let jit_byte = Opcode::JumpIfTrue as u8;
let jump_byte = Opcode::Jump as u8;
let add_byte = Opcode::Add as u8;
let fadd_byte = Opcode::FAdd as u8;
if i + 4 <= code.len() && code[i] == const_byte && code[i + 3] == pop_byte {
return Some(Rewrite {
window_bytes_in: 4,
replacement_bytes: Vec::new(),
surviving_line: 0,
new_constant: None,
});
}
if i + 2 <= code.len() && code[i] == dup_byte && code[i + 1] == pop_byte {
return Some(Rewrite {
window_bytes_in: 2,
replacement_bytes: Vec::new(),
surviving_line: 0,
new_constant: None,
});
}
if i + 4 <= code.len() && code[i] == not_byte && code[i + 1] == jif_byte {
return Some(Rewrite {
window_bytes_in: 4,
replacement_bytes: vec![jit_byte, code[i + 2], code[i + 3]],
surviving_line: chunk.source_lines[i + 1],
new_constant: None,
});
}
if i + 6 <= code.len() && code[i] == const_byte {
let cidx = chunk.read_u16(i + 1) as usize;
let cond = chunk.constants.get(cidx);
let jump_op = code[i + 3];
if jump_op == jif_byte && cond == Some(&ConstValue::Bool(true)) {
return Some(Rewrite {
window_bytes_in: 6,
replacement_bytes: Vec::new(),
surviving_line: 0,
new_constant: None,
});
}
if jump_op == jif_byte && cond == Some(&ConstValue::Bool(false)) {
return Some(Rewrite {
window_bytes_in: 6,
replacement_bytes: vec![jump_byte, code[i + 4], code[i + 5]],
surviving_line: chunk.source_lines[i + 3],
new_constant: None,
});
}
if jump_op == jit_byte && cond == Some(&ConstValue::Bool(true)) {
return Some(Rewrite {
window_bytes_in: 6,
replacement_bytes: vec![jump_byte, code[i + 4], code[i + 5]],
surviving_line: chunk.source_lines[i + 3],
new_constant: None,
});
}
if jump_op == jit_byte && cond == Some(&ConstValue::Bool(false)) {
return Some(Rewrite {
window_bytes_in: 6,
replacement_bytes: Vec::new(),
surviving_line: 0,
new_constant: None,
});
}
}
if i + 7 <= code.len()
&& code[i] == const_byte
&& code[i + 3] == const_byte
&& (code[i + 6] == add_byte || code[i + 6] == fadd_byte)
{
let a_idx = chunk.read_u16(i + 1) as usize;
let b_idx = chunk.read_u16(i + 4) as usize;
let a = chunk.constants.get(a_idx);
let b_val = chunk.constants.get(b_idx);
let folded = match (a, b_val, code[i + 6]) {
(Some(ConstValue::I64(x)), Some(ConstValue::I64(y)), op_byte)
if op_byte == add_byte =>
{
x.checked_add(*y).map(ConstValue::I64)
}
(Some(ConstValue::F64(x)), Some(ConstValue::F64(y)), op_byte)
if op_byte == fadd_byte =>
{
Some(ConstValue::F64(x + y))
}
_ => None,
};
if let Some(value) = folded {
return Some(Rewrite {
window_bytes_in: 7,
replacement_bytes: vec![const_byte, 0, 0],
surviving_line: chunk.source_lines[i + 6],
new_constant: Some(value),
});
}
}
None
}
fn apply_rewrite(chunk: &mut Chunk, at: usize, mut rewrite: Rewrite) -> bool {
let delta = rewrite.window_bytes_in - rewrite.replacement_bytes.len();
if delta > 0 && !prevalidate_jump_offsets(chunk, at, delta) {
return false;
}
if let Some(value) = rewrite.new_constant.take() {
let idx = chunk.add_constant(value);
let bytes = idx.to_le_bytes();
rewrite.replacement_bytes[1] = bytes[0];
rewrite.replacement_bytes[2] = bytes[1];
}
if delta > 0 {
fix_jump_offsets(chunk, at, delta);
}
let end = at + rewrite.window_bytes_in;
chunk
.code
.splice(at..end, rewrite.replacement_bytes.iter().copied());
chunk.source_lines.splice(
at..end,
std::iter::repeat_n(rewrite.surviving_line, rewrite.replacement_bytes.len()),
);
debug_assert_eq!(
chunk.code.len(),
chunk.source_lines.len(),
"peephole broke the code/source_lines lockstep invariant"
);
true
}
fn jump_offset_pos(op: Opcode, op_pos: usize) -> Option<usize> {
match op {
Opcode::Jump | Opcode::JumpIfFalse | Opcode::JumpIfTrue => Some(op_pos + 1),
Opcode::MatchVariant => Some(op_pos + 3),
_ => None,
}
}
fn adjusted_offset(
op_pos: usize,
offset_pos: usize,
offset: i16,
cut_at: usize,
delta: usize,
) -> Option<isize> {
let off = offset as isize;
let target = offset_pos as isize + 2 + off;
let src = op_pos as isize;
let cut = cut_at as isize;
let cut_end = (cut_at + delta) as isize;
let d = delta as isize;
if src < cut && target >= cut_end {
Some(off - d)
} else if src >= cut_end && target < cut {
Some(off + d)
} else {
None
}
}
fn fix_jump_offsets(chunk: &mut Chunk, cut_at: usize, delta: usize) {
let mut j = 0usize;
while j < chunk.code.len() {
let Some(op) = Opcode::from_u8(chunk.code[j]) else {
j += 1;
continue;
};
let size = 1 + op.operand_bytes() as usize;
if let Some(offset_pos) = jump_offset_pos(op, j) {
let inside_cut = j >= cut_at && j < cut_at + delta;
if !inside_cut && offset_pos + 2 <= chunk.code.len() {
let offset = chunk.read_i16(offset_pos);
if let Some(new_off) = adjusted_offset(j, offset_pos, offset, cut_at, delta) {
let clamped = new_off.clamp(i16::MIN as isize, i16::MAX as isize) as i16;
let bytes = clamped.to_le_bytes();
chunk.code[offset_pos] = bytes[0];
chunk.code[offset_pos + 1] = bytes[1];
}
}
}
j += size;
}
}
fn prevalidate_jump_offsets(chunk: &Chunk, cut_at: usize, delta: usize) -> bool {
let mut j = 0usize;
while j < chunk.code.len() {
let Some(op) = Opcode::from_u8(chunk.code[j]) else {
j += 1;
continue;
};
let size = 1 + op.operand_bytes() as usize;
if let Some(offset_pos) = jump_offset_pos(op, j) {
let inside_cut = j >= cut_at && j < cut_at + delta;
if !inside_cut && offset_pos + 2 <= chunk.code.len() {
let offset = chunk.read_i16(offset_pos);
if let Some(new_off) = adjusted_offset(j, offset_pos, offset, cut_at, delta)
&& (new_off < i16::MIN as isize || new_off > i16::MAX as isize)
{
return false;
}
}
}
j += size;
}
true
}
#[cfg(test)]
mod tests {
use super::*;
use crate::chunk::Program;
use crate::lexer::Lexer;
use crate::parser::Parser;
use crate::typechecker::check_program;
fn write_const(c: &mut Chunk, idx: u16, line: u32) {
c.write_op(Opcode::Const, line);
c.write_u16(idx, line);
}
fn write_jump(c: &mut Chunk, op: Opcode, offset: i16, line: u32) {
c.write_op(op, line);
c.write_i16(offset, line);
}
fn compile(src: &str) -> Program {
let tokens = Lexer::tokenize(src).expect("lex");
let ast = Parser::parse(&tokens).expect("parse");
let (typed, errors, _) = check_program(&ast, src);
assert!(errors.is_empty(), "typecheck errors: {errors:?}");
crate::codegen::compile_program(&typed, src).expect("compile")
}
#[test]
fn pattern_1_const_pop_pair_is_removed() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(7));
write_const(&mut c, idx, 1);
c.write_op(Opcode::Pop, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code, vec![Opcode::Halt as u8]);
assert_eq!(c.code.len(), 1);
assert_eq!(c.source_lines.len(), 1);
assert_eq!(c.source_lines, vec![2]);
}
#[test]
fn pattern_2_dup_pop_pair_is_removed() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(0));
write_const(&mut c, idx, 1);
c.write_op(Opcode::Dup, 1);
c.write_op(Opcode::Pop, 1);
c.write_op(Opcode::Return, 2);
peephole(&mut c);
assert_eq!(
c.code,
vec![Opcode::Const as u8, 0, 0, Opcode::Return as u8]
);
assert_eq!(c.code.len(), c.source_lines.len());
}
#[test]
fn pattern_3_not_jump_if_false_becomes_jump_if_true() {
let mut c = Chunk::new();
c.write_op(Opcode::Not, 1);
write_jump(&mut c, Opcode::JumpIfFalse, 5, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code[0], Opcode::JumpIfTrue as u8);
assert_eq!(c.read_i16(1), 5);
assert_eq!(c.code[3], Opcode::Halt as u8);
assert_eq!(c.code.len(), 4);
assert_eq!(c.code.len(), c.source_lines.len());
}
#[test]
fn pattern_4_const_true_jump_if_false_is_removed() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::Bool(true));
write_const(&mut c, idx, 1);
write_jump(&mut c, Opcode::JumpIfFalse, 5, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code, vec![Opcode::Halt as u8]);
assert_eq!(c.code.len(), c.source_lines.len());
}
#[test]
fn pattern_5_const_false_jump_if_false_becomes_jump() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::Bool(false));
write_const(&mut c, idx, 1);
write_jump(&mut c, Opcode::JumpIfFalse, 5, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code[0], Opcode::Jump as u8);
assert_eq!(c.read_i16(1), 5);
assert_eq!(c.code[3], Opcode::Halt as u8);
assert_eq!(c.code.len(), 4);
assert_eq!(c.code.len(), c.source_lines.len());
}
#[test]
fn pattern_6_const_true_jump_if_true_becomes_jump() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::Bool(true));
write_const(&mut c, idx, 1);
write_jump(&mut c, Opcode::JumpIfTrue, 9, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code[0], Opcode::Jump as u8);
assert_eq!(c.read_i16(1), 9);
assert_eq!(c.code[3], Opcode::Halt as u8);
assert_eq!(c.code.len(), c.source_lines.len());
}
#[test]
fn pattern_7_const_false_jump_if_true_is_removed() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::Bool(false));
write_const(&mut c, idx, 1);
write_jump(&mut c, Opcode::JumpIfTrue, 5, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code, vec![Opcode::Halt as u8]);
assert_eq!(c.code.len(), c.source_lines.len());
}
#[test]
fn pattern_8_folds_two_i64_const_add() {
let mut c = Chunk::new();
let a = c.add_constant(ConstValue::I64(2));
let b = c.add_constant(ConstValue::I64(3));
write_const(&mut c, a, 1);
write_const(&mut c, b, 1);
c.write_op(Opcode::Add, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code[0], Opcode::Const as u8);
assert_eq!(c.code[3], Opcode::Halt as u8);
assert_eq!(c.code.len(), 4);
let new_idx = c.read_u16(1);
assert_eq!(c.constants[new_idx as usize], ConstValue::I64(5));
assert_eq!(c.constants[0], ConstValue::I64(2));
assert_eq!(c.constants[1], ConstValue::I64(3));
assert_eq!(c.constants.len(), 3);
}
#[test]
fn pattern_8_folds_two_f64_const_fadd() {
let mut c = Chunk::new();
let a = c.add_constant(ConstValue::F64(1.5));
let b = c.add_constant(ConstValue::F64(2.5));
write_const(&mut c, a, 1);
write_const(&mut c, b, 1);
c.write_op(Opcode::FAdd, 1);
c.write_op(Opcode::Halt, 2);
peephole(&mut c);
assert_eq!(c.code[0], Opcode::Const as u8);
assert_eq!(c.code[3], Opcode::Halt as u8);
let new_idx = c.read_u16(1);
assert_eq!(c.constants[new_idx as usize], ConstValue::F64(4.0));
}
#[test]
fn pattern_8_skips_fold_on_i64_overflow() {
let mut c = Chunk::new();
let a = c.add_constant(ConstValue::I64(i64::MAX));
let b = c.add_constant(ConstValue::I64(1));
write_const(&mut c, a, 1);
write_const(&mut c, b, 1);
c.write_op(Opcode::Add, 1);
c.write_op(Opcode::Halt, 2);
let before = c.code.clone();
peephole(&mut c);
assert_eq!(c.code, before);
assert_eq!(c.constants.len(), 2);
}
#[test]
fn pattern_3_preserves_the_surviving_source_line() {
let mut c = Chunk::new();
c.write_op(Opcode::Not, 5);
write_jump(&mut c, Opcode::JumpIfFalse, 3, 5);
peephole(&mut c);
assert_eq!(c.source_lines, vec![5, 5, 5]);
}
#[test]
fn jump_offset_is_fixed_up_after_a_forward_removal() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(0));
write_jump(&mut c, Opcode::Jump, 5, 1);
write_const(&mut c, idx, 1);
c.write_op(Opcode::Pop, 1);
c.write_op(Opcode::Halt, 2);
assert_eq!(c.code.len(), 8);
peephole(&mut c);
assert_eq!(c.code.len(), 4);
assert_eq!(c.code[0], Opcode::Jump as u8);
assert_eq!(c.read_i16(1), 1);
assert_eq!(c.code[3], Opcode::Halt as u8);
}
#[test]
fn backward_jump_offset_is_fixed_up_after_a_removal() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(0));
c.write_op(Opcode::Halt, 1);
write_const(&mut c, idx, 2);
c.write_op(Opcode::Pop, 2);
write_jump(&mut c, Opcode::Jump, -8, 3);
assert_eq!(c.code.len(), 8);
peephole(&mut c);
assert_eq!(c.code.len(), 4);
assert_eq!(c.code[1], Opcode::Jump as u8);
assert_eq!(c.read_i16(2), -4);
}
#[test]
fn multiple_patterns_collapse_in_a_single_pass() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(0));
write_const(&mut c, idx, 1);
c.write_op(Opcode::Pop, 1);
c.write_op(Opcode::Dup, 2);
c.write_op(Opcode::Pop, 2);
peephole(&mut c);
assert!(c.code.is_empty(), "both pairs removed -> empty chunk");
assert_eq!(c.source_lines.len(), 0);
}
#[test]
fn peephole_is_idempotent_within_a_single_pass() {
let mut c = Chunk::new();
let t = c.add_constant(ConstValue::Bool(false));
write_const(&mut c, t, 1);
write_jump(&mut c, Opcode::JumpIfFalse, 7, 1);
c.write_op(Opcode::Not, 2);
write_jump(&mut c, Opcode::JumpIfFalse, 4, 2);
c.write_op(Opcode::Halt, 3);
peephole(&mut c);
let after_first = c.code.clone();
let lines_first = c.source_lines.clone();
peephole(&mut c);
assert_eq!(c.code, after_first, "second pass changed the bytes");
assert_eq!(c.source_lines, lines_first, "second pass changed the lines");
}
#[test]
fn match_variant_offset_is_fixed_up_after_a_removal() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(0));
c.write_op(Opcode::MatchVariant, 1);
c.write_u16(0, 1); c.write_i16(4, 1); write_const(&mut c, idx, 1);
c.write_op(Opcode::Pop, 1);
c.write_op(Opcode::Halt, 2);
assert_eq!(c.code.len(), 10);
peephole(&mut c);
assert_eq!(c.code.len(), 6);
assert_eq!(c.code[0], Opcode::MatchVariant as u8);
assert_eq!(c.read_i16(3), 0);
assert_eq!(c.code[5], Opcode::Halt as u8);
}
#[test]
fn lockstep_invariant_holds_after_rewrites() {
let mut c = Chunk::new();
let n = c.add_constant(ConstValue::I64(1));
let t = c.add_constant(ConstValue::Bool(true));
write_const(&mut c, n, 1);
c.write_op(Opcode::Pop, 1);
write_const(&mut c, t, 2);
write_jump(&mut c, Opcode::JumpIfFalse, 3, 2);
c.write_op(Opcode::Halt, 3);
peephole(&mut c);
assert_eq!(
c.code.len(),
c.source_lines.len(),
"code and source_lines drifted"
);
}
#[test]
fn source_lines_stay_within_the_program_line_range_after_optimize() {
let src = "fn main() is io {\n\
let x = 1 + 2\n\
println(\"{x}\")\n\
}";
let line_count = src.lines().count() as u32;
let mut program = compile(src);
program.optimize();
for (ci, chunk) in program.chunks.iter().enumerate() {
for (bi, &line) in chunk.source_lines.iter().enumerate() {
assert!(
(1..=line_count).contains(&line),
"chunk {ci} byte {bi}: source line {line} outside 1..={line_count}"
);
}
assert_eq!(chunk.code.len(), chunk.source_lines.len());
}
}
#[test]
fn six_bundled_examples_optimize_without_regression() {
for name in [
"hello",
"fibonacci",
"effects",
"pattern-matching",
"pipeline",
"defer-demo",
] {
let path = format!(
"{}/../../playground/public/examples/{name}.qala",
env!("CARGO_MANIFEST_DIR"),
);
let src = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("read {path}: {e}"));
let tokens =
Lexer::tokenize(&src).unwrap_or_else(|e| panic!("{name}.qala: lex: {e:?}"));
let ast =
Parser::parse(&tokens).unwrap_or_else(|e| panic!("{name}.qala: parse: {e:?}"));
let (typed, terrors, _) = check_program(&ast, &src);
assert!(terrors.is_empty(), "{name}.qala: typecheck: {terrors:?}");
let mut program = crate::codegen::compile_program(&typed, &src)
.unwrap_or_else(|e| panic!("{name}.qala: compile: {e:?}"));
program.optimize();
assert!(
!program.chunks.is_empty(),
"{name}.qala: no chunks after optimize"
);
for (i, chunk) in program.chunks.iter().enumerate() {
assert!(
!chunk.code.is_empty(),
"{name}.qala: chunk {i} empty after optimize"
);
assert_eq!(
chunk.code.len(),
chunk.source_lines.len(),
"{name}.qala: chunk {i} lockstep broken after optimize"
);
let last = *chunk.code.last().expect("non-empty chunk");
assert!(
last == Opcode::Return as u8 || last == Opcode::Halt as u8,
"{name}.qala: chunk {i} lost its terminator"
);
}
assert!(
!program.disassemble().is_empty(),
"{name}.qala: disassembly empty after optimize"
);
}
}
#[test]
fn optimize_is_byte_deterministic_across_runs() {
let src = "fn add(a: i64, b: i64) -> i64 { a + b }\n\
fn main() is io {\n\
let r = add(1, 2)\n\
println(\"{r}\")\n\
}";
let mut first = compile(src);
let mut second = compile(src);
first.optimize();
second.optimize();
assert_eq!(first.chunks.len(), second.chunks.len());
for (a, b) in first.chunks.iter().zip(second.chunks.iter()) {
assert_eq!(a.code, b.code, "optimized code differs across runs");
assert_eq!(
a.source_lines, b.source_lines,
"optimized source map differs across runs"
);
}
assert_eq!(
first.disassemble(),
second.disassemble(),
"optimized disassembly is non-deterministic"
);
}
#[test]
fn peephole_leaves_an_already_tight_chunk_unchanged() {
let mut c = Chunk::new();
let idx = c.add_constant(ConstValue::I64(1));
write_const(&mut c, idx, 1);
c.write_op(Opcode::Return, 1);
let before = c.code.clone();
let lines = c.source_lines.clone();
peephole(&mut c);
assert_eq!(c.code, before);
assert_eq!(c.source_lines, lines);
}
#[test]
fn peephole_on_an_empty_chunk_is_a_no_op() {
let mut c = Chunk::new();
peephole(&mut c);
assert!(c.code.is_empty());
assert!(c.source_lines.is_empty());
}
}