Skip to main content

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 parameters(&self) -> Vec<(DataType, String)> {
31        vec![(DataType::U128, "value".to_string())]
32    }
33
34    fn return_values(&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 rand::rngs::StdRng;
80
81    use super::*;
82    use crate::test_prelude::*;
83
84    impl<const N: u8> Closure for ShiftLeftStatic<N> {
85        type Args = u128;
86
87        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
88            let v = pop_encodable::<Self::Args>(stack)?;
89            push_encodable(stack, &(v << N));
90            Ok(())
91        }
92
93        fn pseudorandom_args(
94            &self,
95            seed: [u8; 32],
96            bench_case: Option<BenchmarkCase>,
97        ) -> Self::Args {
98            match bench_case {
99                Some(BenchmarkCase::CommonCase) => 0x1282,
100                Some(BenchmarkCase::WorstCase) => 0x123456789abcdef,
101                None => StdRng::from_seed(seed).random(),
102            }
103        }
104
105        fn corner_case_args(&self) -> Vec<Self::Args> {
106            vec![0, 1, 8, u32::MAX.into(), u64::MAX.into(), u128::MAX]
107        }
108    }
109
110    #[macro_rules_attr::apply(test)]
111    fn rust_shadow() {
112        macro_rules! test_shift_left_static {
113            ($($i:expr),*$(,)?) => {
114                $(ShadowedClosure::new(ShiftLeftStatic::<$i>).test();)*
115            };
116        }
117
118        test_shift_left_static!(0, 1, 2, 3, 4, 5, 6, 7);
119        test_shift_left_static!(8, 9, 10, 11, 12, 13, 14, 15);
120        test_shift_left_static!(16, 17, 18, 19, 20, 21, 22, 23);
121        test_shift_left_static!(24, 25, 26, 27, 28, 29, 30, 31);
122        test_shift_left_static!(32);
123    }
124
125    #[macro_rules_attr::apply(test)]
126    #[should_panic]
127    fn shift_beyond_limit() {
128        ShadowedClosure::new(ShiftLeftStatic::<33>).test();
129    }
130}
131
132#[cfg(test)]
133mod benches {
134    use super::*;
135    use crate::test_prelude::*;
136
137    #[macro_rules_attr::apply(test)]
138    fn benchmark() {
139        ShadowedClosure::new(ShiftLeftStatic::<5>).bench();
140    }
141}