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

handwired_rs/dsl/
control.rs
