JIT-compiled oplossing voor de Infi Advent of Code 2024 puzzel

stackmachine.rs

/*
Deze oplossing implementeert een JIT compiler met Cranelift, een compiler backend
oorspronkelijk gemaakt voor wasmtime.

De input wordt naar machine code gecompiled en daarna uitgevoerd. In principe kan
het ieder programma aan in de taal beschreven in de Infi opdracht, maar er worden
speciale optimizations gebruikt om nog efficiënter met de echte input om te gaan.
Op mijn PC worden alle blokjes berekend in 900µs, dus 33 nanoseconden per blokje.
Daar gaat wel 450ms compile time aan vooraf.

Er zijn o.a. command line options voor optimizations, om de gegenereerde assembly
te zien, en om een visualisatie van de wolken te bekijken. Probeer deze commands:

cargo run --release -- --visualize --verbose input.txt
cargo run --release -- --disassemble input.txt
cargo run --release -- --skip-midlevel --disassemble input.txt

Probeer ook ongeldige code als input. En bekijk dit programma met --visualize:

    push X
    push Y
    push Z
    push Z
    push -57
    add
    add
    add
    add
    jmpos 2
    push 0
    ret
    push X
    push Y
    push Z
    push Z
    push -59
    add
    add
    add
    add
    jmpos 2
    push 1
    ret
    push 0
    ret

Omdat de input language in het Engels is is het programma ook in het Engels. Maar
de comments zijn allemaal Nederlands.

Cargo.toml:

[package]
name = "infi-aoc24"
version = "0.1.0"
edition = "2021"

[dependencies]
anyhow = "1.0.94"
cranelift-codegen = "0.114.0"
cranelift-frontend = "0.114.0"
cranelift-jit = "0.114.0"
cranelift-module = "0.114.0"
cranelift-native = "0.114.0"
lexopt = "0.3.0"
libc = "0.2.169"

[target.'cfg(unix)'.dependencies]
libc = "0.2.169"

[target.'cfg(windows)'.dependencies]
windows-sys = { version = "0.59.0", features = [
    "Win32_System_SystemInformation",
] }
*/

use std::{
    fmt::Write as _,
    mem::{replace, transmute},
    ops::Range,
    path::{Path, PathBuf},
    time::Instant,
};

use anyhow::{anyhow, bail, Context, Result};
use cranelift_codegen::{
    entity::EntityRef,
    ir::{condcodes::IntCC, types::I32, AbiParam, InstBuilder, MemFlags, Type, Value},
    settings::{self, Configurable},
    verify_function,
};
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{Linkage, Module};

#[derive(Debug)]
struct Args {
    filename: PathBuf,
    visualize: bool,
    stack_size: usize,
    cranelift_settings: Vec<(String, Option<String>)>,
    optimize_midlevel: bool,
    disassemble: bool,
    verbosity: u8,
}

static HELP: &str = "\
Usage: {bin_name} [OPTIONS] FILENAME

Options:
    -V, --visualize              Show an animation of the cloud space
    -s, --stack-size SIZE        Set the stack size (default: 1024)
    -C, --cranelift KEY[=VALUE]  Set a cranelift setting
    -l, --list-cranelift         List all cranelift settings
    -O, --optimize               Enable cranelift optimizations
    -0, --skip-midlevel          Disable mid-level optimizations
    -d, --disassemble            Print assembly of the generated code
    -v, --verbose                Print timing information
";

