tasm_lib/arithmetic/u128/
shift_right_static.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5/// [Shift right][shr] 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/// [shr]: core::ops::Shr
26#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
27pub struct ShiftRightStatic<const N: u8>;
28
29impl<const N: u8> BasicSnippet for ShiftRightStatic<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_right_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: _ v_3 v_2 v_1 v_0
47            // AFTER:  _ (v >> s)_3 (v >> s)_2 (v >> s)_1 (v >> s)_0
48            {self.entrypoint()}:
49                pick 3
50                push {1_u64 << (32 - N)}
51                mul
52                place 3             // _ v3s v2 v1 v0
53                push {1_u64 << (32 - N)}
54                xb_mul              // _ v3s v2s v1s v0s
55
56                /* see non-static version for stack annotations */
57                pick 3
58                split
59                pick 4
60                split
61                pick 5
62                split
63                pick 6
64                split
65
66                pop 1
67                add
68                place 4
69                add
70                place 3
71                add
72                place 2
73
74                return
75        )
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use crate::arithmetic::u128::shift_left_static::ShiftLeftStatic;
83    use crate::test_prelude::*;
84
85    impl<const N: u8> Closure for ShiftRightStatic<N> {
86        type Args = <ShiftLeftStatic<N> as Closure>::Args;
87
88        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
89            let v = pop_encodable::<Self::Args>(stack);
90            push_encodable(stack, &(v >> N));
91        }
92
93        fn pseudorandom_args(
94            &self,
95            seed: [u8; 32],
96            bench_case: Option<BenchmarkCase>,
97        ) -> Self::Args {
98            ShiftLeftStatic::<N>.pseudorandom_args(seed, bench_case)
99        }
100
101        fn corner_case_args(&self) -> Vec<Self::Args> {
102            ShiftLeftStatic::<N>.corner_case_args()
103        }
104    }
105
106    #[test]
107    fn rust_shadow() {
108        macro_rules! test_shift_right_static {
109            ($($i:expr),*$(,)?) => {
110                $(ShadowedClosure::new(ShiftRightStatic::<$i>).test();)*
111            };
112        }
113
114        test_shift_right_static!(0, 1, 2, 3, 4, 5, 6, 7);
115        test_shift_right_static!(8, 9, 10, 11, 12, 13, 14, 15);
116        test_shift_right_static!(16, 17, 18, 19, 20, 21, 22, 23);
117        test_shift_right_static!(24, 25, 26, 27, 28, 29, 30, 31);
118        test_shift_right_static!(32);
119    }
120
121    #[test]
122    #[should_panic]
123    fn shift_beyond_limit() {
124        ShadowedClosure::new(ShiftRightStatic::<33>).test();
125    }
126}
127
128#[cfg(test)]
129mod benches {
130    use super::*;
131    use crate::test_prelude::*;
132
133    #[test]
134    fn benchmark() {
135        ShadowedClosure::new(ShiftRightStatic::<5>).bench();
136    }
137}
tasm_lib/arithmetic/u128/shift_right_static.rs

tasm_lib/arithmetic/u128/
shift_right_static.rs
