refactor sumcheck (#112)

* refactor sumcheck

* fix ci

* clean up

Author: zhenfeizhang

.github/workflows/ci.yml

@@ -85,4 +85,4 @@ jobs:
         run: |
           python3 ./scripts/install.py
           cargo run --bin=dev-setup --release
-          cargo run --release -- -t ${{ matrix.os == 'macos-latest' && 2 || 16 }} -f ${{ matrix.field }}
+          cargo run --bin=gkr --release -- -t ${{ matrix.os == 'macos-latest' && 2 || 16 }} -f ${{ matrix.field }}

Cargo.toml

@@ -1,65 +1,19 @@
-[package]
-name = "expander-rs"
-version = "0.1.0"
-edition = "2021"
-default-run = "expander-rs" # default
-
-[dependencies]
-arith = { path = "arith" }
-gf2 = { path = "arith/gf2" }
-gf2_128 = { path = "arith/gf2_128" }
-mersenne31 = { path = "arith/mersenne31" }
-transcript = { path = "transcript" }
-
-ark-std.workspace = true
-clap.workspace = true
-env_logger.workspace = true
-log.workspace = true
-mpi.workspace = true
-rand.workspace = true
-sha2.workspace = true
-halo2curves.workspace = true
-thiserror.workspace = true
-ethnum.workspace = true
-
-# for the server
-bytes.workspace = true
-tokio.workspace = true
-warp.workspace = true
-chrono.workspace = true
-tynm.workspace = true
-tiny-keccak.workspace = true
-
-[dev-dependencies]
-criterion = "0.5.1"
-
-[[bin]]
-name = "expander-rs-mpi"
-path = "src/main_mpi.rs"
-
-[[bin]]
-name = "expander-exec"
-path = "src/exec.rs"
-
-[[bin]]
-name = "dev-setup"
-path = "src/utils.rs"
-
-[features]
-default = []
-# default = [ "grinding" ]
-grinding = []
-
 [workspace]
 members = [
     "arith", # definitions of Field, Extensions, SIMDS
     "arith/gf2",    
     "arith/gf2_128",
     "arith/goldilocks",
     "arith/mersenne31",
-    "bi-kzg",
+    "bi-kzg", 
+    "circuit", 
+    "config", 
+    "gkr", 
+    "sumcheck",
     "transcript"
 ]
+resolver = "2"
+
 
 [workspace.dependencies]
 ark-std = "0.4"
@@ -86,28 +40,3 @@ tynm = { version = "0.1.6", default-features = false }
 warp = "0.3.7"
 thiserror = "1.0.63"
 ethnum = "1.5.0"
-
-[[bench]]
-name = "bn254_fr"
-harness = false
-path = "arith/benches/fr.rs"
-
-[[bench]]
-name = "gf2"
-harness = false
-path = "arith/gf2/benches/gf2.rs"
-
-[[bench]]
-name = "gf2_128"
-harness = false
-path = "arith/gf2_128/benches/gf2_128.rs"
-
-[[bench]]
-name = "m31"
-harness = false
-path = "arith/mersenne31/benches/m31.rs"
-
-[[bench]]
-name = "gkr-hashes"
-harness = false
-path = "benches/gkr_hashes.rs"

bi-kzg/Cargo.toml

@@ -26,6 +26,6 @@ harness = false
 
 [features]
 default = [ 
-    "ark-std/print-trace",
+    # "ark-std/print-trace",
     "ark-ec/parallel",
     ]

circuit/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "circuit"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+arith = { path = "../arith" }
+config = { path = "../config" }
+transcript = { path = "../transcript" }
+
+ark-std.workspace = true
+ethnum.workspace = true
+log.workspace = true
+thiserror.workspace = true

circuit/src/ecc_circuit.rs

@@ -0,0 +1,145 @@
+use config::GKRConfig;
+use std::{cmp::max, collections::HashMap, fs, io::Cursor};
+
+use crate::*;
+
+// recursive format used in compiler
+pub type SegmentId = usize;
+
+pub struct Allocation {
+    pub i_offset: usize,
+    pub o_offset: usize,
+}
+
+#[derive(Default)]
+pub struct Segment<C: GKRConfig> {
+    pub i_var_num: usize,
+    pub o_var_num: usize,
+    pub child_segs: Vec<(SegmentId, Vec<Allocation>)>,
+    pub gate_muls: Vec<GateMul<C>>,
+    pub gate_adds: Vec<GateAdd<C>>,
+    pub gate_consts: Vec<GateConst<C>>,
+    pub gate_uni: Vec<GateUni<C>>,
+}
+
+impl<C: GKRConfig> Segment<C> {
+    pub fn contain_gates(&self) -> bool {
+        !self.gate_muls.is_empty()
+            || !self.gate_adds.is_empty()
+            || !self.gate_consts.is_empty()
+            || !self.gate_uni.is_empty()
+    }
+
+    pub fn scan_leaf_segments(
+        &self,
+        rc: &RecursiveCircuit<C>,
+        cur_id: SegmentId,
+    ) -> HashMap<SegmentId, Vec<Allocation>> {
+        let mut ret = HashMap::new();
+        if self.contain_gates() {
+            ret.insert(
+                cur_id,
+                vec![Allocation {
+                    i_offset: 0,
+                    o_offset: 0,
+                }],
+            );
+        }
+        for (child_seg_id, child_allocs) in &self.child_segs {
+            let leaves = rc.segments[*child_seg_id].scan_leaf_segments(rc, *child_seg_id);
+            for (leaf_seg_id, leaf_allocs) in leaves {
+                ret.entry(leaf_seg_id).or_insert_with(Vec::new);
+                for child_alloc in child_allocs {
+                    for leaf_alloc in &leaf_allocs {
+                        ret.get_mut(&leaf_seg_id).unwrap().push(Allocation {
+                            i_offset: child_alloc.i_offset + leaf_alloc.i_offset,
+                            o_offset: child_alloc.o_offset + leaf_alloc.o_offset,
+                        });
+                    }
+                }
+            }
+        }
+        ret
+    }
+}
+
+#[derive(Default)]
+pub struct RecursiveCircuit<C: GKRConfig> {
+    pub num_public_inputs: usize,
+    pub num_outputs: usize,
+    pub expected_num_output_zeros: usize,
+
+    pub segments: Vec<Segment<C>>,
+    pub layers: Vec<SegmentId>,
+}
+
+impl<C: GKRConfig> RecursiveCircuit<C> {
+    pub fn load(filename: &str) -> std::result::Result<Self, CircuitError> {
+        let file_bytes = fs::read(filename)?;
+        let cursor = Cursor::new(file_bytes);
+
+        Ok(Self::deserialize_from(cursor))
+    }
+
+    pub fn flatten(&self) -> Circuit<C> {
+        let mut ret = Circuit::<C> {
+            expected_num_output_zeros: self.expected_num_output_zeros,
+            ..Default::default()
+        };
+        // layer-by-layer conversion
+        for layer_id in &self.layers {
+            let layer_seg = &self.segments[*layer_id];
+            let leaves = layer_seg.scan_leaf_segments(self, *layer_id);
+            let mut ret_layer = CircuitLayer {
+                input_var_num: max(layer_seg.i_var_num, 1), // var_num >= 1
+                output_var_num: max(layer_seg.o_var_num, 1), // var_num >= 1
+                ..Default::default()
+            };
+            for (leaf_seg_id, leaf_allocs) in leaves {
+                let leaf_seg = &self.segments[leaf_seg_id];
+                for alloc in leaf_allocs {
+                    for gate in &leaf_seg.gate_muls {
+                        let mut gate = gate.clone();
+                        gate.i_ids[0] += alloc.i_offset;
+                        gate.i_ids[1] += alloc.i_offset;
+                        gate.o_id += alloc.o_offset;
+                        ret_layer.mul.push(gate);
+                    }
+                    for gate in &leaf_seg.gate_adds {
+                        let mut gate = gate.clone();
+                        gate.i_ids[0] += alloc.i_offset;
+                        gate.o_id += alloc.o_offset;
+                        ret_layer.add.push(gate);
+                    }
+                    for gate in &leaf_seg.gate_consts {
+                        let mut gate = gate.clone();
+                        gate.o_id += alloc.o_offset;
+                        ret_layer.const_.push(gate);
+                    }
+                    for gate in &leaf_seg.gate_uni {
+                        let mut gate = gate.clone();
+                        gate.i_ids[0] += alloc.i_offset;
+                        gate.o_id += alloc.o_offset;
+                        ret_layer.uni.push(gate);
+                    }
+                }
+            }
+            // debug print layer
+            log::trace!(
+                "layer {} mul: {} add: {} const:{} uni:{} i_var_num: {} o_var_num: {}",
+                ret.layers.len(),
+                ret_layer.mul.len(),
+                ret_layer.add.len(),
+                ret_layer.const_.len(),
+                ret_layer.uni.len(),
+                ret_layer.input_var_num,
+                ret_layer.output_var_num,
+            );
+            ret.layers.push(ret_layer);
+        }
+
+        ret.identify_rnd_coefs();
+        ret.identify_structure_info();
+        ret
+    }
+}