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::{HwType, 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 Rust type a circuit value stands for, e.g. `u8` for `Val<u8>`: what `Some(value)` wraps when nothing else
167/// says.
168pub trait Native {
169    type Of: HwType<Val = Self>;
170}
171
172impl<T: Scalar + HwType<Val = Val<T>>> Native for Val<T> {
173    type Of = T;
174}
175impl Native for () {
176    type Of = ();
177}
178impl<V: Native, const N: usize> Native for [V; N] {
179    type Of = [V::Of; N];
180}
181
182/// The common type of an `if`'s two branch values: native when both are
183/// native, a circuit value when either is.
184pub trait Unify<B>: Sized {
185    type Out;
186    fn left(self) -> Self::Out;
187    fn right(other: B) -> Self::Out;
188}
189
190impl<T: Scalar> Unify<Val<T>> for Val<T> {
191    type Out = Val<T>;
192    fn left(self) -> Val<T> {
193        self
194    }
195    fn right(other: Val<T>) -> Val<T> {
196        other
197    }
198}
199impl<T: Scalar> Unify<T> for Val<T> {
200    type Out = Val<T>;
201    fn left(self) -> Val<T> {
202        self
203    }
204    fn right(other: T) -> Val<T> {
205        Val::constant(other)
206    }
207}
208impl Unify<()> for () {
209    type Out = ();
210    fn left(self) {}
211    fn right(_: ()) {}
212}
213impl<A: Unify<B>, B, const N: usize> Unify<[B; N]> for [A; N] {
214    type Out = [A::Out; N];
215    fn left(self) -> [A::Out; N] {
216        self.map(A::left)
217    }
218    fn right(other: [B; N]) -> [A::Out; N] {
219        other.map(A::right)
220    }
221}
222
223macro_rules! native_conversions {
224    ($($t:ident),*) => {$(
225        impl Conv<Val<$t>> for $t {
226            fn conv(self) -> Val<$t> {
227                Val::constant(self)
228            }
229        }
230        impl Lift for $t {
231            type Out = Val<$t>;
232            fn lift(self) -> Val<$t> {
233                Val::constant(self)
234            }
235        }
236        impl Unify<$t> for $t {
237            type Out = $t;
238            fn left(self) -> $t {
239                self
240            }
241            fn right(other: $t) -> $t {
242                other
243            }
244        }
245        impl Unify<Val<$t>> for $t {
246            type Out = Val<$t>;
247            fn left(self) -> Val<$t> {
248                Val::constant(self)
249            }
250            fn right(other: Val<$t>) -> Val<$t> {
251                other
252            }
253        }
254    )*};
255}
256native_conversions!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
257
258macro_rules! tuple_traits {
259    ($(($($name:ident $other:ident $index:tt),+)),*) => {$(
260        impl<$($name: Native),+> Native for ($($name,)+) {
261            type Of = ($($name::Of,)+);
262        }
263        impl<$($name: Placeholder),+> Placeholder for ($($name,)+) {
264            fn placeholder() -> Self {
265                ($($name::placeholder(),)+)
266            }
267        }
268        impl<$($name: Mux),+> Mux for ($($name,)+) {
269            fn select(bit: Val<bool>, off: Self, on: Self) -> Self {
270                ($($name::select(bit, off.$index, on.$index),)+)
271            }
272        }
273        impl<$($name: Conv<$other>, $other),+> Conv<($($other,)+)> for ($($name,)+) {
274            fn conv(self) -> ($($other,)+) {
275                ($(self.$index.conv(),)+)
276            }
277        }
278        impl<$($name: Lift),+> Lift for ($($name,)+) {
279            type Out = ($($name::Out,)+);
280            fn lift(self) -> Self::Out {
281                ($(self.$index.lift(),)+)
282            }
283        }
284        impl<$($name: Unify<$other>, $other),+> Unify<($($other,)+)> for ($($name,)+) {
285            type Out = ($($name::Out,)+);
286            fn left(self) -> Self::Out {
287                ($(self.$index.left(),)+)
288            }
289            fn right(other: ($($other,)+)) -> Self::Out {
290                ($($name::right(other.$index),)+)
291            }
292        }
293        impl<$($name: Copy),+> Elem<usize> for ($($name,)+) {
294            type Out = Self;
295            fn index(items: &[Self], index: usize) -> Self {
296                items[index]
297            }
298        }
299        impl<$($name: Copy),+, I: IntScalar> Elem<Val<I>> for ($($name,)+)
300        where
301            Self: Lift,
302            <Self as Lift>::Out: Mux,
303        {
304            type Out = <Self as Lift>::Out;
305            fn index(items: &[Self], index: Val<I>) -> Self::Out {
306                let lifted: Vec<Self::Out> = items.iter().map(|&item| item.lift()).collect();
307                Mux::read_dynamic(&lifted, index)
308            }
309        }
310    )*};
311}
312tuple_traits!(
313    (A0 B0 0),
314    (A0 B0 0, A1 B1 1),
315    (A0 B0 0, A1 B1 1, A2 B2 2),
316    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3),
317    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4),
318    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5),
319    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6),
320    (A0 B0 0, A1 B1 1, A2 B2 2, A3 B3 3, A4 B4 4, A5 B5 5, A6 B6 6, A7 B7 7),
321    (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),
322    (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),
323    (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),
324    (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)
325);
326
327// ---- Branching --------------------------------------------------------------------------------------------------------
328
329/// How an `if` runs: a native condition takes one branch, like Rust; a
330/// circuit condition runs both and selects.
331pub enum Branch {
332    Static(bool),
333    Dynamic(Val<bool>),
334}
335
336pub trait Cond: Copy {
337    fn branch(self) -> Branch;
338}
339impl Cond for bool {
340    fn branch(self) -> Branch {
341        Branch::Static(self)
342    }
343}
344impl Cond for Val<bool> {
345    fn branch(self) -> Branch {
346        Branch::Dynamic(self)
347    }
348}
349
350/// An `if` whose branches only compute a value.
351pub trait IfElse<A, B> {
352    type Out;
353    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> Self::Out;
354}
355
356impl<A: Unify<B>, B> IfElse<A, B> for bool {
357    type Out = A::Out;
358    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> A::Out {
359        if self { then().left() } else { A::right(otherwise()) }
360    }
361}
362
363impl<A: Unify<B>, B> IfElse<A, B> for Val<bool>
364where
365    A::Out: Lift,
366    <A::Out as Lift>::Out: Mux,
367{
368    type Out = <A::Out as Lift>::Out;
369    fn if_else(self, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> Self::Out {
370        let on = then().left().lift();
371        let off = A::right(otherwise()).lift();
372        Mux::select(self, off, on)
373    }
374}
375
376pub fn if_else<C: IfElse<A, B>, A, B>(condition: C, then: impl FnOnce() -> A, otherwise: impl FnOnce() -> B) -> C::Out {
377    condition.if_else(then, otherwise)
378}
379
380/// An `if` whose branches both make the type of where its value goes, e.g. a struct field: an `if` of plain literals
381/// there takes the field's type rather than `i32`.
382pub fn choose<C: Cond, R: Mux>(condition: C, then: impl FnOnce() -> R, otherwise: impl FnOnce() -> R) -> R {
383    match condition.branch() {
384        Branch::Static(flag) => {
385            if flag {
386                then()
387            } else {
388                otherwise()
389            }
390        }
391        Branch::Dynamic(bit) => {
392            let on = then();
393            let off = otherwise();
394            Mux::select(bit, off, on)
395        }
396    }
397}
398
399// ---- Indexing ---------------------------------------------------------------------------------------------------------
400
401/// `items[index]`: native for a native index, a table read for a circuit index.
402pub trait Elem<I>: Sized {
403    type Out;
404    fn index(items: &[Self], index: I) -> Self::Out;
405}
406
407impl<T: Scalar> Elem<usize> for Val<T> {
408    type Out = Val<T>;
409    fn index(items: &[Val<T>], index: usize) -> Val<T> {
410        items[index]
411    }
412}
413impl<T: Scalar, I: IntScalar> Elem<Val<I>> for Val<T> {
414    type Out = Val<T>;
415    fn index(items: &[Val<T>], index: Val<I>) -> Val<T> {
416        Mux::read_dynamic(items, index)
417    }
418}
419impl<E: Copy, const N: usize> Elem<usize> for [E; N] {
420    type Out = [E; N];
421    fn index(items: &[[E; N]], index: usize) -> [E; N] {
422        items[index]
423    }
424}
425impl<E: Copy, const N: usize, I: IntScalar> Elem<Val<I>> for [E; N]
426where
427    [E; N]: Lift,
428    <[E; N] as Lift>::Out: Mux,
429{
430    type Out = <[E; N] as Lift>::Out;
431    fn index(items: &[[E; N]], index: Val<I>) -> Self::Out {
432        let lifted: Vec<Self::Out> = items.iter().map(|&item| item.lift()).collect();
433        Mux::read_dynamic(&lifted, index)
434    }
435}
436
437macro_rules! native_elements {
438    ($($t:ident),*) => {$(
439        impl Elem<usize> for $t {
440            type Out = $t;
441            fn index(items: &[$t], index: usize) -> $t {
442                items[index]
443            }
444        }
445        impl<I: IntScalar> Elem<Val<I>> for $t {
446            type Out = Val<$t>;
447            fn index(items: &[$t], index: Val<I>) -> Val<$t> {
448                let constants: Vec<Val<$t>> = items.iter().map(|&item| Val::constant(item)).collect();
449                Mux::read_dynamic(&constants, index)
450            }
451        }
452    )*};
453}
454native_elements!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
455
456pub fn index<E: Elem<I>, I, const N: usize>(items: &[E; N], index: I) -> E::Out {
457    E::index(items, index)
458}
459
460/// `items[index] = value`.
461pub trait StoreElem<I>: Sized {
462    fn store(items: &mut [Self], index: I, value: Self);
463}
464impl<E> StoreElem<usize> for E {
465    fn store(items: &mut [E], index: usize, value: E) {
466        items[index] = value;
467    }
468}
469impl<E: Mux, I: IntScalar> StoreElem<Val<I>> for E {
470    fn store(items: &mut [E], index: Val<I>, value: E) {
471        for (position, item) in items.iter_mut().enumerate() {
472            *item = E::select(eq_position(index, position), *item, value);
473        }
474    }
475}
476
477pub fn store<E: StoreElem<I>, I, V: Conv<E>, const N: usize>(items: &mut [E; N], index: I, value: V) {
478    E::store(items, index, value.conv());
479}
480
481// ---- Casts ------------------------------------------------------------------------------------------------------------
482
483/// `value as T`.
484pub trait CastTo<T> {
485    type Out;
486    fn cast_to(self) -> Self::Out;
487}
488
489#[track_caller]
490pub fn cast<T, S: CastTo<T>>(value: S) -> S::Out {
491    value.cast_to()
492}
493
494impl<S: Scalar, T: Scalar> CastTo<T> for Val<S> {
495    type Out = Val<T>;
496    #[track_caller]
497    fn cast_to(self) -> Val<T> {
498        cast_value::<S, T>(self, Location::caller())
499    }
500}
501
502fn cast_value<S: Scalar, T: Scalar>(value: Val<S>, at: &'static Location<'static>) -> Val<T> {
503    // A value with bytes casts through them between integers, unless its one wire already fits the target.
504    let through_bytes = value.is_wide()
505        || value.has_bytes() && S::KIND == Kind::Int && T::KIND == Kind::Int && {
506            let range = with_circuit(|circuit| circuit.range(value.wire()));
507            is_wide::<T>() || !(range.min >= T::MIN && range.max <= T::MAX)
508        };
509    if through_bytes {
510        // `as` between integers keeps the low bits, sign-extended from a signed source: the bytes, resized.
511        return match T::KIND {
512            Kind::Int if is_wide::<T>() => Val::from_limbs(wide::limbs_as::<S, T>(value)),
513            Kind::Int => wide::narrow(&wide::limbs_as::<S, S>(value)),
514            _ => panic!("{at}: casting an integer beyond ±2^24 to {} isn't supported yet", T::NAME),
515        };
516    }
517    let same = Val::<T>::from_wire(value.wire());
518    match (S::KIND, T::KIND) {
519        (Kind::Bool, _) => same,
520        (_, Kind::Bool) => unsupported("only bool casts to bool"),
521        (Kind::Float, Kind::Float) => same,
522        (Kind::Int, Kind::Float) => {
523            let range = with_circuit(|circuit| circuit.range(value.wire()));
524            if range.min >= -16_777_216.0 && range.max <= 16_777_216.0 {
525                same
526            } else {
527                panic!("{at}: casting integers beyond ±2^24 to f32 rounds, which isn't supported yet");
528            }
529        }
530        (Kind::Int, Kind::Int) => {
531            if S::MIN >= T::MIN && S::MAX <= T::MAX {
532                return same;
533            }
534            let range = with_circuit(|circuit| circuit.range(value.wire()));
535            if range.min >= T::MIN && range.max <= T::MAX {
536                return same;
537            }
538            if is_wide::<T>() {
539                // The wrapped value may leave ±2^24, e.g. `-1i32 as u32`, so it is held as bytes.
540                return Val::from_limbs(wide::limbs_as::<S, T>(value));
541            }
542            // `as` between integers keeps the low bits: the same value modulo the target's span.
543            let span = T::MAX - T::MIN + 1.0;
544            integer_table(Val::<T>::from_wire(value.wire()), at, |wide| {
545                let wrapped = (wide as f64 - T::MIN).rem_euclid(span) + T::MIN;
546                wrapped as i64
547            })
548        }
549        (Kind::Float, Kind::Int) => truncate::<T>(value.wire(), at),
550    }
551}
552
553/// `x as T` for an f32 x: truncate toward zero, then saturate to T's range.
554fn truncate<T: Scalar>(value: Wire, at: &'static Location<'static>) -> Val<T> {
555    let range = with_circuit(|circuit| circuit.range(value));
556    if !(range.min > -4_194_304.0 && range.max < 4_194_304.0) {
557        panic!("{at}: casting f32 to an integer needs a declared range within ±2^22, got {}..{}", range.min, range.max);
558    }
559    let truncated = with_circuit(|circuit| {
560        let rounded = emit_round(circuit, value);
561        // Round to nearest, then step back toward zero where rounding went away from it.
562        let rounded_up = emit_compare(circuit, rounded, Relation::Greater, value);
563        let rounded_down = emit_compare(circuit, rounded, Relation::Less, value);
564        let zero = circuit.constant(0.0);
565        let positive = emit_compare(circuit, value, Relation::AtLeast, zero);
566        let too_high = circuit.multiply(rounded_up, positive);
567        let negative = circuit.affine(1.0 - positive);
568        let too_low = circuit.multiply(rounded_down, negative);
569        let wire = circuit.affine(Lin::new(vec![(rounded, 1.0), (too_high, -1.0), (too_low, 1.0)], 0.0));
570        circuit.assume_grid_allowing_readers(wire, 1.0);
571        Val::<T>::from_wire(wire)
572    });
573    if range.min >= T::MIN && range.max <= T::MAX {
574        return truncated;
575    }
576    let low = Val::<T>::constant(T::from_f64(T::MIN));
577    let high = Val::<T>::constant(T::from_f64(T::MAX));
578    let above = Mux::select(gt(truncated, high), truncated, high);
579    Mux::select(gt(low, above), above, low)
580}
581
582macro_rules! native_casts {
583    ($($source:ident),*) => {$(
584        native_casts!(@to $source: u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
585    )*};
586    (@to $source:ident: $($target:ident),*) => {$(
587        impl CastTo<$target> for $source {
588            type Out = $target;
589            fn cast_to(self) -> $target {
590                self as $target
591            }
592        }
593    )*};
594}
595native_casts!(u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);
596
597macro_rules! bool_casts {
598    ($($target:ident),*) => {$(
599        impl CastTo<$target> for bool {
600            type Out = $target;
601            fn cast_to(self) -> $target {
602                self as $target
603            }
604        }
605    )*};
606}
607bool_casts!(u8, u16, u32, u64, usize, i8, i16, i32, i64);
608
609#[cfg(test)]
610mod tests {
611    use super::*;
612    use crate::circuit::Circuit;
613
614    #[test]
615    fn conversions_infer_from_context() {
616        super::super::CONTEXT.with(|context| *context.borrow_mut() = Some(Circuit::new(0)));
617        let value: Val<u8> = Conv::conv(5u8);
618        let picked = if_else(true, || value, || 7);
619        let _: Val<u8> = picked;
620        let native = if_else(false, || 1, || 2);
621        assert_eq!(native, 2);
622        super::super::CONTEXT.with(|context| *context.borrow_mut() = None);
623    }
624}
handwired_rs/dsl/control.rs

handwired_rs/dsl/
control.rs