fn main() -> Result<()> {
    let args = parse_args()?;
    let before_frontend = Instant::now();

    // Alvast in vogelvlucht wat we zometeen gaan zien.

    // Eerst moeten we de input natuurlijk parsen.
    let instructions = parse(args.filename.as_ref())?;

    // Daarna moeten we de geparsede input verwerken om structuur aan te brengen.
    let mut lowered = lower(&instructions)?;

    // We herkennen bepaalde patronen en maken ze efficiënter. Dit is niet nodig,
    // maar zowel het compilen als het uitvoeren wordt er sneller van.
    if args.optimize_midlevel {
        lowered.iter_mut().for_each(optimize);
    }

    if args.verbosity > 0 {
        println!("Frontend time: {:.2?}", before_frontend.elapsed());
    }

    // We vertalen naar Cranelift instructies en krijgen een function pointer.
    let before_cranelift = Instant::now();
    let program = compile(&lowered, &args)?;
    if args.verbosity > 0 {
        println!("Cranelift time: {:.2?}", before_cranelift.elapsed());
    }

    // We zetten beschermd geheugen klaar voor het programma. Iedere waarde op de
    // stack is 4 bytes.
    let mut memory = GuardedMemory::new(args.stack_size * 4)?;

    // En nu kunnen we het programma eindelijk uitvoeren.
    let before_run = Instant::now();
    let mut part1 = 0;
    let mut clouded = Grid::default();
    for coord in coords() {
        let [x, y, z] = coord;
        let result = unsafe { program(memory.memory(), x as i32, y as i32, z as i32) };
        part1 += result;
        clouded[x][y][z] = result > 0;
    }

    if args.verbosity > 0 {
        let runtime = before_run.elapsed();
        println!(
            "Runtime: {:.2?} ({:?}/coord)",
            runtime,
            runtime / NUM_COORDS as u32,
        );
    }

    if args.visualize {
        visualize(&clouded);
    }

    println!("Part 1: {part1}");

    // We markeren blokjes als verwerkt door ze op `false` te zetten. Blokjes die
    // nog niet verwerkt waren worden het begin van een floodfill.
    let mut num_clouds = 0;
    let mut this_cloud = Vec::new();
    for coord in coords() {
        let [x, y, z] = coord;
        if !replace(&mut clouded[x][y][z], false) {
            continue;
        }
        num_clouds += 1;

        this_cloud.push(coord);
        while let Some(coord) = this_cloud.pop() {
            for neighbor in neighbors(coord) {
                let [x, y, z] = neighbor;
                if replace(&mut clouded[x][y][z], false) {
                    this_cloud.push(neighbor);
                }
            }
        }
    }

    println!("Part 2: {num_clouds}");

    Ok(())
}

type Coord = [usize; 3];
type Grid = [[[bool; GRID_SIZE]; GRID_SIZE]; GRID_SIZE];

const GRID_SIZE: usize = 30;
const GRID_RANGE: Range<usize> = 0..GRID_SIZE;
const NUM_COORDS: usize = GRID_SIZE * GRID_SIZE * GRID_SIZE;

/// Alle coördinaten binnen het 30×30×30 gebied.
fn coords() -> impl Iterator<Item = Coord> {
    GRID_RANGE.flat_map(|x| GRID_RANGE.flat_map(move |y| GRID_RANGE.map(move |z| [x, y, z])))
}

/// Alle buurblokjes van een coördinaat (binnen het 30×30×30 gebied).
fn neighbors([x, y, z]: Coord) -> impl Iterator<Item = Coord> {
    [
        [x.checked_sub(1), Some(y), Some(z)],
        [Some(x + 1), Some(y), Some(z)],
        [Some(x), y.checked_sub(1), Some(z)],
        [Some(x), Some(y + 1), Some(z)],
        [Some(x), Some(y), z.checked_sub(1)],
        [Some(x), Some(y), Some(z + 1)],
    ]
    .into_iter()
    .filter_map(|[x, y, z]| {
        (x? < GRID_SIZE && y? < GRID_SIZE && z? < GRID_SIZE).then_some([x?, y?, z?])
    })
}

fn parse_args() -> Result<Args> {
    use lexopt::prelude::*;

    let mut filename = None;
    let mut visualize = false;
    let mut stack_size = 1024;
    let mut cranelift_settings = Vec::new();
    let mut optimize_midlevel = true;
    let mut disassemble = false;
    let mut list_cranelift = false;
    let mut verbosity: u8 = 0;

    let mut parser = lexopt::Parser::from_env();
    while let Some(arg) = parser.next()? {
        match arg {
            Short('V') | Long("visualize") => visualize = true,
            Short('s') | Long("stack-size") => stack_size = parser.value()?.parse()?,
            Short('C') | Long("cranelift") => {
                let setting = parser.value()?.string()?;
                if let Some((key, value)) = setting.split_once('=') {
                    cranelift_settings.push((key.into(), Some(value.into())));
                } else {
                    cranelift_settings.push((setting, None));
                }
            }
            Short('O') | Long("optimize") => {
                cranelift_settings.push(("opt_level".into(), Some("speed".into())));
            }
            Short('0') | Long("skip-midlevel") => optimize_midlevel = false,
            Short('d') | Long("disassemble") => disassemble = true,
            Short('v') | Long("verbose") => verbosity = verbosity.saturating_add(1),
            Value(value) if filename.is_none() => filename = Some(value.into()),
            Short('l') | Long("list-cranelift") => list_cranelift = true,
            Long("help") => usage(&parser, 0),
            _ => return Err(arg.unexpected().into()),
        }
    }

    // Hiermee wachten we tot na de loop zodat --help altijd voorrang heeft.
    if list_cranelift {
        list_cranelift_settings();
    }

    Ok(Args {
        filename: filename.unwrap_or_else(|| usage(&parser, 1)),
        visualize,
        stack_size,
        cranelift_settings,
        optimize_midlevel,
        disassemble,
        verbosity,
    })
}

