wai-quantum 0.6.0

//! A small RV32IM encoder and the gate's kernels for a 32-bit RISC-V core that
//! runs with no operating system: each program is built here, as machine words
//! and as the assembly listing they encode (the implementation id's input).
//!
//! The kernels read a parameter block at a base address in the core's memory and
//! write their results beside it ([`Layout`]); a done word tells the host they
//! finished. Every 64-bit value is a pair of words: the product is `mul` (low)
//! and `mulh` (high, signed); the rounding constant's carry is `sltu`; the
//! arithmetic shift is built from word shifts; a row's carries and borrows are
//! `sltu` too.

use std::collections::BTreeMap;
use std::fmt::Write as _;

/// A register, `x0` to `x31`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct R(pub u8);

/// The ABI names, for the listing.
const NAMES: [&str; 32] =
    ["zero", "ra", "sp", "gp", "tp", "t0", "t1", "t2", "s0", "s1", "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10", "s11", "t3", "t4", "t5", "t6"];

impl std::fmt::Display for R {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(NAMES[self.0 as usize])
    }
}

pub const ZERO: R = R(0);
pub const T0: R = R(5);
pub const T1: R = R(6);
pub const T2: R = R(7);
pub const S0: R = R(8);
pub const S1: R = R(9);
pub const A0: R = R(10);
pub const A1: R = R(11);
pub const A2: R = R(12);
pub const A3: R = R(13);
pub const A4: R = R(14);
pub const A5: R = R(15);
pub const A6: R = R(16);
pub const A7: R = R(17);
pub const S2: R = R(18);
pub const S3: R = R(19);
pub const T3: R = R(28);
pub const T4: R = R(29);
pub const T5: R = R(30);
pub const T6: R = R(31);

fn r_type(f7: u32, rs2: R, rs1: R, f3: u32, rd: R, op: u32) -> u32 {
    (f7 << 25) | (u32::from(rs2.0) << 20) | (u32::from(rs1.0) << 15) | (f3 << 12) | (u32::from(rd.0) << 7) | op
}

fn i_type(imm: i32, rs1: R, f3: u32, rd: R, op: u32) -> u32 {
    assert!((-2048..2048).contains(&imm), "an I-type immediate is 12 bits: {imm}");
    (((imm as u32) & 0xfff) << 20) | (u32::from(rs1.0) << 15) | (f3 << 12) | (u32::from(rd.0) << 7) | op
}

/// One branch waiting for its label.
struct Fixup {
    at: usize,
    label: &'static str,
    f3: u32,
    rs1: R,
    rs2: R,
}

/// A program being built: machine words and their listing.
#[derive(Default)]
pub struct Asm {
    words: Vec<u32>,
    lines: Vec<String>,
    labels: BTreeMap<&'static str, usize>,
    fixups: Vec<Fixup>,
}

macro_rules! r_ops {
    ($($name:ident = $f7:literal, $f3:literal;)*) => {$(
        pub fn $name(&mut self, rd: R, rs1: R, rs2: R) -> &mut Self {
            self.emit(r_type($f7, rs2, rs1, $f3, rd, 0x33), format!("{} {rd}, {rs1}, {rs2}", stringify!($name)))
        }
    )*};
}

impl Asm {
    fn emit(&mut self, word: u32, line: String) -> &mut Self {
        self.words.push(word);
        self.lines.push(line);
        self
    }

    r_ops! {
        add = 0x00, 0;
        sub = 0x20, 0;
        sltu = 0x00, 3;
        or = 0x00, 6;
        mul = 0x01, 0;
        mulh = 0x01, 1;
    }

    pub fn addi(&mut self, rd: R, rs1: R, imm: i32) -> &mut Self {
        self.emit(i_type(imm, rs1, 0, rd, 0x13), format!("addi {rd}, {rs1}, {imm}"))
    }

    pub fn slli(&mut self, rd: R, rs1: R, sh: u32) -> &mut Self {
        assert!(sh < 32);
        self.emit(i_type(sh as i32, rs1, 1, rd, 0x13), format!("slli {rd}, {rs1}, {sh}"))
    }

    pub fn srli(&mut self, rd: R, rs1: R, sh: u32) -> &mut Self {
        assert!(sh < 32);
        self.emit(i_type(sh as i32, rs1, 5, rd, 0x13), format!("srli {rd}, {rs1}, {sh}"))
    }

    pub fn srai(&mut self, rd: R, rs1: R, sh: u32) -> &mut Self {
        assert!(sh < 32);
        self.emit(i_type(0x400 | sh as i32, rs1, 5, rd, 0x13), format!("srai {rd}, {rs1}, {sh}"))
    }

    pub fn lw(&mut self, rd: R, off: i32, rs1: R) -> &mut Self {
        self.emit(i_type(off, rs1, 2, rd, 0x03), format!("lw {rd}, {off}({rs1})"))
    }

