handwired-rs 0.1.0

//! Parameters, return values, and the compiled program a `#[handwired]`
//! function produces.

use std::cell::Cell;
use std::fmt::Debug;
use std::panic::{self, AssertUnwindSafe};
use std::sync::Once;

use super::{CONTEXT, Kind, Mux, Scalar, Val, is_wide, with_circuit};
use crate::Error;
use crate::circuit::Circuit;
use crate::tensor::{TensorExecutor, TensorProgram};
use crate::wire::Wire;

/// What a parameter declares with `#[range(lo, hi)]` and `#[grid(step)]`. A
/// parameter of a compound type, e.g. `[u8; 4]`, applies them to each scalar.
#[derive(Copy, Clone, Debug, Default)]
pub struct Attrs {
    pub range: Option<(f64, f64)>,
    pub grid: Option<f64>,
}

impl Attrs {
    pub const NONE: Attrs = Attrs { range: None, grid: None };
}

/// Hands out source cells while a function's parameters are declared.
pub struct Inputs {
    next: usize,
}

/// A deterministic random source for sampling inputs.
pub struct Rng(u64);

impl Rng {
    pub fn new(seed: u64) -> Rng {
        Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1)
    }

    pub fn next_u64(&mut self) -> u64 {
        self.0 ^= self.0 >> 12;
        self.0 ^= self.0 << 25;
        self.0 ^= self.0 >> 27;
        self.0.wrapping_mul(0x2545_F491_4F6C_DD1D)
    }

    /// Uniform in 0..1.
    pub fn unit(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
    }
}

/// A type a `#[handwired]` function can take or return.
pub trait HwType: Sized {
    /// The circuit form of the type, e.g. `Val<u8>` for `u8`.
    type Val: Mux;
    const CELLS: usize;
    fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Self::Val;
    fn encode(&self, out: &mut Vec<f32>);
    fn decode(values: &mut dyn Iterator<Item = f32>) -> Self;
    fn outputs(value: &Self::Val, out: &mut Vec<Wire>);
    fn sample(rng: &mut Rng, attrs: &Attrs) -> Self;
    /// A start value for a binding declared without one.
    fn zero() -> Self::Val;
}

fn scalar_range<T: Scalar>(attrs: &Attrs) -> (f64, f64) {
    match (T::KIND, attrs.range) {
        (Kind::Bool, _) => (0.0, 1.0),
        (_, Some((low, high))) => {
            assert!(
                low >= T::MIN && high <= T::MAX && low <= high,
                "#[range({low}, {high})] doesn't fit in {}",
                T::NAME
            );
            if T::KIND == Kind::Float { (f32_at_most(low), f32_at_least(high)) } else { (low, high) }
        }
        (Kind::Int, None) if !is_wide::<T>() => (T::MIN, T::MAX),
        (Kind::Int, None) => panic!("a {} parameter needs #[range(lo, hi)] within ±2^24", T::NAME),
        (Kind::Float, None) => panic!("an f32 parameter needs #[range(lo, hi)]"),
    }
}

/// The largest f32 at or below `value`, so a declared range like 1e-3
/// covers every f32 a caller can pass for it.
fn f32_at_most(value: f64) -> f64 {
    let rounded = value as f32;
    (if rounded as f64 > value { rounded.next_down() } else { rounded }) as f64
}

fn f32_at_least(value: f64) -> f64 {
    let rounded = value as f32;
    (if (rounded as f64) < value { rounded.next_up() } else { rounded }) as f64
}

fn declare_scalar<T: Scalar>(inputs: &mut Inputs, attrs: &Attrs) -> Val<T> {
    let (low, high) = scalar_range::<T>(attrs);
    if T::KIND == Kind::Int {
        assert!(low > -16_777_216.0 && high < 16_777_216.0, "a {} parameter's range must lie within ±2^24", T::NAME);
    }
    let index = inputs.next;
    inputs.next += 1;
    with_circuit(|circuit| {
        let wire = circuit.source(index);
        match T::KIND {
            Kind::Bool => {
                circuit.assume_bit(wire);
                circuit.assume_grid(wire, 1.0);
            }
            Kind::Int => {
                circuit.assume_range(wire, low, high);
                circuit.assume_grid(wire, attrs.grid.unwrap_or(1.0).max(1.0));
            }
            Kind::Float => {
                circuit.assume_range(wire, low, high);
                if let Some(grid) = attrs.grid {
                    circuit.assume_grid(wire, grid);
                }
            }
        }
        Val::from_wire(wire)
    })
}

fn sample_scalar<T: Scalar>(rng: &mut Rng, attrs: &Attrs) -> T {
    let (low, high) = scalar_range::<T>(attrs);
    let pick = rng.next_u64() % 8;
    let value = match pick {
        0 => low,
        1 => high,
        2 if low <= 0.0 && high >= 0.0 => 0.0,
        _ => low + rng.unit() * (high - low),
    };
    let snapped = match (T::KIND, attrs.grid) {
        (Kind::Float, None) => value,
        (Kind::Float, Some(grid)) => ((value / grid).round() * grid).clamp(low, high),
        _ => value.round().clamp(low, high),
    };
    T::from_f64(snapped)
}