fn usage(parser: &lexopt::Parser, exitcode: i32) -> ! {
    let bin_name = parser.bin_name().unwrap_or(env!("CARGO_PKG_NAME"));
    println!("{}", HELP.replace("{bin_name}", bin_name));
    std::process::exit(exitcode);
}

fn list_cranelift_settings() -> ! {
    let width = settings::builder()
        .iter()
        .map(|s| s.name.len())
        .max()
        .unwrap();
    for setting in settings::builder().iter() {
        println!(" {:>width$}   {}", setting.name, setting.description);
    }
    std::process::exit(0);
}

// Een animatietje van het wolkendek, laag voor laag.
fn visualize(clouded: &Grid) {
    let mut frame = String::new();

    for z in GRID_RANGE {
        write!(frame, "\nz = {z}\n").unwrap();
        frame.push_str("   012345678901234567890123456789\n");
        for y in GRID_RANGE {
            write!(frame, "{y:>2} ").unwrap();
            for x in GRID_RANGE {
                if clouded[x][y][z] {
                    frame.push('#');
                } else {
                    frame.push(' ');
                }
            }
            frame.push('\n');
        }

        println!("{frame}");
        std::thread::sleep(std::time::Duration::from_millis(100));
        frame.clear();

        // Beweeg de cursor in de volgende frame terug naar de startpositie.
        frame.push_str("\r\x1B[34A");
    }
}

/// Alle mogelijke instructies zoals beschreven in de opdracht.
///
/// Later wijken we hier wat van af.
#[derive(Clone, Copy, Debug)]
enum Instruction {
    PushLiteral(i32),
    PushSymbol(Symbol),
    Add,
    JumpOffset(i32),
    Return,
}

#[derive(Clone, Copy, Debug)]
enum Symbol {
    X,
    Y,
    Z,
}

fn parse(filename: &Path) -> Result<Vec<Instruction>> {
    let text = std::fs::read_to_string(filename)?;

    // Parsen is simpel, maar we doen ons best om goede errors te geven.
    fn parse_instruction(line: &str) -> Result<Instruction> {
        let mut line = line.split_whitespace();
        let Some(instruction) = line.next() else {
            bail!("Empty line");
        };
        let operand = line.next();
        if let Some(rest) = line.next() {
            bail!("Superfluous argument {rest:?}");
        }

        Ok(match (instruction, operand) {
            ("add", None) => Instruction::Add,
            ("ret", None) => Instruction::Return,
            ("push", Some("X" | "x")) => Instruction::PushSymbol(Symbol::X),
            ("push", Some("Y" | "y")) => Instruction::PushSymbol(Symbol::Y),
            ("push", Some("Z" | "z")) => Instruction::PushSymbol(Symbol::Z),
            ("push", Some(operand)) => Instruction::PushLiteral(operand.parse()?),
            ("jmpos", Some(operand)) => Instruction::JumpOffset(operand.parse()?),
            ("add" | "ret", Some(operand)) => bail!("Superfluous argument {:?}", operand),
            ("push" | "jmpos", None) => bail!("Missing argument to '{instruction}'"),
            _ => bail!("Unknown instruction {instruction:?}"),
        })
    }

    let mut instructions = Vec::new();
    for (i, line) in text.lines().enumerate() {
        let parsed =
            parse_instruction(line).with_context(|| format!("Line {}: {}", i + 1, line))?;
        instructions.push(parsed);
    }

    if !matches!(instructions.last(), Some(Instruction::Return)) {
        bail!("Missing return instruction at end of program");
    }

    Ok(instructions)
}

