Draw large queries in sha256 and shift signed mask points

src/air/bitcoin_script.rs

@@ -39,6 +39,39 @@
         }
     }
 
+    /// Mask a point by shifting it with signed mask
+    ///
+    /// Input:
+    /// -  point (in qm31)
+    ///
+    /// Output:
+    /// -  shifted point (in qm31)
+    pub fn shifted_signed_mask_points(
+        mask: &ColumnVec<Vec<isize>>,
+        domains: &[CanonicCoset],
+    ) -> Script {
+        let mut shifted = vec![];
+        for (mask_entry, domain) in mask.iter().zip(domains.iter()) {
+            let step_point = domain.step();
+            for mask_item in mask_entry.iter() {
+                shifted.push(step_point.mul_signed(*mask_item));
+            }
+        }
+
+        script! {
+            if !shifted.is_empty() {
+                for elem in shifted.iter().take(shifted.len() - 1) {
+                    { CirclePointGadget::dup() }
+                    { CirclePointGadget::add_constant_m31_point(elem) }
+                    { CirclePointGadget::swap() }
+                }
+                { CirclePointGadget::add_constant_m31_point(shifted.last().unwrap()) }
+            } else {
+                { CirclePointGadget::drop() }
+            }
+        }
+    }
+
     /// Combine the evaluation on four points into one.
     ///
     /// Input:
@@ -65,13 +98,16 @@
 
 #[cfg(test)]
 mod test {
-    use crate::air::AirGadget;
+    use crate::air::{shifted_signed_mask_points, AirGadget};
     use crate::circle::CirclePointGadget;
     use crate::treepp::*;
+    use crate::utils::get_rand_qm31;
     use rand::{Rng, SeedableRng};
     use rand_chacha::ChaCha20Rng;
+    use rust_bitcoin_m31::qm31_equalverify;
     use stwo_prover::core::air::mask::shifted_mask_points;
     use stwo_prover::core::circle::{CirclePoint, SECURE_FIELD_CIRCLE_ORDER};
+    use stwo_prover::core::fields::qm31::QM31;
     use stwo_prover::core::poly::circle::CanonicCoset;
 
     #[test]
@@ -104,4 +140,66 @@
             assert!(exec_result.success);
         }
     }
+
+    #[test]
+    fn test_shifted_signed_mask_points() {
+        let mut prng = ChaCha20Rng::seed_from_u64(0);
+
+        for _ in 0..100 {
+            let p = CirclePoint::get_point(prng.gen::<u128>() % SECURE_FIELD_CIRCLE_ORDER);
+
+            let mask = vec![vec![0, -1, 2]];
+            let domains = [CanonicCoset::new(5)];
+
+            let expected = shifted_signed_mask_points(&mask, &domains, p);
+            assert_eq!(expected.len(), 1);
+            assert_eq!(expected[0].len(), 3);
+
+            let script = script! {
+                { p }
+                { AirGadget::shifted_signed_mask_points(&mask, &domains) }
+                { expected[0][2] }
+                { CirclePointGadget::equalverify() }
+                { expected[0][1] }
+                { CirclePointGadget::equalverify() }
+                { expected[0][0] }
+                { CirclePointGadget::equalverify() }
+                OP_TRUE
+            };
+
+            let exec_result = execute_script(script);
+            assert!(exec_result.success);
+        }
+    }
+
+    #[test]
+    fn test_eval_from_partial_evals() {
+        let mut prng = ChaCha20Rng::seed_from_u64(0);
+
+        for _ in 0..100 {
+            let a = get_rand_qm31(&mut prng);
+            let b = get_rand_qm31(&mut prng);
+            let c = get_rand_qm31(&mut prng);
+            let d = get_rand_qm31(&mut prng);
+
+            let mut res = a;
+            res += b * QM31::from_u32_unchecked(0, 1, 0, 0);
+            res += c * QM31::from_u32_unchecked(0, 0, 1, 0);
+            res += d * QM31::from_u32_unchecked(0, 0, 0, 1);
+
+            let script = script! {
+                { a }
+                { b }
+                { c }
+                { d }
+                { AirGadget::eval_from_partial_evals() }
+                { res }
+                qm31_equalverify
+                OP_TRUE
+            };
+
+            let exec_result = execute_script(script);
+            assert!(exec_result.success);
+        }
+    }
 }

src/air/mod.rs

@@ -2,7 +2,10 @@ mod bitcoin_script;
 
 use crate::treepp::pushable::{Builder, Pushable};
 pub use bitcoin_script::*;
+use stwo_prover::core::circle::CirclePoint;
 use stwo_prover::core::fields::qm31::SecureField;
+use stwo_prover::core::poly::circle::CanonicCoset;
+use stwo_prover::core::ColumnVec;
 
 #[derive(Clone)]
 /// Hint for the two eval quotient results involved in the composition polynomial.
@@ -20,3 +23,21 @@ impl Pushable for CompositionHint {
         builder
     }
 }