macro_rules! scalar_types {
    ($($t:ident),*) => {$(
        impl HwType for $t {
            type Val = Val<$t>;
            const CELLS: usize = 1;
            fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Val<$t> {
                declare_scalar::<$t>(inputs, attrs)
            }
            fn encode(&self, out: &mut Vec<f32>) {
                out.push(self.to_f64() as f32);
            }
            fn decode(values: &mut dyn Iterator<Item = f32>) -> $t {
                <$t as Scalar>::from_f64(values.next().expect("too few output values") as f64)
            }
            fn outputs(value: &Val<$t>, out: &mut Vec<Wire>) {
                out.push(value.wire);
            }
            fn sample(rng: &mut Rng, attrs: &Attrs) -> $t {
                sample_scalar::<$t>(rng, attrs)
            }
            fn zero() -> Val<$t> {
                Val::constant(<$t as Scalar>::from_f64(0.0))
            }
        }
    )*};
}
scalar_types!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);

impl<T: HwType, const N: usize> HwType for [T; N] {
    type Val = [T::Val; N];
    const CELLS: usize = N * T::CELLS;
    fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Self::Val {
        std::array::from_fn(|_| T::declare(inputs, attrs))
    }
    fn encode(&self, out: &mut Vec<f32>) {
        for item in self {
            item.encode(out);
        }
    }
    fn decode(values: &mut dyn Iterator<Item = f32>) -> Self {
        std::array::from_fn(|_| T::decode(values))
    }
    fn outputs(value: &Self::Val, out: &mut Vec<Wire>) {
        for item in value {
            T::outputs(item, out);
        }
    }
    fn sample(rng: &mut Rng, attrs: &Attrs) -> Self {
        std::array::from_fn(|_| T::sample(rng, attrs))
    }
    fn zero() -> Self::Val {
        std::array::from_fn(|_| T::zero())
    }
}

impl HwType for () {
    type Val = ();
    const CELLS: usize = 0;
    fn declare(_: &mut Inputs, _: &Attrs) {}
    fn encode(&self, _: &mut Vec<f32>) {}
    fn decode(_: &mut dyn Iterator<Item = f32>) {}
    fn outputs(_: &(), _: &mut Vec<Wire>) {}
    fn sample(_: &mut Rng, _: &Attrs) {}
    fn zero() {}
}

/// A function's parameter list, as a tuple, with one `Attrs` per parameter.
pub trait Params: Sized {
    const CELLS: usize;
    fn encode(&self, out: &mut Vec<f32>);
    fn sample(rng: &mut Rng, attrs: &[Attrs]) -> Self;
}

impl Params for () {
    const CELLS: usize = 0;
    fn encode(&self, _: &mut Vec<f32>) {}
    fn sample(_: &mut Rng, _: &[Attrs]) {}
}

macro_rules! tuple_types {
    ($(($($name:ident $index:tt),+)),*) => {$(
        impl<$($name: HwType),+> HwType for ($($name,)+) {
            type Val = ($($name::Val,)+);
            const CELLS: usize = 0 $(+ $name::CELLS)+;
            fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Self::Val {
                ($($name::declare(inputs, attrs),)+)
            }
            fn encode(&self, out: &mut Vec<f32>) {
                $(self.$index.encode(out);)+
            }
            fn decode(values: &mut dyn Iterator<Item = f32>) -> Self {
                ($($name::decode(values),)+)
            }
            fn outputs(value: &Self::Val, out: &mut Vec<Wire>) {
                $($name::outputs(&value.$index, out);)+
            }
            fn sample(rng: &mut Rng, attrs: &Attrs) -> Self {
                ($($name::sample(rng, attrs),)+)
            }
            fn zero() -> Self::Val {
                ($($name::zero(),)+)
            }
        }
    )*};
}
tuple_types!(
    (A0 0),
    (A0 0, A1 1),
    (A0 0, A1 1, A2 2),
    (A0 0, A1 1, A2 2, A3 3),
    (A0 0, A1 1, A2 2, A3 3, A4 4),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5)
);

macro_rules! tuple_params {
    ($(($($name:ident $index:tt),+)),*) => {$(
        impl<$($name: HwType),+> Params for ($($name,)+) {
            const CELLS: usize = 0 $(+ $name::CELLS)+;
            fn encode(&self, out: &mut Vec<f32>) {
                $(self.$index.encode(out);)+
            }
            fn sample(rng: &mut Rng, attrs: &[Attrs]) -> Self {
                ($($name::sample(rng, &attrs[$index]),)+)
            }
        }
    )*};
}
tuple_params!(
    (A0 0),
    (A0 0, A1 1),
    (A0 0, A1 1, A2 2),
    (A0 0, A1 1, A2 2, A3 3),
    (A0 0, A1 1, A2 2, A3 3, A4 4),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11)
);

