extra_safe/
traits.rs

1//! This module contains the traits for the IOPattern type, along with the implementation of type-level operations checking on their correct usage.
2//! The IOPattern type is a type-level HList of IOWords, which are either Absorb or Squeeze.
3//! The main operations are Normalize, which merges successive words of the same type, and Consume, which takes a word and an IOPattern, and checks whether
4//! it is legal to use the operation of this word on the tip of the IOPattern.
5//! This is explained in more detail in [the spec document][1].
6//!
7//! [1]: https://hackmd.io/bHgsH6mMStCVibM_wYvb2w#SAFE-Sponge-API-for-Field-Elements-%E2%80%93-A-Toolbox-for-ZK-Hash-Applications
8
9use core::ops::{Add, Sub};
10use std::marker::PhantomData;
11pub use typenum;
12use typenum::{Bit, Diff, Sum, UInt, Unsigned, U0};
13
14// Our two alternatives for the IOPattern, i.e. these are IOWords
15// Note the phantom type avoids allocating actual data.
16
17/// The type-level Absorb operation
18#[derive(Debug)]
19pub struct Absorb<N>(PhantomData<N>);
20/// The type-level Squeeze operation
21#[derive(Debug)]
22pub struct Squeeze<N>(PhantomData<N>);
23
24/// Our trait for common treatment of both patterns
25pub trait IOWord: private::Sealed {}
26
27impl<N: Unsigned> IOWord for Absorb<N> {}
28impl<N: Unsigned> IOWord for Squeeze<N> {}
29
30/// Type-level HList, specialized to IOWord
31/// using  a sealed trait
32/// See e.g. `<https://hackage.haskell.org/package/heterolist>` (or frunk) for
33/// what a HList is.
34pub trait List: private::Sealed {
35    /// This is an inhabitant of the List type corresponding to the
36    /// Self type
37    fn unit() -> Self;
38
39    /// THis returns whether the list is empty
40    fn is_empty() -> bool;
41}
42
43impl<Item: IOWord, Next: List> List for Cons<Item, Next> {
44    fn unit() -> Self {
45        Cons {
46            _phantom: PhantomData,
47        }
48    }
49    fn is_empty() -> bool {
50        false
51    }
52}
53
54impl List for Nil {
55    fn unit() -> Self {
56        Nil
57    }
58    fn is_empty() -> bool {
59        true
60    }
61}
62
63/// The concrete type constructor for our HList trait
64#[derive(Debug)]
65pub struct Cons<Item, Next: List> {
66    _phantom: PhantomData<(Item, Next)>,
67}
68
69/// The concrete type for our empty HList representant
70#[derive(Debug)]
71pub struct Nil;
72
73/// Convenience helper for creating an instance of List
74#[macro_export]
75macro_rules! iopat {
76    () => { $crate::traits::Nil };
77    (...$rest:ty) => { $rest };
78    ($a:ty) => { $crate::iopat![$a,] };
79    ($a:ty, $($tok:tt)*) => {
80        $crate::traits::Cons<$a, $crate::iopat![$($tok)*]>
81    };
82}
83
84// an IOPattern is a List of IOWords .. (TODO: does this need elaboration?)
85
86/// Normalizing an IOPattern with merge operations applied recursively
87pub trait Normalize: List {
88    /// The output of the normalization
89    type Output: List;
90}
91
92/// Convenience trait for projection of Normalize
93pub type Norm<T> = <T as Normalize>::Output;
94
95// We unfold the type-level cases of the recurrence
96impl Normalize for Nil {
97    type Output = Nil;
98}
99
100// Head zero elimination
101impl<L: Normalize> Normalize for Cons<Squeeze<U0>, L> {
102    type Output = Norm<L>;
103}
104
105impl<L: Normalize> Normalize for Cons<Absorb<U0>, L> {
106    type Output = Norm<L>;
107}
108
109// Base cases
110impl<U: Unsigned, B: Bit> Normalize for Cons<Absorb<UInt<U, B>>, Nil> {
111    type Output = Cons<Absorb<UInt<U, B>>, Nil>;
112}
113impl<U: Unsigned, B: Bit> Normalize for Cons<Squeeze<UInt<U, B>>, Nil> {
114    type Output = Cons<Squeeze<UInt<U, B>>, Nil>;
115}
116
117// Non-head-zero recursive cases: concatenation
118impl<U: Unsigned, B: Bit, M: Unsigned, T: List> Normalize
119    for Cons<Absorb<UInt<U, B>>, Cons<Absorb<M>, T>>
120where
121    UInt<U, B>: Add<M>, // present for all values in practice
122    Cons<Absorb<Sum<UInt<U, B>, M>>, T>: Normalize,
123{
124    type Output = Norm<Cons<Absorb<Sum<UInt<U, B>, M>>, T>>;
125}
126
127impl<U: Unsigned, B: Bit, M: Unsigned, T: List> Normalize
128    for Cons<Squeeze<UInt<U, B>>, Cons<Squeeze<M>, T>>
129where
130    UInt<U, B>: Add<M>, // present for all reasonable values in practice
131    Cons<Squeeze<Sum<UInt<U, B>, M>>, T>: Normalize,
132{
133    type Output = Norm<Cons<Squeeze<Sum<UInt<U, B>, M>>, T>>;
134}
135
136// Non-head-zero recursive cases: no concatenation
137// This requires introspection into the UInt / UTerm type
138impl<U: Unsigned, B: Bit, U2: Unsigned, B2: Bit, T: List> Normalize
139    for Cons<Squeeze<UInt<U, B>>, Cons<Absorb<UInt<U2, B2>>, T>>
140where
141    Cons<Absorb<UInt<U2, B2>>, T>: Normalize,
142{
143    type Output = Cons<Squeeze<UInt<U, B>>, Norm<Cons<Absorb<UInt<U2, B2>>, T>>>;
144}
145
146impl<U: Unsigned, B: Bit, U2: Unsigned, B2: Bit, T: List> Normalize
147    for Cons<Absorb<UInt<U, B>>, Cons<Squeeze<UInt<U2, B2>>, T>>
148where
149    Cons<Squeeze<UInt<U2, B2>>, T>: Normalize,
150{
151    type Output = Cons<Absorb<UInt<U, B>>, Norm<Cons<Squeeze<UInt<U2, B2>>, T>>>;
152}
153
154// and in case the head is zero
155impl<U: Unsigned, B: Bit, T: List> Normalize for Cons<Squeeze<UInt<U, B>>, Cons<Absorb<U0>, T>>
156where
157    Cons<Squeeze<UInt<U, B>>, T>: Normalize,
158{
159    type Output = Norm<Cons<Squeeze<UInt<U, B>>, T>>;
160}
161
162impl<U: Unsigned, B: Bit, T: List> Normalize for Cons<Absorb<UInt<U, B>>, Cons<Squeeze<U0>, T>>
163where
164    Cons<Absorb<UInt<U, B>>, T>: Normalize,
165{
166    type Output = Norm<Cons<Absorb<UInt<U, B>>, T>>;
167}
168
169/// Emptying an IOPattern using an IOWord. This assumes that it is working
170/// with a list in head-normal form (i.e. the first element cannot be merged
171/// with the immediately following list). All lists that have been normalized
172/// are in head-normal form.
173pub trait Consume<Op: IOWord> {
174    /// The output of the consumption
175    type Output: List;
176}
177
178/// Convenience trait for projection of Consume
179#[allow(dead_code)]
180pub type Use<T, U> = <T as Consume<U>>::Output;
181
182// We unfold the type-level cases of the recurrence
183
184// If the consumer is larger than the head pattern, we get to something
185// impossible, because we assume this is only called on normalized lists
186
187// If we get to U0, we end
188impl<N, T: List> Consume<Absorb<U0>> for Cons<Absorb<N>, T>
189where
190    N: Unsigned,
191{
192    type Output = Self;
193}
194
195impl<N, T: List> Consume<Squeeze<U0>> for Cons<Squeeze<N>, T>
196where
197    N: Unsigned,
198{
199    type Output = Self;
200}
201
202// Otherwise, we simplify
203impl<U, B, N, T> Consume<Absorb<UInt<U, B>>> for Cons<Absorb<N>, T>
204where
205    U: Unsigned,
206    B: Bit,
207    N: Unsigned,
208    T: List,
209    N: Sub<UInt<U, B>>, // present for N >= UInt<U, B>
210    Cons<Absorb<Diff<N, UInt<U, B>>>, T>: Normalize,
211{
212    type Output = Norm<Cons<Absorb<Diff<N, UInt<U, B>>>, T>>;
213}
214
215impl<U, B, N, T> Consume<Squeeze<UInt<U, B>>> for Cons<Squeeze<N>, T>
216where
217    U: Unsigned,
218    B: Bit,
219    N: Unsigned,
220    T: List,
221    N: Sub<UInt<U, B>>, // present for N >= UInt<U, B>
222    Cons<Squeeze<Diff<N, UInt<U, B>>>, T>: Normalize,
223{
224    type Output = Norm<Cons<Squeeze<Diff<N, UInt<U, B>>>, T>>;
225}
226
227// Seal the traits so that the above defines admissible implementations of sealed traits
228mod private {
229    pub trait Sealed {}
230
231    impl<N> Sealed for super::Absorb<N> {}
232    impl<N> Sealed for super::Squeeze<N> {}
233
234    impl Sealed for super::Nil {}
235    impl<H, T: super::List> Sealed for super::Cons<H, T> {}
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use typenum::assert_type_eq;
242    use typenum::{U1, U2, U3, U4, U5, U6};
243
244    #[test]
245    fn normalizes() {
246        // Sum cases
247        assert_type_eq!(Norm<iopat![Absorb<U2>, Absorb<U3>]>, iopat![Absorb<U5>]);
248        assert_type_eq!(Norm<iopat![Squeeze<U2>, Squeeze<U3>]>, iopat![Squeeze<U5>]);
249        // recursion cases
250        assert_type_eq!(
251            Norm<iopat![Squeeze<U2>, Absorb<U3>]>,
252            iopat![Squeeze<U2>, Absorb<U3>]
253        );
254        assert_type_eq!(
255            Norm<iopat![Squeeze<U2>, Squeeze<U3>, Absorb<U1>]>,
256            iopat![Squeeze<U5>, Absorb<U1>]
257        );
258        // zero elision at the head
259        assert_type_eq!(Norm<iopat![Absorb<U0>, Absorb<U3>]>, iopat![Absorb<U3>]);
260        assert_type_eq!(Norm<iopat![Absorb<U0>, Squeeze<U3>]>, iopat![Squeeze<U3>]);
261        assert_type_eq!(Norm<iopat![Squeeze<U0>, Squeeze<U3>]>, iopat![Squeeze<U3>]);
262        assert_type_eq!(Norm<iopat![Squeeze<U0>, Absorb<U3>]>, iopat![Absorb<U3>]);
263        // zero elision in recursive cases
264        assert_type_eq!(
265            Norm<iopat![Absorb<U3>, Squeeze<U0>, Absorb<U1>]>,
266            iopat![Absorb<U4>]
267        );
268        assert_type_eq!(
269            Norm<iopat![Squeeze<U3>, Absorb<U0>, Squeeze<U1>]>,
270            iopat![Squeeze<U4>]
271        );
272    }
273
274    #[test]
275    fn uses() {
276        // Substraction
277        assert_type_eq!(
278            Use<iopat![Absorb<U5>], Absorb<U2>>,
279            iopat![Absorb<U3>]
280        );
281        assert_type_eq!(
282            Use<iopat![Absorb<U5>, Squeeze<U2>], Absorb<U2>>,
283            iopat![Absorb<U3>, Squeeze<U2>]
284        );
285
286        assert_type_eq!(
287            Use<iopat![Squeeze<U5>], Squeeze<U2>>,
288            iopat![Squeeze<U3>]
289        );
290        assert_type_eq!(
291            Use<iopat![Squeeze<U5>, Absorb<U2>], Squeeze<U2>>,
292            iopat![Squeeze<U3>, Absorb<U2>]
293        );
294
295        // Zero-simplification
296        assert_type_eq!(Use<iopat![Absorb<U5>], Absorb<U5>>, Nil);
297        assert_type_eq!(
298            Use<iopat![Absorb<U5>], Absorb<U0>>,
299            iopat![Absorb<U5>]
300        );
301        assert_type_eq!(
302            Use<iopat![Squeeze<U5>], Squeeze<U0>>,
303            iopat![Squeeze<U5>]
304        );
305        assert_type_eq!(
306            Use<iopat![Absorb<U3>, Squeeze<U2>], Absorb<U3>>,
307            iopat![Squeeze<U2>]
308        );
309
310        // This doesn't work: you have to call on (head-)normalized lists
311        // initially
312        /*
313        assert_type_eq!(
314            Use<iopat![Absorb<U5>, Absorb<U1>], Absorb<U6>>,
315            Nil
316        );
317        */
318
319        // This, however, works
320        assert_type_eq!(
321            Use<Use<iopat![Squeeze<U3>, Absorb<U5>, Absorb<U1>], Squeeze<U3>>, Absorb<U6>>,
322            Nil
323        );
324    }
325}
extra_safe/traits.rs

extra_safe/
traits.rs