    pub fn sw(&mut self, rs2: R, off: i32, rs1: R) -> &mut Self {
        assert!((-2048..2048).contains(&off), "an S-type offset is 12 bits: {off}");
        let imm = off as u32 & 0xfff;
        let word = ((imm >> 5) << 25) | (u32::from(rs2.0) << 20) | (u32::from(rs1.0) << 15) | (2 << 12) | ((imm & 0x1f) << 7) | 0x23;
        self.emit(word, format!("sw {rs2}, {off}({rs1})"))
    }

    pub fn lui(&mut self, rd: R, upper: u32) -> &mut Self {
        assert!(upper < 1 << 20);
        self.emit((upper << 12) | (u32::from(rd.0) << 7) | 0x37, format!("lui {rd}, {upper:#x}"))
    }

    /// `rd = value`, by `lui` and `addi` (or `addi` alone).
    pub fn li(&mut self, rd: R, value: u32) -> &mut Self {
        let low = ((value << 20) as i32) >> 20; // the low 12 bits, sign-extended
        let upper = value.wrapping_sub(low as u32) >> 12;
        if upper != 0 {
            self.lui(rd, upper);
            if low != 0 {
                self.addi(rd, rd, low);
            }
            self
        } else {
            self.addi(rd, ZERO, low)
        }
    }

    pub fn label(&mut self, name: &'static str) -> &mut Self {
        assert!(self.labels.insert(name, self.words.len()).is_none(), "label {name} twice");
        self.lines.push(format!("{name}:"));
        self
    }

    fn branch(&mut self, f3: u32, mnemonic: &str, rs1: R, rs2: R, label: &'static str) -> &mut Self {
        self.fixups.push(Fixup { at: self.words.len(), label, f3, rs1, rs2 });
        self.emit(0, format!("{mnemonic} {rs1}, {rs2}, {label}"))
    }

    pub fn beq(&mut self, rs1: R, rs2: R, label: &'static str) -> &mut Self {
        self.branch(0, "beq", rs1, rs2, label)
    }

    pub fn bne(&mut self, rs1: R, rs2: R, label: &'static str) -> &mut Self {
        self.branch(1, "bne", rs1, rs2, label)
    }

    /// A word for the tile, carried inline: the core passes on any word in its
    /// instruction stream whose low two bits are not `0b11`, rotated right by
    /// two, so the stream holds the word rotated left by two. `what` names it
    /// in the listing.
    pub fn tt(&mut self, word: u32, what: &str) -> &mut Self {
        let inline = word.rotate_left(2);
        assert_ne!(inline & 3, 3, "{word:#010x} would be a RISC-V instruction");
        self.emit(inline, format!(".word {inline:#010x}  # {what}"))
    }

    /// `jal zero, label`.
    pub fn j(&mut self, label: &'static str) -> &mut Self {
        self.fixups.push(Fixup { at: self.words.len(), label, f3: u32::MAX, rs1: ZERO, rs2: ZERO });
        self.emit(0, format!("j {label}"))
    }

    /// The machine words, every branch resolved, and the listing.
    pub fn finish(mut self) -> (Vec<u32>, String) {
        for f in std::mem::take(&mut self.fixups) {
            let target = *self.labels.get(f.label).unwrap_or_else(|| panic!("no label {}", f.label));
            let off = (target as i64 - f.at as i64) * 4;
            let o = off as u32;
            self.words[f.at] = if f.f3 == u32::MAX {
                assert!((-(1 << 20)..(1 << 20)).contains(&off), "a jump reaches 1 MiB");
                (((o >> 20) & 1) << 31) | (((o >> 1) & 0x3ff) << 21) | (((o >> 11) & 1) << 20) | (((o >> 12) & 0xff) << 12) | 0x6f
            } else {
                assert!((-(1 << 12)..(1 << 12)).contains(&off), "a branch reaches 4 KiB");
                (((o >> 12) & 1) << 31)
                    | (((o >> 5) & 0x3f) << 25)
                    | (u32::from(f.rs2.0) << 20)
                    | (u32::from(f.rs1.0) << 15)
                    | (f.f3 << 12)
                    | (((o >> 1) & 0xf) << 8)
                    | (((o >> 11) & 1) << 7)
                    | 0x63
            };
        }
        let mut text = String::new();
        for l in &self.lines {
            let indent = if l.ends_with(':') { "" } else { "    " };
            let _ = writeln!(text, "{indent}{l}");
        }
        (self.words, text)
    }
}

/// Where a kernel's parameters and results live in the core's memory, from
/// `base`: the pair count, the gate's eight entries, the done word, four
/// columns of `capacity` operands, and the results.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Layout {
    pub base: u32,
    pub capacity: u32,
}