/// A `#[handwired]` function's circuit, compiled for its outputs, with the
/// original function kept as the oracle.
pub struct Compiled<Args, R> {
    pub circuit: Circuit,
    pub program: TensorProgram,
    pub outputs: Vec<Wire>,
    oracle: fn(Args) -> R,
    attrs: &'static [Attrs],
}

/// Restores the previous build context, also when a build panics.
struct ContextGuard(Option<Option<Circuit>>);

impl Drop for ContextGuard {
    fn drop(&mut self) {
        if let Some(previous) = self.0.take() {
            CONTEXT.with(|context| *context.borrow_mut() = previous);
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Report {
    /// Inputs where the circuit matched the function.
    pub checked: usize,
    /// Inputs where the function panicked, e.g. on overflow, so any circuit value is allowed.
    pub skipped: usize,
}

thread_local! {
    static QUIET_PANICS: Cell<bool> = const { Cell::new(false) };
}

/// Silence panic messages from oracle calls on this thread only.
fn install_quiet_hook() {
    static INSTALL: Once = Once::new();
    INSTALL.call_once(|| {
        let previous = panic::take_hook();
        panic::set_hook(Box::new(move |info| {
            if !QUIET_PANICS.with(Cell::get) {
                previous(info);
            }
        }));
    });
}

impl<Args: Params + Clone + Debug, R: HwType + PartialEq + Debug> Compiled<Args, R> {
    /// Build the circuit: `body` declares the parameters and returns the outputs.
    pub fn build(oracle: fn(Args) -> R, attrs: &'static [Attrs], body: impl FnOnce(&mut Inputs) -> R::Val) -> Self {
        let previous = CONTEXT.with(|context| context.replace(Some(Circuit::new(Args::CELLS))));
        let guard = ContextGuard(Some(previous));
        let output = body(&mut Inputs { next: 0 });
        let circuit = CONTEXT.with(|context| context.borrow_mut().take()).expect("the build context disappeared");
        drop(guard);
        let mut outputs = Vec::new();
        R::outputs(&output, &mut outputs);
        let program = circuit.compile(Some(&outputs));
        Compiled { circuit, program, outputs, oracle, attrs }
    }

    fn source(&self, args: &Args) -> Vec<f32> {
        let mut source = Vec::with_capacity(Args::CELLS);
        args.encode(&mut source);
        source
    }

    /// Run the compiled program.
    pub fn eval(&self, args: &Args) -> Result<R, Error> {
        let source = self.source(args);
        self.circuit.validate_source_ranges(&source)?;
        let mut executor = TensorExecutor::new(&self.program);
        Ok(R::decode(&mut executor.run(&source).iter().copied()))
    }

    /// Run the reference evaluator, one operation at a time.
    pub fn eval_reference(&self, args: &Args) -> Result<R, Error> {
        let values = self.circuit.evaluate_reference(&self.source(args))?;
        Ok(R::decode(&mut self.outputs.iter().map(|wire| values[wire.index()])))
    }

    /// Run the original function.
    pub fn oracle(&self, args: Args) -> R {
        (self.oracle)(args)
    }

    /// Compare the compiled program with the original function on `cases`
    /// sampled inputs, including each parameter's range ends. Inputs where
    /// the function panics are skipped. Run it in a build with overflow
    /// checks, e.g. a debug build, so that overflow panics as the contract
    /// assumes instead of wrapping.
    pub fn verify(&self, cases: usize) -> Result<Report, String> {
        self.verify_seeded(cases, 0x5EED)
    }

    pub fn verify_seeded(&self, cases: usize, seed: u64) -> Result<Report, String> {
        if !cfg!(debug_assertions) {
            eprintln!(
                "warning: verify in a build without overflow checks: integer overflow wraps in the function instead \
                 of panicking, and the circuit isn't required to match there"
            );
        }
        install_quiet_hook();
        let mut rng = Rng::new(seed);
        let mut executor = TensorExecutor::new(&self.program);
        let mut report = Report { checked: 0, skipped: 0 };
        for _ in 0..cases {
            let args = Args::sample(&mut rng, self.attrs);
            QUIET_PANICS.with(|quiet| quiet.set(true));
            let expected = panic::catch_unwind(AssertUnwindSafe(|| (self.oracle)(args.clone())));
            QUIET_PANICS.with(|quiet| quiet.set(false));
            let Ok(expected) = expected else {
                report.skipped += 1;
                continue;
            };
            let source = self.source(&args);
            self.circuit.validate_source_ranges(&source).map_err(|error| error.to_string())?;
            let actual = R::decode(&mut executor.run(&source).iter().copied());
            if actual != expected {
                return Err(format!(
                    "for inputs {args:?}, the function returns {expected:?} and the circuit {actual:?}"
                ));
            }
            report.checked += 1;
        }
        Ok(report)
    }
}