// Hier is diepere uitleg nodig.
//
// Cranelift werkt op basis van "blocks" van instructies. Control flow kan alleen
// naar het begin van een blok springen en nooit naar het midden, zelfs niet naar
// een later deel van hetzelfde blok. Dat werkt redelijk logisch met higher-level
// talen (maar niet 1:1, want blocks kunnen niet genest worden). Voor de jumps in
// onze stack machine werkt het minder handig.
//
// Alle jumps zijn gelukkig statisch te bepalen: ze gaan naar een vast adres. Dus
// we kunnen het hele programma doorlopen en de bestemmingen van jumps bijhouden.
// Bij iedere bestemming van een jump begint een nieuw blok.
//
// De instructie na een `jmpos` moet ook het begin van een blok worden, want deze
// instructie heeft in feite twee bestemmingen: één voor een positieve input, één
// voor een negatieve.
//
// Zodra alle instructies in blokken zijn opgedeeld vertalen we de jumps zodat ze
// naar een blok wijzen in plaats van naar een adres.

#[derive(Clone, Copy, Debug)]
enum LowLevelInstruction {
    PushLiteral(i32),
    PushSymbol(Symbol),
    Add,
    JumpIfPositive(usize),
    Return,
    // Deze twee instructies komen niet voor in de opdracht.
    // We genereren ze in `fn optimize()` verderop.
    JumpIfGeq {
        symbol: Symbol,
        comparator: i32,
        to_block_idx: usize,
    },
    ReturnLiteral(i32),
}

impl LowLevelInstruction {
    fn is_control_flow(&self) -> bool {
        match self {
            LowLevelInstruction::PushLiteral(..)
            | LowLevelInstruction::PushSymbol(..)
            | LowLevelInstruction::Add => false,
            LowLevelInstruction::JumpIfPositive(..)
            | LowLevelInstruction::JumpIfGeq { .. }
            | LowLevelInstruction::Return
            | LowLevelInstruction::ReturnLiteral(..) => true,
        }
    }
}

type Block = Vec<LowLevelInstruction>;

fn lower(instructions: &[Instruction]) -> Result<Vec<Block>> {
    let mut block_boundaries: Vec<usize> = vec![0];

    // Eerst zoeken we alle bestemmingen.
    for (addr, &instruction) in instructions.iter().enumerate() {
        if let Instruction::JumpOffset(offset) = instruction {
            let offset: isize = offset.try_into().expect("infallible on 32bit+ machines");
            // De program counter gaat na de jump nog impliciet met 1 omhoog.
            let offset = offset + 1;
            let Some(target) = addr.checked_add_signed(offset) else {
                bail!("Jump with offset {offset} targets a negative address");
            };
            if target >= instructions.len() {
                bail!("Jump with offset {offset} targets beyond the end of the program");
            };
            block_boundaries.push(target);
            block_boundaries.push(addr + 1);
        }
    }

    block_boundaries.sort_unstable();
    block_boundaries.push(instructions.len());
    block_boundaries.dedup();

    // Nu kunnen we de instructies opdelen en omzetten.
    let mut blocks = Vec::new();
    for window in block_boundaries.windows(2) {
        let &[start, end] = window else {
            unreachable!()
        };
        let mut block = Vec::new();
        for (addr, &instruction) in (start..).zip(&instructions[start..end]) {
            match instruction {
                Instruction::JumpOffset(offset) => {
                    let offset: isize = offset.try_into().unwrap();
                    let offset = offset + 1;
                    let target = addr.checked_add_signed(offset).unwrap();
                    let block_idx = block_boundaries.binary_search(&target).unwrap();
                    block.push(LowLevelInstruction::JumpIfPositive(block_idx));
                }
                Instruction::PushLiteral(n) => block.push(LowLevelInstruction::PushLiteral(n)),
                Instruction::PushSymbol(symbol) => {
                    block.push(LowLevelInstruction::PushSymbol(symbol))
                }
                Instruction::Add => block.push(LowLevelInstruction::Add),
                Instruction::Return => {
                    block.push(LowLevelInstruction::Return);
                    // Cranelift vindt het niet leuk als de code in een block nog
                    // door gaat na een return.
                    // De rest van dit blok kan blijkbaar niet uitgevoerd worden,
                    // want er waren geen jumps naartoe. Dus we kunnen het net zo
                    // goed negeren.
                    break;
                }
            }
        }
        blocks.push(block);
    }

    Ok(blocks)
}

