chore: optimize memory usage (#176)

* chore: drop memory

* chore: opt interpolate_prepare and bit_reverse

* chore: opt interpolate_prepare and bit_reverse

* chore: fix mem

* fix: number may consists of hex and negtives

* chore: opt calculate h1h2

* chore: opt calculate h1h2

* chore: opt fft

* fix: str to digit

* fix: str to digit

* chore: recover UT

* chore: code polish & fix rayon hashmap

* chore: code polish & fix rayon hashmap

* chore: drop memory

* fix: handle powdr all constant case

* feat: seralize setup info

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* chore: dev

* doc: add todo

* doc: add todo

* doc: add todo

* fix: rebase

* chore: stash

* fix: ut

* chore: stash code

* chore: stash code

algebraic/src/reader.rs

@@ -91,7 +91,6 @@ pub fn load_witness_from_bin_reader<E: ScalarEngine, R: Read>(mut reader: R) ->
     reader.read_exact(&mut wtns_header)?;
     if wtns_header != [119, 116, 110, 115] {
         // python -c 'print([ord(c) for c in "wtns"])' => [119, 116, 110, 115]
-        // bail!("Invalid file header".to_string()));
         bail!("Invalid file header");
     }
     let version = reader.read_u32::<LittleEndian>()?;

starky/Cargo.toml

@@ -24,6 +24,7 @@ lazy_static = "1.0"
 ## threading
 rayon = { version = "1.5"}
 num_cpus = "1.0"
+hashbrown = { version = "0.14.3", features = ["rayon"] }
 
 # error and log
 thiserror="1.0"

starky/benches/batch_inverse.rs

@@ -4,7 +4,8 @@ extern crate criterion;
 use criterion::{BenchmarkId, Criterion};
 use starky::dev::gen_rand_goldfields;
 use starky::f3g::F3G;
-use starky::traits::{batch_inverse, FieldExtension};
+use starky::polutils::batch_inverse;
+use starky::traits::FieldExtension;
 
 const MIN_K: usize = 6;
 const MAX_K: usize = 24;

starky/src/digest.rs

@@ -1,11 +1,14 @@
 #![allow(non_snake_case)]
+use crate::errors::Result;
 use crate::field_bls12381::Fr as Fr_bls12381;
 use crate::field_bls12381::FrRepr as FrRepr_bls12381;
 use crate::field_bn128::{Fr, FrRepr};
 use crate::traits::MTNodeType;
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use ff::*;
 use plonky::field_gl::Fr as FGL;
 use std::fmt::Display;
+use std::io::{Read, Write};
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
@@ -46,6 +49,21 @@ impl<const N: usize> MTNodeType for ElementDigest<N> {
         }
         y
     }