impl Layout {
    pub const COUNT: u32 = 0;
    pub const ENTRIES: u32 = 4;
    pub const DONE: u32 = 36;
    const COLUMNS: u32 = 0x40;

    /// The address of column `k` (0 to 3).
    pub fn column(&self, k: u32) -> u32 {
        self.base + Self::COLUMNS + k * self.capacity * 4
    }

    /// The address of the results: for a gate, eight words per pair (each
    /// row's low and high words, rows 0r, 0i, 1r, 1i); for the probe, four
    /// (the product's low and high words, then `fx`'s).
    pub fn results(&self) -> u32 {
        self.column(4)
    }

    /// The bytes from `base` the layout spans.
    pub fn span(&self) -> u32 {
        Self::COLUMNS + 4 * self.capacity * 4 + 8 * self.capacity * 4
    }
}

/// `fx(m, a)` into (`T4` low, `T5` high): `A5` holds the rounding constant.
fn fx(p: &mut Asm, m: R, a: R, sh: u32) {
    p.mul(T4, m, a).mulh(T5, m, a);
    p.add(T4, T4, A5).sltu(T6, T4, A5).add(T5, T5, T6);
    p.srli(T4, T4, sh).slli(T6, T5, 32 - sh).or(T4, T4, T6).srai(T5, T5, sh);
}

/// The accumulator (`A6` low, `A7` high) plus or minus (`T4`, `T5`).
fn acc(p: &mut Asm, add: bool) {
    if add {
        p.add(A6, A6, T4).sltu(T6, A6, T4).add(A7, A7, T5).add(A7, A7, T6);
    } else {
        p.sltu(T6, A6, T4).sub(A6, A6, T4).sub(A7, A7, T5).sub(A7, A7, T6);
    }
}

/// The program that applies one gate to every pair: `rnd` the rounding
/// constant, `sh` the shift, `last_add` whether row 0's last real term is
/// added (a mutant).
pub fn rows_program(l: Layout, rnd: u32, sh: u32, last_add: bool) -> (Vec<u32>, String) {
    // s0 base, s1 pairs left, s2..s9 the entries, a0..a3 the columns, a4 the
    // results, a5 the rounding constant, t0..t3 one pair's operands.
    const M: [R; 8] = [R(18), R(19), R(20), R(21), R(22), R(23), R(24), R(25)];
    let mut p = Asm::default();
    p.li(S0, l.base).lw(S1, Layout::COUNT as i32, S0);
    for (k, &m) in M.iter().enumerate() {
        p.lw(m, (Layout::ENTRIES + 4 * k as u32) as i32, S0);
    }
    p.li(A0, l.column(0)).li(A1, l.column(1)).li(A2, l.column(2)).li(A3, l.column(3)).li(A4, l.results()).li(A5, rnd);
    p.label("next").beq(S1, ZERO, "done");
    p.lw(T0, 0, A0).lw(T1, 0, A1).lw(T2, 0, A2).lw(T3, 0, A3);
    // Rows 0r, 0i, 1r, 1i: their terms (entry, operand, added) in order.
    let rows: [[(usize, R, bool); 4]; 4] = [
        [(0, T0, true), (1, T1, false), (2, T2, true), (3, T3, last_add)],
        [(0, T1, true), (1, T0, true), (2, T3, true), (3, T2, true)],
        [(4, T0, true), (5, T1, false), (6, T2, true), (7, T3, false)],
        [(4, T1, true), (5, T0, true), (6, T3, true), (7, T2, true)],
    ];
    for (r, terms) in rows.iter().enumerate() {
        for (i, &(e, a, add)) in terms.iter().enumerate() {
            fx(&mut p, M[e], a, sh);
            if i == 0 {
                p.addi(A6, T4, 0).addi(A7, T5, 0);
            } else {
                acc(&mut p, add);
            }
        }
        p.sw(A6, 8 * r as i32, A4).sw(A7, 8 * r as i32 + 4, A4);
    }
    p.addi(A0, A0, 4).addi(A1, A1, 4).addi(A2, A2, 4).addi(A3, A3, 4).addi(A4, A4, 32).addi(S1, S1, -1).j("next");
    p.label("done").addi(T6, ZERO, 1).sw(T6, Layout::DONE as i32, S0);
    p.label("idle").j("idle");
    p.finish()
}