// Nu komt een speciaal trucje. De input voor de opdracht bestaat in feite al uit
// blokken, altijd in een van deze twee vormen:
//
//   push SYMBOL
//   push NUMBER
//   add
//   jmpos ADDRESS
//
// en:
//
//   push NUMBER
//   ret
//
// Deze blokken hebben de mooie eigenschap dat ze geen stack nodig hebben. Nummer
// een pusht twee waardes op de stack maar haalt ze er allebei weer van af. En de
// tweede vorm pusht en returnt zonder bestaande waardes van de stack te lezen.
//
// We voeren ze uit zonder het geheugen van de stack machine aan te raken. Op die
// manier hoeft Cranelift minder werk te doen. En de processor hoeft bijna alleen
// maar `cmpl` en `jnl` instructies uit te voeren.
//
// Het is dus eigenlijk ook mogelijk om het programma te vertalen naar een simpel
// diagram van "{X,Y,Z} >= n" beslissingen. Maar het mooie van deze aanpak is dat
// we andere inputs niet verliezen.
//
// Omdat blocks netjes van elkaar geïsoleerd zijn kunnen ze makkelijk individueel
// aangepast worden zonder daarmee andere blocks dwars te zitten.

fn optimize(block: &mut Block) {
    use LowLevelInstruction::*;

    match block[..] {
        [PushSymbol(symbol), PushLiteral(number), Add, JumpIfPositive(to_block_idx)] => {
            block[0] = JumpIfGeq {
                symbol,
                comparator: -number,
                to_block_idx,
            };
            block.truncate(1);
        }
        [PushLiteral(number), Return] => {
            block[0] = ReturnLiteral(number);
            block.truncate(1);
        }
        _ => (),
    }
}

// Cranelift spuugt straks een function pointer uit die we kunnen aanroepen alsof
// het een normale C functie is.
//
// Het eerste argument is een pointer naar het geheugen dat voor de stack machine
// gebruikt mag worden. De stack groeit naar boven, en iedere waarde is 32 bits.

type Executable = unsafe extern "C" fn(*mut (), i32, i32, i32) -> i32;

struct Compiler<'a> {
    builder: FunctionBuilder<'a>,
    // Deze variable bevat een pointer naar de bovenkant van de stack, oftewel de
    // eerstvolgende lege plek. Deze schuift op als de stack groeit of krimpt.
    stack_head: Variable,
    // Deze variables bevatten X, Y, en Z. Ze veranderen niet.
    params: [Variable; 3],
    // Dit zijn de blocks zoals Cranelift ze ziet. We maken ze allemaal meteen en
    // vullen ze een voor een.
    blocks: Vec<cranelift_codegen::ir::Block>,
    pointer_type: Type,
}

impl<'a> Compiler<'a> {
    fn new(mut builder: FunctionBuilder<'a>, num_blocks: usize, pointer_type: Type) -> Self {
        let stack_head = Variable::new(0);
        builder.declare_var(stack_head, pointer_type);
        let arg_x = Variable::new(1);
        builder.declare_var(arg_x, I32);
        let arg_y = Variable::new(2);
        builder.declare_var(arg_y, I32);
        let arg_z = Variable::new(3);
        builder.declare_var(arg_z, I32);

        // De entry block slaat de arguments op in variables en jumpt daarna naar
        // het echte begin. (De arguments worden in `fn compile` gedefiniëerd.)
        let entry_block = builder.create_block();

        let blocks: Vec<_> = (0..num_blocks).map(|_| builder.create_block()).collect();

        builder.append_block_params_for_function_params(entry_block);
        builder.switch_to_block(entry_block);
        builder.seal_block(entry_block);
        builder.def_var(stack_head, builder.block_params(entry_block)[0]);
        builder.def_var(arg_x, builder.block_params(entry_block)[1]);
        builder.def_var(arg_y, builder.block_params(entry_block)[2]);
        builder.def_var(arg_z, builder.block_params(entry_block)[3]);

        builder.ins().jump(blocks[0], &[]);

        Compiler {
            builder,
            stack_head,
            params: [arg_x, arg_y, arg_z],
            blocks,
            pointer_type,
        }
    }

    fn finish(mut self) {
        self.builder.seal_all_blocks();
        self.builder.finalize();
    }