+
+    fn save<W: Write>(&self, writer: &mut W) -> Result<()> {
+        for i in &self.0 {
+            writer.write_u64::<LittleEndian>(i.as_int())?;
+        }
+        Ok(())
+    }
+
+    fn load<R: Read>(reader: &mut R) -> Result<Self> {
+        let e1 = reader.read_u64::<LittleEndian>()?;
+        let e2 = reader.read_u64::<LittleEndian>()?;
+        let e3 = reader.read_u64::<LittleEndian>()?;
+        let e4 = reader.read_u64::<LittleEndian>()?;
+        Ok(Self::new(&[e1.into(), e2.into(), e3.into(), e4.into()]))
+    }
 }
 
 impl<const N: usize> Display for ElementDigest<N> {

starky/src/f3g.rs

@@ -629,7 +629,8 @@ impl Display for F3G {
 #[cfg(test)]
 pub mod tests {
     use crate::f3g::F3G;
-    use crate::traits::{batch_inverse, FieldExtension};
+    use crate::polutils::batch_inverse;
+    use crate::traits::FieldExtension;
     use plonky::field_gl::Fr;
     use plonky::Field;
     use std::ops::{Add, Mul};

starky/src/f5g.rs

@@ -725,7 +725,8 @@ impl F5G {
 #[cfg(test)]
 pub mod tests {
     use crate::f5g::F5G;
-    use crate::traits::{batch_inverse, FieldExtension};
+    use crate::polutils::batch_inverse;
+    use crate::traits::FieldExtension;
     use plonky::field_gl::Fr;
     use plonky::Field;
     use std::ops::{Add, Mul};

starky/src/fft.rs

@@ -3,6 +3,7 @@
 use crate::constant::MG;
 use crate::helper::log2_any;
 use crate::traits::FieldExtension;
+use rayon::prelude::*;
 
 #[allow(clippy::upper_case_acronyms)]
 #[derive(Default)]
@@ -77,9 +78,12 @@ impl<F: FieldExtension> FFT<F> {
         let n = p.len();
         let n2inv = F::from(p.len()).inv();
         let mut res = vec![F::ZERO; q.len()];
-        for i in 0..n {
-            res[(n - i) % n] = q[i] * n2inv;
-        }
+
+        res[0] = q[0] * n2inv;
+        res[1..]
+            .par_iter_mut()
+            .enumerate()
+            .for_each(|(i, out)| *out = q[n - i - 1] * n2inv);
         res
     }
 }

starky/src/fft_p.rs

@@ -97,11 +97,14 @@ pub fn bit_reverse<F: FieldExtension>(
 
     let len = n * n_pols;
     assert_eq!(len, buffdst.len());
-    for j in 0..len {
-        let i = j / n_pols;
-        let k = j % n_pols;
-        buffdst[j] = buffsrc[ris[i] * n_pols + k];
-    }
+    buffdst[0..len]
+        .par_iter_mut()
+        .enumerate()
+        .for_each(|(j, out)| {
+            let i = j / n_pols;
+            let k = j % n_pols;
+            *out = buffsrc[ris[i] * n_pols + k];
+        });
 }
 
 pub fn interpolate_bit_reverse<F: FieldExtension>(
@@ -113,12 +116,15 @@ pub fn interpolate_bit_reverse<F: FieldExtension>(
     let n = 1 << nbits;
     let ris = BRs(0, n, nbits); // move it outside the loop. obtain it from cache.
 
-    for i in 0..n {
-        let rii = (n - ris[i]) % n;
-        for k in 0..n_pols {
-            buffdst[i * n_pols + k] = buffsrc[rii * n_pols + k];
-        }
-    }
+    buffdst[0..n * n_pols]
+        .par_chunks_mut(n_pols)
+        .enumerate()
+        .for_each(|(i, out)| {
+            let rii = (n - ris[i]) % n;
+            for k in 0..n_pols {
+                out[k] = buffsrc[rii * n_pols + k];
+            }
+        });
 }
 
 pub fn inv_bit_reverse<F: FieldExtension>(
@@ -133,12 +139,15 @@ pub fn inv_bit_reverse<F: FieldExtension>(
 
     let len = n * n_pols;
     assert_eq!(len, buffdst.len());
-    for j in 0..len {
-        let i = j / n_pols;
-        let k = j % n_pols;
-        let rii = (n - ris[i]) % n;
-        buffdst[j] = buffsrc[rii * n_pols + k] * n_inv;
-    }
+    buffdst[0..len]
+        .par_iter_mut()
+        .enumerate()
+        .for_each(|(j, out)| {
+            let i = j / n_pols;
+            let k = j % n_pols;
+            let rii = (n - ris[i]) % n;
+            *out = buffsrc[rii * n_pols + k] * n_inv;
+        });
 }
 
 pub fn interpolate_prepare<F: FieldExtension>(buff: &mut Vec<F>, n_pols: usize, nbits: usize) {

starky/src/fri.rs

@@ -58,7 +58,6 @@ impl FRI {
         let mut pol = pol.to_owned();
         let mut standard_fft = FFT::new();
         let mut pol_bits = log2_any(pol.len());
-        log::trace!("fri prove {} {}", pol.len(), 1 << pol_bits);
         assert_eq!(1 << pol_bits, pol.len());
         assert_eq!(pol_bits, self.in_nbits);
 
@@ -93,7 +92,6 @@ impl FRI {
                     sinv *= wi;
                 }
             }
-            log::trace!("pol2_e 0={}, 1={}", pol2_e[0], pol2_e[1]);
             if si < self.steps.len() - 1 {
                 let n_groups = 1 << self.steps[si + 1].nBits;
                 let group_size = (1 << stepi.nBits) / n_groups;
@@ -185,7 +183,6 @@ impl FRI {
         let n_queries = self.n_queries;
         let mut ys = transcript.get_permutations(self.n_queries, self.steps[0].nBits)?;
         let mut pol_bits = self.in_nbits;
-        log::trace!("ys: {:?}, pol_bits {}", ys, self.in_nbits);
         let mut shift = F::from(*SHIFT);
 
         let check_query_fn = |si: usize,

starky/src/interpreter.rs

@@ -4,6 +4,7 @@ use crate::starkinfo::StarkInfo;
 use crate::starkinfo_codegen::Node;
 use crate::starkinfo_codegen::Section;
 use crate::traits::FieldExtension;
+use crate::types::parse_pil_number;
 use std::fmt;
 
 #[derive(Clone, Debug)]
@@ -471,12 +472,10 @@ fn get_ref<F: FieldExtension>(
                 panic!("Invalid dom");
             }
         }
-        "number" => Expr::new(
-            Ops::Vari(F::from(r.value.clone().unwrap().parse::<u64>().unwrap())),
-            vec![],
-            vec![],
-            vec![],
-        ),
+        "number" => {
+            let n_val = parse_pil_number(r.value.as_ref().unwrap());
+            Expr::new(Ops::Vari(F::from(n_val)), vec![], vec![], vec![])
+        }
         "public" => Expr::new(
             Ops::Refer,
             vec!["publics".to_string()],

starky/src/lib.rs

@@ -5,7 +5,7 @@
 #![cfg_attr(feature = "avx512", feature(stdsimd))]
 
 pub mod polsarray;
-mod polutils;
+pub mod polutils;
 pub mod stark_verifier_circom;
 pub mod stark_verifier_circom_bn128;
 pub mod traits;

starky/src/merklehash.rs

@@ -28,8 +28,10 @@ use crate::poseidon_opt::Poseidon;
 use crate::traits::MTNodeType;
 use crate::traits::MerkleTree;
 use anyhow::{bail, Result};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use plonky::field_gl::Fr as FGL;
 use rayon::prelude::*;
+use std::io::{Read, Write};
 use std::time::Instant;
 
 #[derive(Default)]
@@ -262,6 +264,42 @@ impl MerkleTree for MerkleTreeGL {
         }
     }
 
+    // TODO: https://github.com/0xEigenLabs/eigen-zkvm/issues/187
+    fn save<W: Write>(&self, writer: &mut W) -> Result<()> {
+        writer.write_u64::<LittleEndian>(self.width as u64)?;
+        writer.write_u64::<LittleEndian>(self.height as u64)?;
+        writer.write_u64::<LittleEndian>(self.elements.len() as u64)?;
+        for i in &self.elements {
+            writer.write_u64::<LittleEndian>(i.as_int())?;
+        }
+        writer.write_u64::<LittleEndian>(self.nodes.len() as u64)?;
+        for i in &self.nodes {
+            i.save(writer)?;
+        }
+        Ok(())
+    }
+
+    fn load<R: Read>(reader: &mut R) -> Result<Self> {
+        let mut mt = Self::new();
+        mt.width = reader.read_u64::<LittleEndian>()? as usize;
+        mt.height = reader.read_u64::<LittleEndian>()? as usize;
+
+        let es = reader.read_u64::<LittleEndian>()? as usize;
+        mt.elements = vec![FGL::ZERO; es];
+        for i in 0..es {
+            let e = reader.read_u64::<LittleEndian>()?;
+            mt.elements[i] = FGL::from(e);
+        }
+
+        let ns = reader.read_u64::<LittleEndian>()? as usize;
+        mt.nodes = vec![ElementDigest::<4>::new(&[FGL::ZERO, FGL::ZERO, FGL::ZERO, FGL::ZERO]); ns];
+        for i in 0..ns {
+            mt.nodes[i] = ElementDigest::<4>::load(reader)?;
+        }
+
+        Ok(mt)
+    }
+
     fn element_size(&self) -> usize {
         self.elements.len()
     }

starky/src/merklehash_bls12381.rs

@@ -8,9 +8,11 @@ use crate::poseidon_bls12381_opt::Poseidon;
 use crate::traits::MTNodeType;
 use crate::traits::MerkleTree;
 use anyhow::{bail, Result};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use ff::Field;
 use plonky::field_gl::Fr as FGL;
 use rayon::prelude::*;
+use std::io::{Read, Write};
 use std::time::Instant;
 
 #[derive(Default)]
@@ -156,6 +158,40 @@ impl MerkleTree for MerkleTreeBLS12381 {
             poseidon: Poseidon::new(),
         }
     }
+    fn save<W: Write>(&self, writer: &mut W) -> Result<()> {
+        writer.write_u64::<LittleEndian>(self.width as u64)?;
+        writer.write_u64::<LittleEndian>(self.height as u64)?;
+        writer.write_u64::<LittleEndian>(self.elements.len() as u64)?;
+        for i in &self.elements {
+            writer.write_u64::<LittleEndian>(i.as_int())?;
+        }
+        writer.write_u64::<LittleEndian>(self.nodes.len() as u64)?;
+        for i in &self.nodes {
+            i.save(writer)?;
+        }
+        Ok(())
+    }
+
+    fn load<R: Read>(reader: &mut R) -> Result<Self> {
+        let mut mt = Self::new();
+        mt.width = reader.read_u64::<LittleEndian>()? as usize;
+        mt.height = reader.read_u64::<LittleEndian>()? as usize;
+
+        let es = reader.read_u64::<LittleEndian>()? as usize;
+        mt.elements = vec![FGL::ZERO; es];
+        for i in 0..es {
+            let e = reader.read_u64::<LittleEndian>()?;
+            mt.elements[i] = FGL::from(e);
+        }
+
+        let ns = reader.read_u64::<LittleEndian>()? as usize;
+        mt.nodes = vec![ElementDigest::<4>::new(&[FGL::ZERO, FGL::ZERO, FGL::ZERO, FGL::ZERO]); ns];
+        for i in 0..ns {
+            mt.nodes[i] = ElementDigest::<4>::load(reader)?;
+        }
+
+        Ok(mt)
+    }
 
     fn element_size(&self) -> usize {
         self.elements.len()