tasm_lib/arithmetic/u128/
shift_left_static.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5/// [Shift left][shl] for unsigned 128-bit integers, with the shift amount
6/// specified at compile time.
7///
8/// The shift amount `N` must be in range `0..=32`.
9///
10/// ### Behavior
11///
12/// ```text
13/// BEFORE: _ [v: u128]
14/// AFTER:  _ [v << SHIFT_AMOUNT: u128]
15/// ```
16///
17/// ### Preconditions
18///
19/// - input argument `arg` is properly [`BFieldCodec`] encoded
20///
21/// ### Postconditions
22///
23/// - the output is properly [`BFieldCodec`] encoded
24///
25/// [shl]: core::ops::Shl
26#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
27pub struct ShiftLeftStatic<const N: u8>;
28
29impl<const N: u8> BasicSnippet for ShiftLeftStatic<N> {
30    fn inputs(&self) -> Vec<(DataType, String)> {
31        vec![(DataType::U128, "value".to_string())]
32    }
33
34    fn outputs(&self) -> Vec<(DataType, String)> {
35        vec![(DataType::U128, "shifted_value".to_string())]
36    }
37
38    fn entrypoint(&self) -> String {
39        format!("tasmlib_arithmetic_u128_shift_left_static_{N}")
40    }
41
42    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
43        assert!(N <= 32, "shift amount must be in range 0..=32");
44
45        triton_asm!(
46            // BEFORE: _ [value: u128]
47            // AFTER:  _ [value << shift: u128]
48            {self.entrypoint()}:
49                push {1_u64 << N}
50                xb_mul          // _ v3 (v2<<shift) (v1<<shift) (v0<<shift)
51                pick 3          // _ (v2<<shift) (v1<<shift) (v0<<shift) v3
52                push {1_u64 << N}
53                mul             // _ (v2<<shift) (v1<<shift) (v0<<shift) (v3<<shift)
54                                // _  v2s         v1s         v0s         v3s
55
56                split           // _ v2s v1s v0s v3s_hi v3s_lo
57                pick 4          // _ v1s v0s v3s_hi v3s_lo v2s
58                split           // _ v1s v0s v3s_hi v3s_lo v2s_hi v2s_lo
59                pick 5          // _ v0s v3s_hi v3s_lo v2s_hi v2s_lo v1s
60                split           // _ v0s v3s_hi v3s_lo v2s_hi v2s_lo v1s_hi v1s_lo
61                pick 6          // _ v3s_hi v3s_lo v2s_hi v2s_lo v1s_hi v1s_lo v0s
62                split           // _ v3s_hi v3s_lo v2s_hi v2s_lo v1s_hi v1s_lo v0s_hi v0'
63
64                place 7         // _ v0' v3s_hi v3s_lo v2s_hi v2s_lo v1s_hi v1s_lo v0s_hi
65                add             // _ v0' v3s_hi v3s_lo v2s_hi v2s_lo v1s_hi v1'
66                place 6         // _ v1' v0' v3s_hi v3s_lo v2s_hi v2s_lo v1s_hi
67                add             // _ v1' v0' v3s_hi v3s_lo v2s_hi v2'
68                place 5         // _ v2' v1' v0' v3s_hi v3s_lo v2s_hi
69                add             // _ v2' v1' v0' v3s_hi v3'
70                place 4         // _ v3' v2' v1' v0' v3s_hi
71                pop 1           // _ v3' v2' v1' v0'
72                return
73        )
74    }
75}
76
77#[cfg(test)]
78pub(crate) mod tests {
79    use super::*;
80    use crate::test_prelude::*;
81    use rand::rngs::StdRng;
82
83    impl<const N: u8> Closure for ShiftLeftStatic<N> {
84        type Args = u128;
85
86        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
87            let v = pop_encodable::<Self::Args>(stack);
88            push_encodable(stack, &(v << N));
89        }
90
91        fn pseudorandom_args(
92            &self,
93            seed: [u8; 32],
94            bench_case: Option<BenchmarkCase>,
95        ) -> Self::Args {
96            match bench_case {
97                Some(BenchmarkCase::CommonCase) => 0x1282,
98                Some(BenchmarkCase::WorstCase) => 0x123456789abcdef,
99                None => StdRng::from_seed(seed).random(),
100            }
101        }
102
103        fn corner_case_args(&self) -> Vec<Self::Args> {
104            vec![0, 1, 8, u32::MAX.into(), u64::MAX.into(), u128::MAX]
105        }
106    }
107
108    #[test]
109    fn rust_shadow() {
110        macro_rules! test_shift_left_static {
111            ($($i:expr),*$(,)?) => {
112                $(ShadowedClosure::new(ShiftLeftStatic::<$i>).test();)*
113            };
114        }
115
116        test_shift_left_static!(0, 1, 2, 3, 4, 5, 6, 7);
117        test_shift_left_static!(8, 9, 10, 11, 12, 13, 14, 15);
118        test_shift_left_static!(16, 17, 18, 19, 20, 21, 22, 23);
119        test_shift_left_static!(24, 25, 26, 27, 28, 29, 30, 31);
120        test_shift_left_static!(32);
121    }
122
123    #[test]
124    #[should_panic]
125    fn shift_beyond_limit() {
126        ShadowedClosure::new(ShiftLeftStatic::<33>).test();
127    }
128}
129
130#[cfg(test)]
131mod benches {
132    use super::*;
133    use crate::test_prelude::*;
134
135    #[test]
136    fn benchmark() {
137        ShadowedClosure::new(ShiftLeftStatic::<5>).bench();
138    }
139}