    // De code in deze method is nog aardig high-level, met abstractions voor het
    // gebruiken van de stack.
    fn add_block(&mut self, block: &[LowLevelInstruction], block_idx: usize) {
        self.builder.switch_to_block(self.blocks[block_idx]);

        for &instruction in block {
            match instruction {
                LowLevelInstruction::PushLiteral(n) => {
                    let value = self.load_literal(n);
                    self.push(value);
                }
                LowLevelInstruction::PushSymbol(symbol) => {
                    let value = self.load_symbol(symbol);
                    self.push(value);
                }
                LowLevelInstruction::Add => {
                    let x = self.pop();
                    let y = self.pop();
                    let result = self.builder.ins().iadd(x, y);
                    self.push(result);
                }
                LowLevelInstruction::JumpIfPositive(to_block_idx) => {
                    let value = self.pop();
                    let zero = self.load_literal(0);
                    let is_not_negative =
                        self.builder
                            .ins()
                            .icmp(IntCC::SignedGreaterThanOrEqual, value, zero);
                    // Een voorwaardelijke jump heeft twee bestemmingen. We geven
                    // expliciet aan dat we voor een negatieve waarde verder gaan
                    // in het volgende blok.
                    self.builder.ins().brif(
                        is_not_negative,
                        self.blocks[to_block_idx],
                        &[],
                        self.blocks[block_idx + 1],
                        &[],
                    );
                }
                LowLevelInstruction::JumpIfGeq {
                    symbol,
                    comparator,
                    to_block_idx,
                } => {
                    let value = self.load_symbol(symbol);
                    let comparator = self.load_literal(comparator);
                    let is_geq =
                        self.builder
                            .ins()
                            .icmp(IntCC::SignedGreaterThanOrEqual, value, comparator);
                    self.builder.ins().brif(
                        is_geq,
                        self.blocks[to_block_idx],
                        &[],
                        self.blocks[block_idx + 1],
                        &[],
                    );
                }
                LowLevelInstruction::Return => {
                    let value = self.pop();
                    self.builder.ins().return_(&[value]);
                }
                LowLevelInstruction::ReturnLiteral(number) => {
                    let value = self.load_literal(number);
                    self.builder.ins().return_(&[value]);
                }
            }
        }
        if !block.last().unwrap().is_control_flow() {
            self.builder.ins().jump(self.blocks[block_idx + 1], &[]);
        }
    }

    fn load_symbol(&mut self, symbol: Symbol) -> Value {
        self.builder.use_var(match symbol {
            Symbol::X => self.params[0],
            Symbol::Y => self.params[1],
            Symbol::Z => self.params[2],
        })
    }

    // Dit is de kern van de stack machine. Instructions produceren values die we
    // weer doorgeven aan andere instructions. Zelfs voor een constante moeten er
    // instructions gebruikt worden.

    fn load_literal(&mut self, number: i32) -> Value {
        self.builder.ins().iconst(I32, i64::from(number))
    }

    fn push(&mut self, value: Value) {
        // Schrijf een waarde naar `stack_head`...
        let cur_head = self.builder.use_var(self.stack_head);
        self.builder
            .ins()
            .store(Self::memory_flags(), value, cur_head, 0);

        // ...en tel 4 bytes bij `stack_head` op.
        let value_width = self.value_width();
        let new_head = self.builder.ins().iadd(cur_head, value_width);
        self.builder.def_var(self.stack_head, new_head);
    }

    fn pop(&mut self) -> Value {
        // Trek 4 bytes van `stack_head` af...
        let cur_head = self.builder.use_var(self.stack_head);
        let value_width = self.value_width();
        let new_head = self.builder.ins().isub(cur_head, value_width);
        self.builder.def_var(self.stack_head, new_head);

        // ...en lees een waarde uit (het nieuwe adres in) `stack_head`.
        self.builder
            .ins()
            .load(I32, Self::memory_flags(), new_head, 0)
    }

    fn value_width(&mut self) -> Value {
        self.builder.ins().iconst(self.pointer_type, 4)
    }

