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

handwired_rs/dsl/
control.rs