/// The program that forms each pair's `a·b` and `fx(a, b)`: `a` in column 0,
/// `b` in column 1.
pub fn probe_program(l: Layout, rnd: u32, sh: u32) -> (Vec<u32>, String) {
    let mut p = Asm::default();
    p.li(S0, l.base).lw(S1, Layout::COUNT as i32, S0);
    p.li(A0, l.column(0)).li(A1, l.column(1)).li(A4, l.results()).li(A5, rnd);
    p.label("next").beq(S1, ZERO, "done");
    p.lw(T0, 0, A0).lw(T1, 0, A1);
    p.mul(T2, T0, T1).mulh(T3, T0, T1).sw(T2, 0, A4).sw(T3, 4, A4);
    fx(&mut p, T1, T0, sh);
    p.sw(T4, 8, A4).sw(T5, 12, A4);
    p.addi(A0, A0, 4).addi(A1, A1, 4).addi(A4, A4, 16).addi(S1, S1, -1).j("next");
    p.label("done").addi(T6, ZERO, 1).sw(T6, Layout::DONE as i32, S0);
    p.label("idle").j("idle");
    p.finish()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn one(f: impl FnOnce(&mut Asm)) -> u32 {
        let mut p = Asm::default();
        f(&mut p);
        p.finish().0[0]
    }

    /// Encodings from the instruction set manual's formats, worked by hand.
    #[test]
    fn encodings_match_the_manual() {
        assert_eq!(one(|p| { p.addi(R(1), ZERO, 1); }), 0x0010_0093);
        assert_eq!(one(|p| { p.add(R(3), R(1), R(2)); }), 0x0020_81b3);
        assert_eq!(one(|p| { p.sub(R(3), R(1), R(2)); }), 0x4020_81b3);
        assert_eq!(one(|p| { p.mul(R(3), R(1), R(2)); }), 0x0220_81b3);
        assert_eq!(one(|p| { p.mulh(R(3), R(1), R(2)); }), 0x0220_91b3);
        assert_eq!(one(|p| { p.sltu(T1, T0, A5); }), 0x00f2_b333);
        assert_eq!(one(|p| { p.or(R(3), R(1), R(2)); }), 0x0020_e1b3);
        assert_eq!(one(|p| { p.lw(T0, 8, A0); }), 0x0085_2283);
        assert_eq!(one(|p| { p.sw(T0, 8, A0); }), 0x0055_2423);
        assert_eq!(one(|p| { p.lui(T0, 0x12345); }), 0x1234_52b7);
        assert_eq!(one(|p| { p.slli(T0, T0, 2); }), 0x0022_9293);
        assert_eq!(one(|p| { p.srli(T0, T0, 30); }), 0x01e2_d293);
        assert_eq!(one(|p| { p.srai(T0, T0, 30); }), 0x41e2_d293);
        assert_eq!(one(|p| { p.addi(T0, T0, -1); }), 0xfff2_8293);
        assert_eq!(one(|p| { p.label("x").j("x"); }), 0x0000_006f);
        assert_eq!(one(|p| { p.label("x").beq(ZERO, ZERO, "x"); }), 0x0000_0063);
    }

    #[test]
    fn branches_reach_forward_and_back() {
        let mut p = Asm::default();
        p.label("top").beq(A0, A1, "end").addi(A0, A0, 1).j("top").label("end");
        let (w, text) = p.finish();
        // beq a0, a1, +12: imm[4:1] = 6.
        assert_eq!(w[0], (6 << 8) | (11 << 20) | (10 << 15) | 0x63);
        // j -8: imm = -8, 21 bits.
        assert_eq!(w[2], 0xff9f_f06f);
        assert!(text.contains("beq a0, a1, end") && text.contains("j top"));
    }

    #[test]
    fn li_builds_every_kind_of_constant() {
        for v in [0u32, 1, 0x7ff, 0x800, 0xfff, 0x1000, 0x2000_0000, 0x1234_5678, 0x8000_0800, 0xffff_ffff, 0x1_0040] {
            let mut p = Asm::default();
            p.li(T0, v);
            let (w, _) = p.finish();
            // Evaluate: lui sets the upper bits, addi adds the sign-extended low.
            let mut x = 0u32;
            for word in w {
                match word & 0x7f {
                    0x37 => x = word & 0xffff_f000,
                    0x13 => x = (if (word >> 15) & 0x1f == 0 { 0 } else { x }).wrapping_add(((word as i32) >> 20) as u32),
                    _ => unreachable!(),
                }
            }
            assert_eq!(x, v, "{v:#x}");
        }
    }

    #[test]
    fn the_programs_and_their_mutants_differ_and_fit() {
        let l = Layout { base: 0x1_0000, capacity: 1024 };
        let (w, text) = rows_program(l, 1 << 29, 30, false);
        assert!(text.contains("mulh") && text.contains("srai") && w.len() > 200);
        assert_ne!(rows_program(l, 1 << 28, 30, false).1, text);
        assert_ne!(rows_program(l, 1 << 29, 29, false).1, text);
        assert_ne!(rows_program(l, 1 << 29, 30, true).1, text);
        assert_ne!(probe_program(l, 1 << 29, 30).1, probe_program(l, 1 << 28, 30).1);
        assert_eq!(l.results(), l.base + 0x40 + 16 * 1024);
        assert_eq!(l.span(), 0x40 + 48 * 1024);
    }
}