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::default_value())
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    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6),
280    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6, A7 B7 7),
281    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6, A7 B7 7, A8 B8 8),
282    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6, A7 B7 7, A8 B8 8, A9 B9 9),
283    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6, A7 B7 7, A8 B8 8, A9 B9 9, A10 B10 10),
284    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6, A7 B7 7, A8 B8 8, A9 B9 9, A10 B10 10, A11 B11 11)
285);
286
287// ---- Branching --------------------------------------------------------------------------------------------------------
288
289/// How an `if` runs: a native condition takes one branch, like Rust; a
290/// circuit condition runs both and selects.
291pub enum Branch {
292    Static(bool),
293    Dynamic(Val<bool>),
294}
295
296pub trait Cond: Copy {
297    fn branch(self) -> Branch;
298}
299impl Cond for bool {
300    fn branch(self) -> Branch {
301        Branch::Static(self)
302    }
303}
304impl Cond for Val<bool> {
305    fn branch(self) -> Branch {
306        Branch::Dynamic(self)
307    }
308}
309
310/// An `if` whose branches only compute a value.
311pub trait IfElse<A, B> {
312    type Out;
313    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> Self::Out;
314}
315
316impl<A: Unify<B>, B> IfElse<A, B> for bool {
317    type Out = A::Out;
318    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> A::Out {
319        if self { then().left() } else { A::right(otherwise()) }
320    }
321}
322
323impl<A: Unify<B>, B> IfElse<A, B> for Val<bool>
324where
325    A::Out: Lift,
326    <A::Out as Lift>::Out: Mux,
327{
328    type Out = <A::Out as Lift>::Out;
329    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> Self::Out {
330        let on = then().left().lift();
331        let off = A::right(otherwise()).lift();
332        Mux::select(self, off, on)
333    }
334}
335
336pub fn if_else<C: IfElse<A, B>, A, B>(condition: C, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> C::Out {
337    condition.if_else(then, otherwise)
338}
339
340// ---- Indexing ---------------------------------------------------------------------------------------------------------
341
342/// `items[index]`: native for a native index, a table read for a circuit index.
343pub trait Elem<I>: Sized {
344    type Out;
345    fn index(items: &[Self], index: I) -> Self::Out;
346}
347
348impl<T: Scalar> Elem<usize> for Val<T> {
349    type Out = Val<T>;
350    fn index(items: &[Val<T>], index: usize) -> Val<T> {
351        items[index]
352    }
353}
354impl<T: Scalar, I: IntScalar> Elem<Val<I>> for Val<T> {
355    type Out = Val<T>;
356    fn index(items: &[Val<T>], index: Val<I>) -> Val<T> {
357        Mux::read_dynamic(items, index)
358    }
359}
360impl<E: Copy, const N: usize> Elem<usize> for [E; N] {
361    type Out = [E; N];
362    fn index(items: &[[E; N]], index: usize) -> [E; N] {
363        items[index]
364    }
365}
366impl<E: Copy, const N: usize, I: IntScalar> Elem<Val<I>> for [E; N]
367where
368    [E; N]: Lift,
369    <[E; N] as Lift>::Out: Mux,
370{
371    type Out = <[E; N] as Lift>::Out;
372    fn index(items: &[[E; N]], index: Val<I>) -> Self::Out {
373        let lifted: Vec<Self::Out> = items.iter().map(|&item| item.lift()).collect();
374        Mux::read_dynamic(&lifted, index)
375    }
376}
377
378macro_rules! native_elements {
379    ($($t:ident),*) => {$(
380        impl Elem<usize> for $t {
381            type Out = $t;
382            fn index(items: &[$t], index: usize) -> $t {
383                items[index]
384            }
385        }
386        impl<I: IntScalar> Elem<Val<I>> for $t {
387            type Out = Val<$t>;
388            fn index(items: &[$t], index: Val<I>) -> Val<$t> {
389                let constants: Vec<Val<$t>> = items.iter().map(|&item| Val::constant(item)).collect();
390                Mux::read_dynamic(&constants, index)
391            }
392        }
393    )*};
394}
395native_elements!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
396
397pub fn index<E: Elem<I>, I, const N: usize>(items: &[E; N], index: I) -> E::Out {
398    E::index(items, index)
399}
400
401/// `items[index] = value`.
402pub trait StoreElem<I>: Sized {
403    fn store(items: &mut [Self], index: I, value: Self);
404}
405impl<E> StoreElem<usize> for E {
406    fn store(items: &mut [E], index: usize, value: E) {
407        items[index] = value;
408    }
409}
410impl<E: Mux, I: IntScalar> StoreElem<Val<I>> for E {
411    fn store(items: &mut [E], index: Val<I>, value: E) {
412        for (position, item) in items.iter_mut().enumerate() {
413            *item = E::select(eq_position(index, position), *item, value);
414        }
415    }
416}
417
418pub fn store<E: StoreElem<I>, I, V: Conv<E>, const N: usize>(items: &mut [E; N], index: I, value: V) {
419    E::store(items, index, value.conv());
420}
421
422// ---- Casts ------------------------------------------------------------------------------------------------------------
423
424/// `value as T`.
425pub trait CastTo<T> {
426    type Out;
427    fn cast_to(self) -> Self::Out;
428}
429
430#[track_caller]
431pub fn cast<T, S: CastTo<T>>(value: S) -> S::Out {
432    value.cast_to()
433}
434
435impl<S: Scalar, T: Scalar> CastTo<T> for Val<S> {
436    type Out = Val<T>;
437    #[track_caller]
438    fn cast_to(self) -> Val<T> {
439        cast_value::<S, T>(self, Location::caller())
440    }
441}
442
443fn cast_value<S: Scalar, T: Scalar>(value: Val<S>, at: &'static Location<'static>) -> Val<T> {
444    let same = Val::<T>::from_wire(value.wire);
445    match (S::KIND, T::KIND) {
446        (Kind::Bool, _) => same,
447        (_, Kind::Bool) => unsupported("only bool casts to bool"),
448        (Kind::Float, Kind::Float) => same,
449        (Kind::Int, Kind::Float) => {
450            let range = with_circuit(|circuit| circuit.range(value.wire));
451            if range.min >= -16_777_216.0 && range.max <= 16_777_216.0 {
452                same
453            } else {
454                panic!("{at}: casting integers beyond ±2^24 to f32 rounds, which isn't supported yet");
455            }
456        }
457        (Kind::Int, Kind::Int) => {
458            if S::MIN >= T::MIN && S::MAX <= T::MAX {
459                return same;
460            }
461            let range = with_circuit(|circuit| circuit.range(value.wire));
462            if range.min >= T::MIN && range.max <= T::MAX {
463                return same;
464            }
465            if is_wide::<T>() {
466                panic!("{at}: this cast to {} may wrap, which is supported for types up to 16 bits", T::NAME);
467            }
468            // `as` between integers keeps the low bits: the same value modulo the target's span.
469            let span = T::MAX - T::MIN + 1.0;
470            integer_table(Val::<T>::from_wire(value.wire), at, |wide| {
471                let wrapped = (wide as f64 - T::MIN).rem_euclid(span) + T::MIN;
472                wrapped as i64
473            })
474        }
475        (Kind::Float, Kind::Int) => truncate::<T>(value.wire, at),
476    }
477}
478
479/// `x as T` for an f32 x: truncate toward zero, then saturate to T's range.
480fn truncate<T: Scalar>(value: Wire, at: &'static Location<'static>) -> Val<T> {
481    let range = with_circuit(|circuit| circuit.range(value));
482    if !(range.min > -4_194_304.0 && range.max < 4_194_304.0) {
483        panic!("{at}: casting f32 to an integer needs a declared range within ±2^22, got {}..{}", range.min, range.max);
484    }
485    let truncated = with_circuit(|circuit| {
486        let rounded = emit_round(circuit, value);
487        // Round to nearest, then step back toward zero where rounding went away from it.
488        let rounded_up = emit_compare(circuit, rounded, Relation::Greater, value);
489        let rounded_down = emit_compare(circuit, rounded, Relation::Less, value);
490        let zero = circuit.constant(0.0);
491        let positive = emit_compare(circuit, value, Relation::AtLeast, zero);
492        let too_high = circuit.multiply(rounded_up, positive);
493        let negative = circuit.affine(1.0 - positive);
494        let too_low = circuit.multiply(rounded_down, negative);
495        let wire = circuit.affine(Lin::new(vec![(rounded, 1.0), (too_high, -1.0), (too_low, 1.0)], 0.0));
496        circuit.assume_grid_allowing_readers(wire, 1.0);
497        Val::<T>::from_wire(wire)
498    });
499    if range.min >= T::MIN && range.max <= T::MAX {
500        return truncated;
501    }
502    let low = Val::<T>::constant(T::from_f64(T::MIN));
503    let high = Val::<T>::constant(T::from_f64(T::MAX));
504    let above = Mux::select(gt(truncated, high), truncated, high);
505    Mux::select(gt(low, above), above, low)
506}
507
508macro_rules! native_casts {
509    ($($source:ident),*) => {$(
510        native_casts!(@to $source: u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
511    )*};
512    (@to $source:ident: $($target:ident),*) => {$(
513        impl CastTo<$target> for $source {
514            type Out = $target;
515            fn cast_to(self) -> $target {
516                self as $target
517            }
518        }
519    )*};
520}
521native_casts!(u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
522
523macro_rules! bool_casts {
524    ($($target:ident),*) => {$(
525        impl CastTo<$target> for bool {
526            type Out = $target;
527            fn cast_to(self) -> $target {
528                self as $target
529            }
530        }
531    )*};
532}
533bool_casts!(u8, u16, u32, u64, usize, i8, i16, i32, i64);
534
535#[cfg(test)]
536mod tests {
537    use super::*;
538    use crate::circuit::Circuit;
539
540    #[test]
541    fn conversions_infer_from_context() {
542        super::super::CONTEXT.with(|context| *context.borrow_mut() = Some(Circuit::new(0)));
543        let value: Val<u8> = Conv::conv(5u8);
544        let picked = if_else(true, || value, || 7);
545        let _: Val<u8> = picked;
546        let native = if_else(false, || 1, || 2);
547        assert_eq!(native, 2);
548        super::super::CONTEXT.with(|context| *context.borrow_mut() = None);
549    }
550}
handwired_rs/dsl/control.rs

handwired_rs/dsl/
control.rs
