Migrate the math structs to be using stwo (#12)

* migrate the math structs to stwo

* clippy, ibutterfly now takes m31 as itwid

Author: weikengchen

Cargo.toml

@@ -12,6 +12,7 @@ sha2 = "0.10.8"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
 stwo-prover = { git = "https://github.com/starkware-libs/stwo", branch = "dev", commit = "2c8b6e5" }
+num-traits = "0.2.0"
 
 [profile.dev]
 opt-level = 3

src/channel/bitcoin_script.rs

@@ -60,12 +60,14 @@ mod test {
     use crate::channel::{Channel, ChannelGadget};
     use crate::channel_commit::Commitment;
     use crate::channel_extract::ExtractorGadget;
-    use crate::math::{CM31, M31, QM31};
     use crate::treepp::*;
     use bitcoin_script::script;
     use rand::{Rng, RngCore, SeedableRng};
     use rand_chacha::ChaCha20Rng;
     use rust_bitcoin_m31::qm31_equalverify;
+    use stwo_prover::core::fields::cm31::CM31;
+    use stwo_prover::core::fields::m31::M31;
+    use stwo_prover::core::fields::qm31::QM31;
 
     #[test]
     fn test_absorb_commitment() {

src/channel/mod.rs

@@ -1,8 +1,8 @@
 use crate::channel_commit::Commitment;
 use crate::channel_extract::{Extraction5M31, ExtractionQM31, Extractor};
-use crate::math::QM31;
 use crate::utils::trim_m31;
 use sha2::{Digest, Sha256};
+use stwo_prover::core::fields::qm31::QM31;
 
 mod bitcoin_script;
 pub use bitcoin_script::*;

src/channel_commit/bitcoin_script.rs

@@ -39,11 +39,13 @@ impl CommitmentGadget {
 #[cfg(test)]
 mod test {
     use crate::channel_commit::{Commitment, CommitmentGadget};
-    use crate::math::{CM31, M31, QM31};
     use crate::treepp::*;
     use bitcoin_script::script;
     use rand::{RngCore, SeedableRng};
     use rand_chacha::ChaCha20Rng;
+    use stwo_prover::core::fields::cm31::CM31;
+    use stwo_prover::core::fields::m31::M31;
+    use stwo_prover::core::fields::qm31::QM31;
 
     #[test]
     fn test_commit_m31() {

src/channel_commit/mod.rs

@@ -1,8 +1,10 @@
 use crate::utils::num_to_bytes;
 use sha2::{Digest, Sha256};
+use stwo_prover::core::fields::cm31::CM31;
+use stwo_prover::core::fields::m31::M31;
+use stwo_prover::core::fields::qm31::QM31;
 
 mod bitcoin_script;
-use crate::math::{CM31, M31, QM31};
 pub use bitcoin_script::*;
 
 /// A commitment, which is a 32-byte SHA256 hash

src/channel_extract/bitcoin_script.rs

@@ -199,11 +199,12 @@ impl ExtractorGadget {
 #[cfg(test)]
 mod test {
     use crate::channel_extract::{Extractor, ExtractorGadget};
-    use crate::math::{Field, M31};
     use crate::treepp::*;
     use bitcoin_script::script;
+    use num_traits::Zero;
     use rand::{Rng, SeedableRng};
     use rand_chacha::ChaCha20Rng;
+    use stwo_prover::core::fields::m31::M31;
 
     #[test]
     fn test_unpack_negative_zero() {
@@ -226,8 +227,6 @@ mod test {
             OP_EQUAL
         };
         let exec_result = execute_script(script);
-        println!("{:8}", exec_result.final_stack);
-        println!("{:?}", exec_result.error);
         assert!(exec_result.success);
     }
 

src/channel_extract/mod.rs

@@ -1,9 +1,11 @@
 use crate::treepp::pushable::{Builder, Pushable};
 use bitcoin::script::PushBytesBuf;
 use core::ops::Neg;
+use stwo_prover::core::fields::cm31::CM31;
+use stwo_prover::core::fields::m31::M31;
+use stwo_prover::core::fields::qm31::QM31;
 
 mod bitcoin_script;
-use crate::math::{CM31, M31, QM31};
 pub use bitcoin_script::*;
 
 /// Basic hint structure for extracting a single qm31 element.

src/circle/bitcoin_script.rs

@@ -1,6 +1,7 @@
-use crate::circle::CirclePoint;
 use crate::treepp::*;
 use rust_bitcoin_m31::{m31_add, m31_mul, m31_neg, m31_sub};
+use stwo_prover::core::circle::CirclePoint;
+use stwo_prover::core::fields::m31::M31;
 
 /// Gadget for points on the circle curve over the m31 field.
 /// This is not the secure field.
@@ -16,7 +17,7 @@ impl CirclePointGadget {
     }
 
     /// Push a constant point.
-    pub fn push(point: &CirclePoint) -> Script {
+    pub fn push(point: &CirclePoint<M31>) -> Script {
         script! {
             { point.x.0 }
             { point.y.0 }
@@ -93,13 +94,13 @@ impl CirclePointGadget {
 
 #[cfg(test)]
 mod test {
-    use crate::circle::CirclePoint;
     use crate::circle::CirclePointGadget;
-    use crate::math::M31;
     use crate::treepp::*;
     use rand_chacha::rand_core::{RngCore, SeedableRng};
     use rand_chacha::ChaCha20Rng;
     use std::ops::{Add, Sub};
+    use stwo_prover::core::circle::CirclePoint;
+    use stwo_prover::core::fields::m31::M31;
 
     #[test]
     fn test_double() {

src/circle/mod.rs

@@ -1,97 +1,2 @@
 mod bitcoin_script;
 pub use bitcoin_script::*;
-
-use std::ops::{Add, Neg, Sub};
-
-use crate::math::M31;
-
-/// A point on the complex circle. Treated as an additive group.
-#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub struct CirclePoint {
-    /// x coordinate.
-    pub x: M31,
-    /// y coordinate.
-    pub y: M31,
-}
-
-impl CirclePoint {
-    /// Push the zero point.
-    pub fn zero() -> Self {
-        Self {
-            x: 1.into(),
-            y: 0.into(),
-        }
-    }
-
-    /// Double a point.
-    pub fn double(&self) -> Self {
-        *self + *self
-    }
-
-    /// Multiply a point with a scalar.
-    pub fn mul(&self, mut scalar: u128) -> CirclePoint {
-        let mut res = Self::zero();
-        let mut cur = *self;
-        while scalar > 0 {
-            if scalar & 1 == 1 {
-                res = res + cur;
-            }
-            cur = cur.double();
-            scalar >>= 1;
-        }
-        res
-    }
-
-    /// Double a point repeatedly for n times.
-    pub fn repeated_double(&self, n: usize) -> Self {
-        let mut res = *self;
-        for _ in 0..n {
-            res = res.double();
-        }
-        res
-    }
-
-    /// Negate a point.
-    pub fn conjugate(&self) -> CirclePoint {
-        Self {
-            x: self.x,
-            y: -self.y,
-        }
-    }
-
-    /// Compute a subgroup generator for points on the circle curve over the m31 field.
-    pub fn subgroup_gen(logn: usize) -> Self {
-        M31_CIRCLE_GEN.repeated_double(M31_CIRCLE_LOG_ORDER - logn)
-    }
-}
-impl Add for CirclePoint {
-    type Output = Self;
-
-    fn add(self, rhs: Self) -> Self::Output {
-        let x = self.x * rhs.x - self.y * rhs.y;
-        let y = self.x * rhs.y + self.y * rhs.x;
-        Self { x, y }
-    }
-}
-impl Neg for CirclePoint {
-    type Output = Self;
-
-    fn neg(self) -> Self::Output {
-        self.conjugate()
-    }
-}
-impl Sub for CirclePoint {
-    type Output = Self;
-
-    fn sub(self, rhs: Self) -> Self::Output {
-        self + (-rhs)
-    }
-}
-
-/// The group order of the points on the circle curve over the m31 field.
-pub const M31_CIRCLE_LOG_ORDER: usize = 31;
-/// A generator of the circle curve over the m31 field.
-pub const M31_CIRCLE_GEN: CirclePoint = CirclePoint {
-    x: M31(2),
-    y: M31(1268011823),
-};

src/circle_secure/bitcoin_script.rs

@@ -1,15 +1,15 @@
 use crate::treepp::*;
+use num_traits::One;
 use rust_bitcoin_m31::{
     m31_add_n31, m31_sub, push_m31_one, push_n31_one, push_qm31_one, qm31_double, qm31_dup,
     qm31_equalverify, qm31_from_bottom, qm31_mul, qm31_neg, qm31_roll, qm31_rot, qm31_square,
     qm31_sub, qm31_swap,
 };
 use std::ops::{Add, Mul, Neg};
+use stwo_prover::core::fields::qm31::QM31;
+use stwo_prover::core::fields::{Field, FieldExpOps};
 
-use crate::{
-    channel::ChannelGadget,
-    math::{Field, QM31},
-};
+use crate::channel::ChannelGadget;
 
 /// Gadget for points on the circle curve in the qm31 field.
 pub struct CirclePointSecureGadget;
@@ -104,18 +104,20 @@ impl CirclePointSecureGadget {
 
 #[cfg(test)]
 mod test {
+    use num_traits::One;
     use std::ops::{Add, Mul, Neg};
 
     use crate::treepp::*;
     use rand::{Rng, RngCore, SeedableRng};
     use rand_chacha::ChaCha20Rng;
     use rust_bitcoin_m31::qm31_equalverify;
+    use stwo_prover::core::fields::m31::M31;
+    use stwo_prover::core::fields::qm31::QM31;
+    use stwo_prover::core::fields::{Field, FieldExpOps};
 
     use crate::{
-        channel::Channel,
-        channel_extract::ExtractorGadget,
+        channel::Channel, channel_extract::ExtractorGadget,
         circle_secure::bitcoin_script::CirclePointSecureGadget,
-        math::{Field, M31, QM31},
     };
 
     #[test]

src/constraints/bitcoin_script.rs

@@ -1,8 +1,7 @@
 use crate::treepp::*;
 use rust_bitcoin_m31::{qm31_add, qm31_mul, qm31_swap};
 use std::ops::{Add, Mul, Neg};
-
-use crate::math::QM31;
+use stwo_prover::core::fields::qm31::QM31;
 
 /// Gadget for constraints over the circle curve
 pub struct ConstraintsGadget;
@@ -48,8 +47,8 @@ mod test {
     use rand::{RngCore, SeedableRng};
     use rand_chacha::ChaCha20Rng;
     use rust_bitcoin_m31::qm31_equalverify;
-
-    use crate::math::{M31, QM31};
+    use stwo_prover::core::fields::m31::M31;
+    use stwo_prover::core::fields::qm31::QM31;
 
     #[test]
     fn test_pair_vanishing() {

src/fri/bitcoin_script.rs

@@ -227,19 +227,21 @@ impl FRIGadget {
 #[cfg(test)]
 mod test {
     use crate::channel::Channel;
-    use crate::circle::CirclePoint;
     use crate::fri;
     use crate::fri::{FRIGadget, N_QUERIES};
-    use crate::math::Field;
     use crate::treepp::*;
     use crate::twiddle_merkle_tree::{TwiddleMerkleTree, TWIDDLE_MERKLE_TREE_ROOT_18};
     use crate::utils::permute_eval;
     use bitcoin::hashes::Hash;
     use bitcoin::{TapLeafHash, Transaction};
     use bitcoin_scriptexec::{Exec, ExecCtx, Experimental, Options, TxTemplate};
+    use num_traits::One;
     use rand::{Rng, SeedableRng};
     use rand_chacha::ChaCha20Rng;
     use rust_bitcoin_m31::qm31_equalverify;
+    use stwo_prover::core::circle::CirclePointIndex;
+    use stwo_prover::core::fields::m31::M31;
+    use stwo_prover::core::fields::FieldExpOps;
 
     #[test]
     fn test_fiat_shamir() {
@@ -255,15 +257,15 @@ mod test {
         let logn = 19;
 
         let proof = {
-            let p = CirclePoint::subgroup_gen(logn + 1);
+            let p = CirclePointIndex::subgroup_gen(logn as u32 + 1).to_point();
 
             let mut prng = ChaCha20Rng::seed_from_u64(0);
 
             let mut channel_init_state = [0u8; 32];
             channel_init_state.iter_mut().for_each(|v| *v = prng.gen());
 
             let evaluation = (0..(1 << logn))
-                .map(|i| (p.mul(i * 2 + 1).x.square().square() + 1.into()).into())
+                .map(|i| (p.mul(i * 2 + 1).x.square().square() + M31::one()).into())
                 .collect();
             let evaluation = permute_eval(evaluation);
 
@@ -325,10 +327,10 @@ mod test {
         channel_init_state.iter_mut().for_each(|v| *v = prng.gen());
 
         let proof = {
-            let p = CirclePoint::subgroup_gen(logn + 1);
+            let p = CirclePointIndex::subgroup_gen(logn as u32 + 1).to_point();
 
             let evaluation = (0..(1 << logn))
-                .map(|i| (p.mul(i * 2 + 1).x.square().square() + 1.into()).into())
+                .map(|i| (p.mul(i * 2 + 1).x.square().square() + M31::one()).into())
                 .collect();
             let evaluation = permute_eval(evaluation);
 
@@ -388,10 +390,10 @@ mod test {
         channel_init_state.iter_mut().for_each(|v| *v = prng.gen());
 
         let proof = {
-            let p = CirclePoint::subgroup_gen(logn + 1);
+            let p = CirclePointIndex::subgroup_gen(logn as u32 + 1).to_point();
 
             let evaluation = (0..(1 << logn))
-                .map(|i| (p.mul(i * 2 + 1).x.square().square() + 1.into()).into())
+                .map(|i| (p.mul(i * 2 + 1).x.square().square() + M31::one()).into())
                 .collect();
             let evaluation = permute_eval(evaluation);
 
@@ -448,10 +450,10 @@ mod test {
         channel_init_state.iter_mut().for_each(|v| *v = prng.gen());
 
         let proof = {
-            let p = CirclePoint::subgroup_gen(logn + 1);
+            let p = CirclePointIndex::subgroup_gen(logn as u32 + 1).to_point();
 
             let evaluation = (0..(1 << logn))
-                .map(|i| (p.mul(i * 2 + 1).x.square().square() + 1.into()).into())
+                .map(|i| (p.mul(i * 2 + 1).x.square().square() + M31::one()).into())
                 .collect();
             let evaluation = permute_eval(evaluation);
 
@@ -527,15 +529,15 @@ mod test {
         let logn = 19;
 
         let proof = {
-            let p = CirclePoint::subgroup_gen(logn + 1);
+            let p = CirclePointIndex::subgroup_gen(logn as u32 + 1).to_point();
 
             let mut prng = ChaCha20Rng::seed_from_u64(0);
 
             let mut channel_init_state = [0u8; 32];
             channel_init_state.iter_mut().for_each(|v| *v = prng.gen());
 
             let evaluation = (0..(1 << logn))
-                .map(|i| (p.mul(i * 2 + 1).x.square().square() + 1.into()).into())
+                .map(|i| (p.mul(i * 2 + 1).x.square().square() + M31::one()).into())
                 .collect();
             let evaluation = permute_eval(evaluation);