+
+/// A helper function to shift points, but with signed masks.
+pub fn shifted_signed_mask_points(
+    mask: &ColumnVec<Vec<isize>>,
+    domains: &[CanonicCoset],
+    point: CirclePoint<SecureField>,
+) -> ColumnVec<Vec<CirclePoint<SecureField>>> {
+    mask.iter()
+        .zip(domains.iter())
+        .map(|(mask_entry, domain)| {
+            let trace_step = CanonicCoset::new(domain.log_size()).step();
+            mask_entry
+                .iter()
+                .map(|mask_item| point + trace_step.mul_signed(*mask_item).into_ef())
+                .collect()
+        })
+        .collect()
+}

src/channel/bitcoin_script.rs

@@ -78,15 +78,27 @@ impl Sha256ChannelGadget {
     /// Draw queries from the channel, each of logn bits, using hints.
     ///
     /// Output:
-    ///    channel digest
     ///    all the numbers (m)
+    ///    channel digest
     pub fn draw_numbers_with_hint(m: usize, logn: usize) -> Script {
         script! {
-            OP_DUP hash OP_SWAP
-            OP_PUSHBYTES_1 OP_PUSHBYTES_0 OP_CAT hash
-            { Self::unpack_multi_m31(m) }
-            for i in 0..m {
-                { i } OP_ROLL { trim_m31_gadget(logn) }
+            for _ in 0..(m / 8) {
+                OP_DUP hash OP_SWAP
+                OP_PUSHBYTES_1 OP_PUSHBYTES_0 OP_CAT hash
+                { Self::unpack_multi_m31(8) }
+                for i in 0..8 {
+                    { i } OP_ROLL { trim_m31_gadget(logn) }
+                }
+                8 OP_ROLL
+            }
+            if m % 8 != 0 {
+                OP_DUP hash OP_SWAP
+                OP_PUSHBYTES_1 OP_PUSHBYTES_0 OP_CAT hash
+                { Self::unpack_multi_m31(m % 8) }
+                for i in 0..m % 8 {
+                    { i } OP_ROLL { trim_m31_gadget(logn) }
+                }
+                { m % 8 } OP_ROLL
             }
         }
     }
@@ -335,33 +347,64 @@ mod test {
     }
 
     #[test]
-    fn test_draw_5numbers_with_hint() {
+    fn test_draw_many_numbers_with_hint() {
         let mut prng = ChaCha20Rng::seed_from_u64(0);
 
-        let channel_script = Sha256ChannelGadget::draw_numbers_with_hint(5, 15);
+        let channel_script = Sha256ChannelGadget::draw_numbers_with_hint(8, 10);
+        report_bitcoin_script_size("Channel", "draw_8numbers_with_hint", channel_script.len());
+
+        for _ in 0..10 {
+            let mut a = [0u8; 32];
+            a.iter_mut().for_each(|v| *v = prng.gen());
+            let a = Sha256Hash::from(a.to_vec());
 
-        report_bitcoin_script_size("Channel", "draw_5numbers_with_hint", channel_script.len());
+            let mut channel = Sha256Channel::default();
+            channel.update_digest(a);
+            let (b, hint) = channel.draw_queries_and_hints(8, 10);
 
-        for _ in 0..100 {
+            let c = channel.digest;
+
+            let script = script! {
+                { hint }
+                { a }
+                { channel_script.clone() }
+                OP_TOALTSTACK
+                for i in 0..8 {
+                    { b[7 - i] }
+                    OP_EQUALVERIFY
+                }
+                OP_FROMALTSTACK
+                { c }
+                OP_EQUAL
+            };
+            let exec_result = execute_script(script);
+            assert!(exec_result.success);
+        }
+
+        let channel_script = Sha256ChannelGadget::draw_numbers_with_hint(12, 10);
+        report_bitcoin_script_size("Channel", "draw_12numbers_with_hint", channel_script.len());
+
+        for _ in 0..10 {
             let mut a = [0u8; 32];
             a.iter_mut().for_each(|v| *v = prng.gen());
             let a = Sha256Hash::from(a.to_vec());
 
             let mut channel = Sha256Channel::default();
             channel.update_digest(a);
-            let (b, hint) = channel.draw_queries_and_hints(5, 15);
+            let (b, hint) = channel.draw_queries_and_hints(12, 10);
 
             let c = channel.digest;
 
             let script = script! {
                 { hint }
                 { a }
                 { channel_script.clone() }
-                { b[4] } OP_EQUALVERIFY
-                { b[3] } OP_EQUALVERIFY
-                { b[2] } OP_EQUALVERIFY
-                { b[1] } OP_EQUALVERIFY
-                { b[0] } OP_EQUALVERIFY
+                OP_TOALTSTACK
+                for i in 0..12 {
+                    { b[11 - i] }
+                    OP_EQUALVERIFY
+                }
+                OP_FROMALTSTACK
                 { c }
                 OP_EQUAL
             };