handwired_rs/dsl/
control.rs

1//! Conversions, selection, branching, indexing and casts.
2
3use std::panic::Location;
4
5use super::ops::integer_table;
6use super::{IntScalar, Kind, Scalar, Val, eq, gt, is_wide, known_bit, settle, with_circuit};
7use crate::circuit::OutOfRange;
8use crate::emit::{Relation, emit_compare, emit_round};
9use crate::wire::{Lin, Wire};
10
11/// The build stops here: the macro puts this where a construct can't be
12/// translated, e.g. a `return` inside a branch whose condition depends on
13/// the inputs.
14#[track_caller]
15pub fn unsupported(what: &str) -> ! {
16    panic!("{}: {what}", Location::caller())
17}
18
19// ---- Selection --------------------------------------------------------------------------------------------------------
20
21/// A value a circuit can choose between: `bit ? on : off`, with both sides computed.
22pub trait Mux: Copy {
23    fn select(bit: Val<bool>, off: Self, on: Self) -> Self;
24
25    /// `items[index]` at an index that depends on the inputs. Out of bounds
26    /// reads the first item; the oracle would have panicked there.
27    fn read_dynamic<I: IntScalar>(items: &[Self], index: Val<I>) -> Self {
28        let mut chosen = items[0];
29        for (position, &item) in items.iter().enumerate().skip(1) {
30            chosen = Self::select(eq_position(index, position), chosen, item);
31        }
32        chosen
33    }
34}
35
36fn eq_position<I: IntScalar>(index: Val<I>, position: usize) -> Val<bool> {
37    eq(index, I::from_f64(position as f64))
38}
39
40impl<T: Scalar> Mux for Val<T> {
41    fn select(bit: Val<bool>, off: Val<T>, on: Val<T>) -> Val<T> {
42        if off.wire == on.wire {
43            return off;
44        }
45        if let Some(value) = known_bit(bit) {
46            return if value { on } else { off };
47        }
48        with_circuit(|circuit| {
49            let (off_range, on_range) = (circuit.range(off.wire), circuit.range(on.wire));
50            let wire = circuit.select(bit.wire, off.wire, on.wire);
51            settle(circuit, wire, off_range.min.min(on_range.min), off_range.max.max(on_range.max))
52        })
53    }
54
55    fn read_dynamic<I: IntScalar>(items: &[Val<T>], index: Val<I>) -> Val<T> {
56        let candidates: Vec<Wire> = items.iter().map(|item| item.wire).collect();
57        with_circuit(|circuit| {
58            let wire = circuit.index_read(index.wire, &candidates, OutOfRange::Clamp);
59            let range = circuit.range(wire);
60            settle(circuit, wire, range.min, range.max)
61        })
62    }
63}
64
65impl Mux for () {
66    fn select(_: Val<bool>, _: (), _: ()) {}
67}
68
69impl<M: Mux, const N: usize> Mux for [M; N] {
70    fn select(bit: Val<bool>, off: [M; N], on: [M; N]) -> [M; N] {
71        std::array::from_fn(|position| M::select(bit, off[position], on[position]))
72    }
73}
74
75/// A stand-in value for a branch that ends in `return`, `break` or
76/// `continue`: the branch's value is never used on the paths that take it,
77/// but both branches of an input-dependent `if` are computed.
78pub trait Placeholder: Copy {
79    fn placeholder() -> Self;
80}
81impl<T: Scalar> Placeholder for Val<T> {
82    fn placeholder() -> Val<T> {
83        Val::constant(T::from_f64(0.0))
84    }
85}
86impl Placeholder for () {
87    fn placeholder() {}
88}
89impl<P: Placeholder, const N: usize> Placeholder for [P; N] {
90    fn placeholder() -> [P; N] {
91        [P::placeholder(); N]
92    }
93}
94
95pub fn placeholder<P: Placeholder>() -> P {
96    P::placeholder()
97}
98
99// ---- Conversions ------------------------------------------------------------------------------------------------------
100
101/// Turn a value into the circuit form of a known type, e.g. a literal `5`
102/// assigned to a `u8` variable that holds a circuit value.
103pub trait Conv<T> {
104    fn conv(self) -> T;
105}
106
107impl<T: Scalar> Conv<Val<T>> for Val<T> {
108    fn conv(self) -> Val<T> {
109        self
110    }
111}
112impl Conv<()> for () {
113    fn conv(self) {}
114}
115impl<A: Conv<B>, B, const N: usize> Conv<[B; N]> for [A; N] {
116    fn conv(self) -> [B; N] {
117        self.map(Conv::conv)
118    }
119}
120
121/// The circuit form of a value whose type is not known yet, for `let mut`
122/// bindings that later take circuit values.
123pub trait Lift {
124    type Out;
125    fn lift(self) -> Self::Out;
126}
127
128impl<T: Scalar> Lift for Val<T> {
129    type Out = Val<T>;
130    fn lift(self) -> Val<T> {
131        self
132    }
133}
134impl Lift for () {
135    type Out = ();
136    fn lift(self) {}
137}
138impl<L: Lift, const N: usize> Lift for [L; N] {
139    type Out = [L::Out; N];
140    fn lift(self) -> [L::Out; N] {
141        self.map(Lift::lift)
142    }
143}
144
145/// The common type of an `if`'s two branch values: native when both are
146/// native, a circuit value when either is.
147pub trait Unify<B>: Sized {
148    type Out;
149    fn left(self) -> Self::Out;
150    fn right(other: B) -> Self::Out;
151}
152
153impl<T: Scalar> Unify<Val<T>> for Val<T> {
154    type Out = Val<T>;
155    fn left(self) -> Val<T> {
156        self
157    }
158    fn right(other: Val<T>) -> Val<T> {
159        other
160    }
161}
162impl<T: Scalar> Unify<T> for Val<T> {
163    type Out = Val<T>;
164    fn left(self) -> Val<T> {
165        self
166    }
167    fn right(other: T) -> Val<T> {
168        Val::constant(other)
169    }
170}
171impl Unify<()> for () {
172    type Out = ();
173    fn left(self) {}
174    fn right(_: ()) {}
175}
176impl<A: Unify<B>, B, const N: usize> Unify<[B; N]> for [A; N] {
177    type Out = [A::Out; N];
178    fn left(self) -> [A::Out; N] {
179        self.map(A::left)
180    }
181    fn right(other: [B; N]) -> [A::Out; N] {
182        other.map(A::right)
183    }
184}
185
186macro_rules! native_conversions {
187    ($($t:ident),*) => {$(
188        impl Conv<Val<$t>> for $t {
189            fn conv(self) -> Val<$t> {
190                Val::constant(self)
191            }
192        }
193        impl Lift for $t {
194            type Out = Val<$t>;
195            fn lift(self) -> Val<$t> {
196                Val::constant(self)
197            }
198        }
199        impl Unify<$t> for $t {
200            type Out = $t;
201            fn left(self) -> $t {
202                self
203            }
204            fn right(other: $t) -> $t {
205                other
206            }
207        }
208        impl Unify<Val<$t>> for $t {
209            type Out = Val<$t>;
210            fn left(self) -> Val<$t> {
211                Val::constant(self)
212            }
213            fn right(other: Val<$t>) -> Val<$t> {
214                other
215            }
216        }
217    )*};
218}
219native_conversions!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
220
221macro_rules! tuple_traits {
222    ($(($($name:ident $other:ident $index:tt),+)),*) => {$(
223        impl<$($name: Placeholder),+> Placeholder for ($($name,)+) {
224            fn placeholder() -> Self {
225                ($($name::placeholder(),)+)
226            }
227        }
228        impl<$($name: Mux),+> Mux for ($($name,)+) {
229            fn select(bit: Val<bool>, off: Self, on: Self) -> Self {
230                ($($name::select(bit, off.$index, on.$index),)+)
231            }
232        }
233        impl<$($name: Conv<$other>, $other),+> Conv<($($other,)+)> for ($($name,)+) {
234            fn conv(self) -> ($($other,)+) {
235                ($(self.$index.conv(),)+)
236            }
237        }
238        impl<$($name: Lift),+> Lift for ($($name,)+) {
239            type Out = ($($name::Out,)+);
240            fn lift(self) -> Self::Out {
241                ($(self.$index.lift(),)+)
242            }
243        }
244        impl<$($name: Unify<$other>, $other),+> Unify<($($other,)+)> for ($($name,)+) {
245            type Out = ($($name::Out,)+);
246            fn left(self) -> Self::Out {
247                ($(self.$index.left(),)+)
248            }
249            fn right(other: ($($other,)+)) -> Self::Out {
250                ($($name::right(other.$index),)+)
251            }
252        }
253        impl<$($name: Copy),+> Elem<usize> for ($($name,)+) {
254            type Out = Self;
255            fn index(items: &[Self], index: usize) -> Self {
256                items[index]
257            }
258        }
259        impl<$($name: Copy),+, I: IntScalar> Elem<Val<I>> for ($($name,)+)
260        where
261            Self: Lift,
262            <Self as Lift>::Out: Mux,
263        {
264            type Out = <Self as Lift>::Out;
265            fn index(items: &[Self], index: Val<I>) -> Self::Out {
266                let lifted: Vec<Self::Out> = items.iter().map(|&item| item.lift()).collect();
267                Mux::read_dynamic(&lifted, index)
268            }
269        }
270    )*};
271}
272tuple_traits!(
273    (A0 B0 0),
274    (A0 B0 0, A1 B1 1),
275    (A0 B0 0, A1 B1 1, A2 B2 2),
276    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3),
277    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4),
278    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5)
279);
280
281// ---- Branching --------------------------------------------------------------------------------------------------------
282
283/// How an `if` runs: a native condition takes one branch, like Rust; a
284/// circuit condition runs both and selects.
285pub enum Branch {
286    Static(bool),
287    Dynamic(Val<bool>),
288}
289
290pub trait Cond: Copy {
291    fn branch(self) -> Branch;
292}
293impl Cond for bool {
294    fn branch(self) -> Branch {
295        Branch::Static(self)
296    }
297}
298impl Cond for Val<bool> {
299    fn branch(self) -> Branch {
300        Branch::Dynamic(self)
301    }
302}
303
304/// An `if` whose branches only compute a value.
305pub trait IfElse<A, B> {
306    type Out;
307    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> Self::Out;
308}
309
310impl<A: Unify<B>, B> IfElse<A, B> for bool {
311    type Out = A::Out;
312    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> A::Out {
313        if self { then().left() } else { A::right(otherwise()) }
314    }
315}
316
317impl<A: Unify<B>, B> IfElse<A, B> for Val<bool>
318where
319    A::Out: Lift,
320    <A::Out as Lift>::Out: Mux,
321{
322    type Out = <A::Out as Lift>::Out;
323    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> Self::Out {
324        let on = then().left().lift();
325        let off = A::right(otherwise()).lift();
326        Mux::select(self, off, on)
327    }
328}
329
330pub fn if_else<C: IfElse<A, B>, A, B>(condition: C, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> C::Out {
331    condition.if_else(then, otherwise)
332}
333
334// ---- Indexing ---------------------------------------------------------------------------------------------------------
335
336/// `items[index]`: native for a native index, a table read for a circuit index.
337pub trait Elem<I>: Sized {
338    type Out;
339    fn index(items: &[Self], index: I) -> Self::Out;
340}
341
342impl<T: Scalar> Elem<usize> for Val<T> {
343    type Out = Val<T>;
344    fn index(items: &[Val<T>], index: usize) -> Val<T> {
345        items[index]
346    }
347}
348impl<T: Scalar, I: IntScalar> Elem<Val<I>> for Val<T> {
349    type Out = Val<T>;
350    fn index(items: &[Val<T>], index: Val<I>) -> Val<T> {
351        Mux::read_dynamic(items, index)
352    }
353}
354impl<E: Copy, const N: usize> Elem<usize> for [E; N] {
355    type Out = [E; N];
356    fn index(items: &[[E; N]], index: usize) -> [E; N] {
357        items[index]
358    }
359}
360impl<E: Copy, const N: usize, I: IntScalar> Elem<Val<I>> for [E; N]
361where
362    [E; N]: Lift,
363    <[E; N] as Lift>::Out: Mux,
364{
365    type Out = <[E; N] as Lift>::Out;
366    fn index(items: &[[E; N]], index: Val<I>) -> Self::Out {
367        let lifted: Vec<Self::Out> = items.iter().map(|&item| item.lift()).collect();
368        Mux::read_dynamic(&lifted, index)
369    }
370}
371
372macro_rules! native_elements {
373    ($($t:ident),*) => {$(
374        impl Elem<usize> for $t {
375            type Out = $t;
376            fn index(items: &[$t], index: usize) -> $t {
377                items[index]
378            }
379        }
380        impl<I: IntScalar> Elem<Val<I>> for $t {
381            type Out = Val<$t>;
382            fn index(items: &[$t], index: Val<I>) -> Val<$t> {
383                let constants: Vec<Val<$t>> = items.iter().map(|&item| Val::constant(item)).collect();
384                Mux::read_dynamic(&constants, index)
385            }
386        }
387    )*};
388}
389native_elements!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
390
391pub fn index<E: Elem<I>, I, const N: usize>(items: &[E; N], index: I) -> E::Out {
392    E::index(items, index)
393}
394
395/// `items[index] = value`.
396pub trait StoreElem<I>: Sized {
397    fn store(items: &mut [Self], index: I, value: Self);
398}
399impl<E> StoreElem<usize> for E {
400    fn store(items: &mut [E], index: usize, value: E) {
401        items[index] = value;
402    }
403}
404impl<E: Mux, I: IntScalar> StoreElem<Val<I>> for E {
405    fn store(items: &mut [E], index: Val<I>, value: E) {
406        for (position, item) in items.iter_mut().enumerate() {
407            *item = E::select(eq_position(index, position), *item, value);
408        }
409    }
410}
411
412pub fn store<E: StoreElem<I>, I, V: Conv<E>, const N: usize>(items: &mut [E; N], index: I, value: V) {
413    E::store(items, index, value.conv());
414}
415
416// ---- Casts ------------------------------------------------------------------------------------------------------------
417
418/// `value as T`.
419pub trait CastTo<T> {
420    type Out;
421    fn cast_to(self) -> Self::Out;
422}
423
424#[track_caller]
425pub fn cast<T, S: CastTo<T>>(value: S) -> S::Out {
426    value.cast_to()
427}
428
429impl<S: Scalar, T: Scalar> CastTo<T> for Val<S> {
430    type Out = Val<T>;
431    #[track_caller]
432    fn cast_to(self) -> Val<T> {
433        cast_value::<S, T>(self, Location::caller())
434    }
435}
436
437fn cast_value<S: Scalar, T: Scalar>(value: Val<S>, at: &'static Location<'static>) -> Val<T> {
438    let same = Val::<T>::from_wire(value.wire);
439    match (S::KIND, T::KIND) {
440        (Kind::Bool, _) => same,
441        (_, Kind::Bool) => unsupported("only bool casts to bool"),
442        (Kind::Float, Kind::Float) => same,
443        (Kind::Int, Kind::Float) => {
444            let range = with_circuit(|circuit| circuit.range(value.wire));
445            if range.min >= -16_777_216.0 && range.max <= 16_777_216.0 {
446                same
447            } else {
448                panic!("{at}: casting integers beyond ±2^24 to f32 rounds, which isn't supported yet");
449            }
450        }
451        (Kind::Int, Kind::Int) => {
452            if S::MIN >= T::MIN && S::MAX <= T::MAX {
453                return same;
454            }
455            let range = with_circuit(|circuit| circuit.range(value.wire));
456            if range.min >= T::MIN && range.max <= T::MAX {
457                return same;
458            }
459            if is_wide::<T>() {
460                panic!("{at}: this cast to {} may wrap, which is supported for types up to 16 bits", T::NAME);
461            }
462            // `as` between integers keeps the low bits: the same value modulo the target's span.
463            let span = T::MAX - T::MIN + 1.0;
464            integer_table(Val::<T>::from_wire(value.wire), at, |wide| {
465                let wrapped = (wide as f64 - T::MIN).rem_euclid(span) + T::MIN;
466                wrapped as i64
467            })
468        }
469        (Kind::Float, Kind::Int) => truncate::<T>(value.wire, at),
470    }
471}
472
473/// `x as T` for an f32 x: truncate toward zero, then saturate to T's range.
474fn truncate<T: Scalar>(value: Wire, at: &'static Location<'static>) -> Val<T> {
475    let range = with_circuit(|circuit| circuit.range(value));
476    if !(range.min > -4_194_304.0 && range.max < 4_194_304.0) {
477        panic!("{at}: casting f32 to an integer needs a declared range within ±2^22, got {}..{}", range.min, range.max);
478    }
479    let truncated = with_circuit(|circuit| {
480        let rounded = emit_round(circuit, value);
481        // Round to nearest, then step back toward zero where rounding went away from it.
482        let rounded_up = emit_compare(circuit, rounded, Relation::Greater, value);
483        let rounded_down = emit_compare(circuit, rounded, Relation::Less, value);
484        let zero = circuit.constant(0.0);
485        let positive = emit_compare(circuit, value, Relation::AtLeast, zero);
486        let too_high = circuit.multiply(rounded_up, positive);
487        let negative = circuit.affine(1.0 - positive);
488        let too_low = circuit.multiply(rounded_down, negative);
489        let wire = circuit.affine(Lin::new(vec![(rounded, 1.0), (too_high, -1.0), (too_low, 1.0)], 0.0));
490        circuit.assume_grid_allowing_readers(wire, 1.0);
491        Val::<T>::from_wire(wire)
492    });
493    if range.min >= T::MIN && range.max <= T::MAX {
494        return truncated;
495    }
496    let low = Val::<T>::constant(T::from_f64(T::MIN));
497    let high = Val::<T>::constant(T::from_f64(T::MAX));
498    let above = Mux::select(gt(truncated, high), truncated, high);
499    Mux::select(gt(low, above), above, low)
500}
501
502macro_rules! native_casts {
503    ($($source:ident),*) => {$(
504        native_casts!(@to $source: u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
505    )*};
506    (@to $source:ident: $($target:ident),*) => {$(
507        impl CastTo<$target> for $source {
508            type Out = $target;
509            fn cast_to(self) -> $target {
510                self as $target
511            }
512        }
513    )*};
514}
515native_casts!(u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
516
517macro_rules! bool_casts {
518    ($($target:ident),*) => {$(
519        impl CastTo<$target> for bool {
520            type Out = $target;
521            fn cast_to(self) -> $target {
522                self as $target
523            }
524        }
525    )*};
526}
527bool_casts!(u8, u16, u32, u64, usize, i8, i16, i32, i64);
528
529#[cfg(test)]
530mod tests {
531    use super::*;
532    use crate::circuit::Circuit;
533
534    #[test]
535    fn conversions_infer_from_context() {
536        super::super::CONTEXT.with(|context| *context.borrow_mut() = Some(Circuit::new(0)));
537        let value: Val<u8> = Conv::conv(5u8);
538        let picked = if_else(true, || value, || 7);
539        let _: Val<u8> = picked;
540        let native = if_else(false, || 1, || 2);
541        assert_eq!(native, 2);
542        super::super::CONTEXT.with(|context| *context.borrow_mut() = None);
543    }
544}
handwired_rs/dsl/control.rs

handwired_rs/dsl/
control.rs