    fn memory_flags() -> MemFlags {
        let mut flags = MemFlags::new();
        flags.set_aligned();
        // Nog een flag is `flags.set_notrap()`. Die zou Cranelift vertellen geen
        // rekening te houden met segfaults, en maakt sommige extra optimizations
        // mogelijk, als Cranelift kan bewijzen dat het niet nodig is om geheugen
        // daadwerkelijk aan te raken. Maar dat willen we niet: we willen zo snel
        // mogelijk een segfault. Zie verderop.
        flags
    }
}

// Cranelift heeft nogal wat boilerplate nodig.
// De meeste interessante dingen hebben we hierboven al gehad.

fn compile(blocks: &[Block], args: &Args) -> Result<Executable> {
    let mut config = settings::builder();
    for (key, value) in &args.cranelift_settings {
        match value {
            Some(value) => config.set(key, value)?,
            None => config.enable(key)?,
        }
    }

    let flags = settings::Flags::new(config);
    let isa_builder =
        cranelift_native::builder().map_err(|msg| anyhow!("unsupported host machine: {msg}"))?;
    let isa = isa_builder.finish(flags.clone())?;

    let jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
    let mut module = JITModule::new(jit_builder);
    let mut module_context = module.make_context();

    if args.disassemble {
        module_context.want_disasm = true;
    }

    let pointer_type = module.target_config().pointer_type();

    let sig = &mut module_context.func.signature;
    sig.params.push(AbiParam::new(pointer_type)); // Geheugen
    sig.params.push(AbiParam::new(I32)); // X
    sig.params.push(AbiParam::new(I32)); // Y
    sig.params.push(AbiParam::new(I32)); // Z
    sig.returns.push(AbiParam::new(I32)); // Resultaat

    let mut ctx = FunctionBuilderContext::new();
    let builder = FunctionBuilder::new(&mut module_context.func, &mut ctx);

    let mut compiler = Compiler::new(builder, blocks.len(), pointer_type);
    for (block_idx, instructions) in blocks.iter().enumerate() {
        compiler.add_block(instructions, block_idx);
    }
    compiler.finish();

    verify_function(&module_context.func, &flags)?;

    let func_id = module.declare_function(
        "stackmachine",
        Linkage::Export,
        &module_context.func.signature,
    )?;
    module.define_function(func_id, &mut module_context)?;
    module.finalize_definitions()?;

    if let Some(ref asm) = module_context.compiled_code().unwrap().vcode {
        println!("{asm}");
    }

    let func_ptr = module.get_finalized_function(func_id);

    Ok(unsafe { transmute::<*const u8, Executable>(func_ptr) })
}

// Nu moeten we het programma nog uitvoeren!
//
// Het is niet zo maar veilig om de code te draaien. Neem bijvoorbeeld:
//
//   add
//   ret
//
// Dit programma is niet geldig, want het voert `add` uit op een lege stack. Maar
// onze executable controleert daar niet op. Die zou geheugen dat voor het stack-
// geheugen ligt lezen en bewerken, met memory corruption als gevolg. Zo kan veel
// stuk gaan. Zelfs arbitrary code execution is niet uitgesloten.
//
// Ook de bovenkant van het stack-geheugen is in gevaar. Stel:
//
//   push 1
//   push 1
//   jmpos -3
//   ret
//
// We hebben wel de garantie dat als het mis gaat het programma vlak voor of vlak
// na het geheugen schrijft. Het kan niet opeens miljoenen bytes verder springen.
//
// En dat betekent dat we memory protection kunnen inzetten. Geheugen is verdeeld
// in pages, normaal gesproken bestaande uit 4096 bytes. Het programma krijgt van
// ons een aantal pages, met daar omheen twee pages die beschermd zijn. Zo dus:
//
//                ↓
//   [XXXXXXXXXX][          ][XXXXXXXXXX]
//    PROT_NONE   PROT_WRITE  PROT_NONE
//
// Het programma mag alles doen met het geheugen dat PROT_WRITE is. Maar geheugen
// dat PROT_NONE is is beschermd: bij schrijven of lezen grijpen de MMU en kernel
// in met een segfault.
//
// Unix en Windows hebben allebei een eigen setje functies hiervoor. Ze lijken op
// elkaar maar zijn net wat anders. Het is natuurlijk ondenkbaar om niet portable
// te zijn (ik weet niet wat de elfjes gebruiken) dus dat trekken we eerst recht.

#[cfg(unix)]
mod mmap {
    use std::io::{Error, Result};

    use libc::{MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, PROT_NONE, PROT_WRITE, _SC_PAGESIZE};

