Add tests of XFE exponentiation snippets

TritonVM · Feb 26, 2024 · 98c3e96 · 98c3e96
1 parent 98d57fc
commit 98c3e96
Show file tree

Hide file tree

Showing 3 changed files with 145 additions and 1 deletion.
diff --git a/tasm-lib/src/arithmetic/xfe/cube.rs b/tasm-lib/src/arithmetic/xfe/cube.rs
@@ -42,14 +42,17 @@ mod tests {
     use itertools::Itertools;
     use num::One;
     use num::Zero;
+    use rand::random;
     use rand::rngs::StdRng;
     use rand::Rng;
     use rand::SeedableRng;
     use triton_vm::twenty_first::shared_math::traits::ModPowU32;
     use triton_vm::twenty_first::shared_math::x_field_element::EXTENSION_DEGREE;
 
     use super::*;
+    use crate::arithmetic::xfe::mod_pow_u32_generic::XfeModPowU32Generic;
     use crate::snippet_bencher::BenchmarkCase;
+    use crate::test_helpers::test_rust_equivalence_given_complete_state;
     use crate::traits::closure::Closure;
     use crate::traits::closure::ShadowedClosure;
     use crate::traits::rust_shadow::RustShadow;
@@ -104,6 +107,51 @@ mod tests {
     fn cube_xfe_pbt() {
         ShadowedClosure::new(Cube).test()
     }
+
+    #[test]
+    fn compare_to_generic_pow_u32() {
+        // Run `cube` snippet
+        let input: XFieldElement = random();
+        let init_stack_to_fourth = [
+            Cube.init_stack_for_isolated_run(),
+            input.coefficients.into_iter().rev().collect_vec(),
+        ]
+        .concat();
+        let final_state_from_to_fourth = test_rust_equivalence_given_complete_state(
+            &ShadowedClosure::new(Cube),
+            &init_stack_to_fourth,
+            &[],
+            &NonDeterminism::default(),
+            &None,
+            None,
+        );
+
+        // Run snippet for generic pow
+        let init_stack_to_generic = [
+            XfeModPowU32Generic.init_stack_for_isolated_run(),
+            vec![BFieldElement::new(3)],
+            input.coefficients.into_iter().rev().collect_vec(),
+        ]
+        .concat();
+        let final_state_from_generic = test_rust_equivalence_given_complete_state(
+            &ShadowedClosure::new(XfeModPowU32Generic),
+            &init_stack_to_generic,
+            &[],
+            &NonDeterminism::default(),
+            &None,
+            None,
+        );
+
+        // Assert that height agrees, and the top-3 elements agree
+        assert_eq!(
+            final_state_from_generic.final_stack.len(),
+            final_state_from_to_fourth.final_stack.len()
+        );
+        assert_eq!(
+            final_state_from_generic.final_stack[16..=18],
+            final_state_from_to_fourth.final_stack[16..=18],
+        );
+    }
 }
 
 #[cfg(test)]

diff --git a/tasm-lib/src/arithmetic/xfe/square.rs b/tasm-lib/src/arithmetic/xfe/square.rs
@@ -38,13 +38,16 @@ mod tests {
     use itertools::Itertools;
     use num::One;
     use num::Zero;
+    use rand::random;
     use rand::rngs::StdRng;
     use rand::Rng;
     use rand::SeedableRng;
     use triton_vm::twenty_first::shared_math::x_field_element::EXTENSION_DEGREE;
 
     use super::*;
+    use crate::arithmetic::xfe::mod_pow_u32_generic::XfeModPowU32Generic;
     use crate::snippet_bencher::BenchmarkCase;
+    use crate::test_helpers::test_rust_equivalence_given_complete_state;
     use crate::traits::closure::Closure;
     use crate::traits::closure::ShadowedClosure;
     use crate::traits::rust_shadow::RustShadow;
@@ -99,6 +102,51 @@ mod tests {
     fn square_xfe_pbt() {
         ShadowedClosure::new(Square).test()
     }
+
+    #[test]
+    fn compare_to_generic_pow_u32() {
+        // Run `cube` snippet
+        let input: XFieldElement = random();
+        let init_stack_to_fourth = [
+            Square.init_stack_for_isolated_run(),
+            input.coefficients.into_iter().rev().collect_vec(),
+        ]
+        .concat();
+        let final_state_from_to_fourth = test_rust_equivalence_given_complete_state(
+            &ShadowedClosure::new(Square),
+            &init_stack_to_fourth,
+            &[],
+            &NonDeterminism::default(),
+            &None,
+            None,
+        );
+
+        // Run snippet for generic pow
+        let init_stack_to_generic = [
+            XfeModPowU32Generic.init_stack_for_isolated_run(),
+            vec![BFieldElement::new(2)],
+            input.coefficients.into_iter().rev().collect_vec(),
+        ]
+        .concat();
+        let final_state_from_generic = test_rust_equivalence_given_complete_state(
+            &ShadowedClosure::new(XfeModPowU32Generic),
+            &init_stack_to_generic,
+            &[],
+            &NonDeterminism::default(),
+            &None,
+            None,
+        );
+
+        // Assert that height agrees, and the top-3 elements agree
+        assert_eq!(
+            final_state_from_generic.final_stack.len(),
+            final_state_from_to_fourth.final_stack.len()
+        );
+        assert_eq!(
+            final_state_from_generic.final_stack[16..=18],
+            final_state_from_to_fourth.final_stack[16..=18],
+        );
+    }
 }
 
 #[cfg(test)]

diff --git a/tasm-lib/src/arithmetic/xfe/to_the_fourth.rs b/tasm-lib/src/arithmetic/xfe/to_the_fourth.rs
@@ -43,14 +43,17 @@ mod tests {
     use itertools::Itertools;
     use num::One;
     use num::Zero;
+    use rand::random;
     use rand::rngs::StdRng;
     use rand::Rng;
     use rand::SeedableRng;
     use triton_vm::twenty_first::shared_math::traits::ModPowU32;
     use triton_vm::twenty_first::shared_math::x_field_element::EXTENSION_DEGREE;
 
     use super::*;
+    use crate::arithmetic::xfe::mod_pow_u32_generic::XfeModPowU32Generic;
     use crate::snippet_bencher::BenchmarkCase;
+    use crate::test_helpers::test_rust_equivalence_given_complete_state;
     use crate::traits::closure::Closure;
     use crate::traits::closure::ShadowedClosure;
     use crate::traits::rust_shadow::RustShadow;
@@ -103,7 +106,52 @@ mod tests {
 
     #[test]
     fn to_the_fourth_xfe_pbt() {
-        ShadowedClosure::new(ToTheFourth).test()
+        ShadowedClosure::new(ToTheFourth).test();
+    }
+
+    #[test]
+    fn compare_to_generic_pow_u32() {
+        // Run `to_the_fourth` snippet
+        let input: XFieldElement = random();
+        let init_stack_to_fourth = [
+            ToTheFourth.init_stack_for_isolated_run(),
+            input.coefficients.into_iter().rev().collect_vec(),
+        ]
+        .concat();
+        let final_state_from_to_fourth = test_rust_equivalence_given_complete_state(
+            &ShadowedClosure::new(ToTheFourth),
+            &init_stack_to_fourth,
+            &[],
+            &NonDeterminism::default(),
+            &None,
+            None,
+        );
+
+        // Run snippet for generic pow
+        let init_stack_to_generic = [
+            XfeModPowU32Generic.init_stack_for_isolated_run(),
+            vec![BFieldElement::new(4)],
+            input.coefficients.into_iter().rev().collect_vec(),
+        ]
+        .concat();
+        let final_state_from_generic = test_rust_equivalence_given_complete_state(
+            &ShadowedClosure::new(XfeModPowU32Generic),
+            &init_stack_to_generic,
+            &[],
+            &NonDeterminism::default(),
+            &None,
+            None,
+        );
+
+        // Assert that height agrees, and the top-3 elements agree
+        assert_eq!(
+            final_state_from_generic.final_stack.len(),
+            final_state_from_to_fourth.final_stack.len()
+        );
+        assert_eq!(
+            final_state_from_generic.final_stack[16..=18],
+            final_state_from_to_fourth.final_stack[16..=18],
+        );
     }
 }