    // Pages zijn meestal 4KiB, maar groter op sommige moderne systemen.
    pub fn page_size() -> usize {
        unsafe { libc::sysconf(_SC_PAGESIZE) as usize }
    }

    // Hiermee maken we een lap beschermd geheugen...
    pub fn map_prot_none(len: usize) -> Result<*mut ()> {
        let addr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                len,
                PROT_NONE,
                MAP_PRIVATE | MAP_ANONYMOUS,
                -1,
                0,
            )
        };
        if addr == MAP_FAILED {
            return Err(Error::last_os_error());
        }
        Ok(addr.cast())
    }

    // ...hiermee halen we de bescherming van een sub-lap af...
    pub unsafe fn protect_write(start: *mut (), len: usize) -> Result<()> {
        let result = unsafe { libc::mprotect(start.cast(), len, PROT_WRITE) };
        if result != 0 {
            return Err(Error::last_os_error());
        }
        Ok(())
    }

    // ...en hiermee geven we de lap weer terug aan het besturingssysteem.
    pub unsafe fn unmap(addr: *mut (), len: usize) {
        let status = unsafe { libc::munmap(addr.cast(), len) };
        debug_assert_eq!(status, 0);
    }
}

#[cfg(windows)]
mod mmap {
    use std::{
        io::{Error, Result},
        mem::MaybeUninit,
    };

    use windows_sys::Win32::System::{
        Memory::{
            VirtualAlloc, VirtualFree, VirtualProtect, MEM_COMMIT, MEM_RELEASE, MEM_RESERVE,
            PAGE_NOACCESS, PAGE_READWRITE,
        },
        SystemInformation::GetSystemInfo,
    };

    pub fn page_size() -> usize {
        unsafe {
            let mut system_info = MaybeUninit::uninit();
            GetSystemInfo(system_info.as_mut_ptr());
            system_info.assume_init().dwPageSize as usize
        }
    }

    pub fn map_prot_none(len: usize) -> Result<*mut ()> {
        let addr = unsafe {
            VirtualAlloc(
                std::ptr::null(),
                len,
                MEM_RESERVE | MEM_COMMIT,
                PAGE_NOACCESS,
            )
        };
        if addr.is_null() {
            return Err(Error::last_os_error());
        }
        Ok(addr.cast())
    }

    pub unsafe fn protect_write(start: *mut (), len: usize) -> Result<()> {
        let mut oldprotect = MaybeUninit::uninit();
        let result =
            unsafe { VirtualProtect(start.cast(), len, PAGE_READWRITE, oldprotect.as_mut_ptr()) };
        if result == 0 {
            return Err(Error::last_os_error());
        }
        Ok(())
    }

    pub unsafe fn unmap(addr: *mut (), _len: usize) {
        let result = unsafe { VirtualFree(addr.cast(), 0, MEM_RELEASE) };
        debug_assert_ne!(result, 0);
    }
}

// Nu knopen we het aan elkaar en pakken we het in voor `fn main()`.

struct GuardedMemory {
    addr: *mut (),
    len: usize,
    usable_start: *mut (),
}

impl GuardedMemory {
    fn new(bytes: usize) -> Result<Self> {
        let page_size = mmap::page_size();

        let usable_pages = bytes.div_ceil(page_size);

        // 1 beschermde page er voor en 1 er na.
        let total_pages = usable_pages + 2;
        let len = page_size * total_pages;

        // We mappen n + 2 pages met PROT_NONE.
        let addr = mmap::map_prot_none(len).context("mmap")?;

        // We maken het object alvast zodat drop ook draait als mprotect() faalt.
        let usable_start = addr.wrapping_byte_add(page_size);
        let usable_len = page_size * usable_pages;
        let this = Self {
            addr,
            usable_start,
            len,
        };

        // We maken de pages in het midden writable. Nu kan main() aan het werk.
        unsafe {
            mmap::protect_write(usable_start, usable_len).context("mprotect")?;
        }

        Ok(this)
    }

    fn memory(&mut self) -> *mut () {
        self.usable_start
    }
}

// Aan het einde geven we het geheugen terug aan het besturingssysteem.
// En dan zijn we klaar.
impl Drop for GuardedMemory {
    fn drop(&mut self) {
        unsafe {
            mmap::unmap(self.addr, self.len);
        }
    }
}