Merge pull request #43 from orochi-network/feature/kzg_setup_and_comm…

…it_trace_record

Add setup + commit + misc methods

zkmemory/examples/256bits-machine.rs

@@ -1,14 +1,14 @@
 use rbtree::RBTree;
-use std::{
-    marker::PhantomData,
-    println,
-};
+use std::{marker::PhantomData, println};
 use zkmemory::{
-    impl_register_machine, impl_stack_machine, impl_state_machine,
-    machine::{AbstractContext, AbstractInstruction, AbstractMachine, Register, CellInteraction, TraceRecord},
-    config::{AllocatedSection, Config, ConfigArgs, DefaultConfig},
     base::{Base, B256},
-    error::Error, kzg::{KZGMemoryCommitment, KZGParams},
+    config::{AllocatedSection, Config, ConfigArgs, DefaultConfig},
+    error::Error,
+    impl_register_machine, impl_stack_machine, impl_state_machine,
+    machine::{
+        AbstractContext, AbstractInstruction, AbstractMachine, CellInteraction, Register,
+        TraceRecord,
+    },
 };
 
 /// My instruction set for the machine
@@ -132,18 +132,17 @@ where
             MyInstruction::Read(addr) => {
                 if !machine.memory_allocated.contain(*addr) {
                     panic!("{}", Error::MemoryAccessDeinied);
-                }
-                else {
+                } else {
                     machine.read(*addr).expect("Unable to read to memory");
                 }
             }
             MyInstruction::Write(addr, val) => {
                 if !machine.memory_allocated.contain(*addr) {
                     panic!("{}", Error::MemoryAccessDeinied);
-                }
-                else {
-                    machine.write(*addr, *val)
-                    .expect("Unable to write to memory");
+                } else {
+                    machine
+                        .write(*addr, *val)
+                        .expect("Unable to write to memory");
                 }
             }
             MyInstruction::Push(value) => {
@@ -203,7 +202,7 @@ where
                             }
                             _ => panic!("Register unable to be two cells"),
                         }
-                    },
+                    }
                     _ => panic!("Register unable to be two cells"),
                 }
             }
@@ -246,12 +245,21 @@ where
 
     /// Show address maps of memory, stack and registers sections
     pub fn show_sections_maps(&self) -> () {
-        println!("Memory section map: from {} to {}", 
-        self.memory_allocated.low(), self.memory_allocated.high());
-        println!("Register section map: from {} to {}", 
-        self.register_allocated.low(), self.register_allocated.high());
-        println!("Stack section map: from {} to {}", 
-        self.stack_allocated.low(), self.stack_allocated.high());
+        println!(
+            "Memory section map: from {} to {}",
+            self.memory_allocated.low(),
+            self.memory_allocated.high()
+        );
+        println!(
+            "Register section map: from {} to {}",
+            self.register_allocated.low(),
+            self.register_allocated.high()
+        );
+        println!(
+            "Stack section map: from {} to {}",
+            self.stack_allocated.low(),
+            self.stack_allocated.high()
+        );
     }
 }
 
@@ -315,7 +323,6 @@ impl_stack_machine!(StateMachine);
 impl_state_machine!(StateMachine);
 
 fn main() {
-
     // Define the desired machine configuration
     let mut machine = StateMachine::<B256, B256, 32, 32>::new(DefaultConfig::default());
 
@@ -347,5 +354,5 @@ fn main() {
     // Print the trace record (prettified), sorted by ascending address by default
     for x in machine.trace().into_iter() {
         println!("{:?}", x);
-    }  
+    }
 }

zkmemory/src/commitment/kzg.rs

@@ -1,31 +1,23 @@
-extern crate alloc;
 extern crate std;
-use alloc::vec::Vec;
 use core::marker::PhantomData;
-use ff::{Field, PrimeField, WithSmallOrderMulGroup};
+use ff::{Field, PrimeField};
 use group::Curve;
 use rand_core::OsRng;
-use rbtree::RBTree;
 
-use crate::{
-    base::Base,
-    machine::TraceRecord,
-    machine::MemoryInstruction,
-};
+use crate::{base::Base, machine::MemoryInstruction, machine::TraceRecord};
 use halo2_proofs::{
-    arithmetic::{eval_polynomial, lagrange_interpolate},
+    arithmetic::lagrange_interpolate,
     poly::{
-        commitment::{Blind, CommitmentScheme, ParamsProver, Prover},
+        commitment::{Blind, ParamsProver},
         kzg::commitment::ParamsKZG,
-        {Coeff, EvaluationDomain, Polynomial, ProverQuery},
+        {Coeff, EvaluationDomain, Polynomial},
     },
-    transcript::{EncodedChallenge, TranscriptWriterBuffer},
 };
 use halo2curves::bn256::{Bn256, Fr, G1Affine};
 
 /// A KZG module that commit to the memory trace through the execution trace
 #[derive(Debug, Clone)]
-pub struct KZGParams<K, V, const S: usize, const T: usize>
+pub struct KZGMemoryCommitment<K, V, const S: usize, const T: usize>
 where
     K: Base<S>,
     V: Base<T>,
@@ -34,76 +26,41 @@ where
     kzg_params: ParamsKZG<Bn256>,
     // Domain used for creating polynomials
     domain: EvaluationDomain<Fr>,
-    // A copy of an execution trace, RBTree for later implementation of sorting
-    trace_record: RBTree<TraceRecord<K, V, S, T>, PhantomData<()>>,
-}
-
-/// KZG trait for committing the memory trace elements
-pub trait KZGMemoryCommitment<K, V, const S: usize, const T: usize>
-where
-    K: Base<S>,
-    V: Base<T>,
-{
-    /// Init a new KZG module
-    fn init() -> KZGParams<K, V, S, T>;
-
-    /// Commit the trace element tuple to G1
-    fn commit_trace_element(&mut self, trace: TraceRecord<K, V, S, T>) -> G1Affine;
-
-    /// Convert trace tuple elements to field elements
-    /// Also, fill with ZERO elements to 8 (the nearest power of 2)
-    fn trace_to_field(&self, trace: TraceRecord<K, V, S, T>) -> [Fr; 8];
-
-    /// Use Lagrange interpolation to form the polynomial
-    /// representing a trace tuple element
-    /// We will use the points as successive powers of the field's primitive root
-    fn get_trace_poly(&self, evals: [Fr; 8]) -> Polynomial<Fr, Coeff>;
-
-    /// Convert raw 32 bytes from big endian to Fr element
-    fn be_bytes_to_field(&self, bytes: &mut [u8]) -> Fr;
-
-    /// Create witness
-    fn create_proof_sh_plonk<
-        'params,
-        Scheme: CommitmentScheme,
-        P: Prover<'params, Scheme>,
-        E: EncodedChallenge<Scheme::Curve>,
-        TW: TranscriptWriterBuffer<Vec<u8>, Scheme::Curve, E>,
-    >(
-        params: &'params Scheme::ParamsProver,
-        // a list of point x_1,x_2,...x_n
-        points_list: Vec<Scheme::Scalar>,
-        // a list of polynomials p_1(x), p_2(x),...,p_n(x)
-        polynomial_list: Vec<Polynomial<Scheme::Scalar, Coeff>>,
-    ) -> Vec<u8>
-    where
-        Scheme::Scalar: WithSmallOrderMulGroup<3>;
+    // Phantom data for K and V
+    _marker_k: PhantomData<K>,
+    _marker_v: PhantomData<V>,
 }
 
-impl<K, V, const S: usize, const T: usize> KZGMemoryCommitment<K, V, S, T> for KZGParams<K, V, S, T>
+impl<K, V, const S: usize, const T: usize> KZGMemoryCommitment<K, V, S, T>
 where
     K: Base<S>,
     V: Base<T>,
 {
-    fn init() -> Self {
+    /// Initialize KZG parameters
+    /// K = 3 since we need the poly degree to be 2^3 = 8
+    pub fn new() -> Self {
         const K: u32 = 3;
         Self {
             kzg_params: ParamsKZG::<Bn256>::new(K),
             domain: EvaluationDomain::new(1, K),
-            trace_record: RBTree::new(),
+            _marker_k: PhantomData::<K>,
+            _marker_v: PhantomData::<V>,
         }
     }
 
-    fn commit_trace_element(&mut self, trace: TraceRecord<K, V, S, T>) -> G1Affine {
-        // Update the trace record set
-        self.trace_record.insert(trace, PhantomData);
-        let field_tuple = self.trace_to_field(trace);
-        let poly = self.get_trace_poly(field_tuple);
+    /// Commit a trace record in an execution trace
+    pub fn commit(&mut self, trace: TraceRecord<K, V, S, T>) -> G1Affine {
+        // Convert trace record into polynomial
+        let poly = self.poly_from_trace(trace);
+
+        // Commit the polynomial using Halo2's code
         let alpha = Blind(Fr::random(OsRng));
         let commitment = self.kzg_params.commit(&poly, alpha);
         commitment.to_affine()
     }
 
+    // Convert a trace record to 8 field elements
+    // The last 3 elements will be ZERO
     fn trace_to_field(&self, trace: TraceRecord<K, V, S, T>) -> [Fr; 8] {
         let (t, s, i, l, d) = trace.get_tuple();
 
@@ -132,77 +89,94 @@ where
         }
     }
 
-    fn get_trace_poly(&self, evals: [Fr; 8]) -> Polynomial<Fr, Coeff> {
+    // Convert the trace record into a polynomial
+    fn poly_from_trace(&self, trace: TraceRecord<K, V, S, T>) -> Polynomial<Fr, Coeff> {
+        let evals = self.trace_to_field(trace);
+        self.poly_from_evals(evals)
+    }
+
+    // Convert 8 field elements of a trace record
+    // into a polynomial
+    fn poly_from_evals(&self, evals: [Fr; 8]) -> Polynomial<Fr, Coeff> {
         let mut points_arr = [Fr::ONE; 8];
         let mut current_point = Fr::ONE;
+
+        // We use successive powers of primitive roots as points
+        // We use elements in trace record to be the evals
+        // 3 last evals should be ZERO
         for i in 1..8 as usize {
             current_point *= Fr::MULTIPLICATIVE_GENERATOR;
             points_arr[i] = current_point;
         }
+
+        // Use Lagrange interpolation
         self.domain.coeff_from_vec(lagrange_interpolate(
             points_arr.as_slice(),
             evals.as_slice(),
         ))
     }
 
+    // Convert 32 bytes number to field elements
+    // This is made compatible with the Fr endianess
     fn be_bytes_to_field(&self, bytes: &mut [u8]) -> Fr {
         bytes.reverse();
         let b = bytes.as_ref();
         let inner = [0, 8, 16, 24].map(|i| u64::from_le_bytes(b[i..i + 8].try_into().unwrap()));
-        Fr::from_raw(inner)
+        let result = Fr::from_raw(inner);
+        result
     }
+}
 
-    fn create_proof_sh_plonk<
-        'params,
-        Scheme: CommitmentScheme,
-        P: Prover<'params, Scheme>,
-        E: EncodedChallenge<Scheme::Curve>,
-        TW: TranscriptWriterBuffer<Vec<u8>, Scheme::Curve, E>,
-    >(
-        params: &'params Scheme::ParamsProver,
-        // a list of point x_1,x_2,...x_n
-        points_list: Vec<Scheme::Scalar>,
-        // a list of polynomials p_1(x), p_2(x),...,p_n(x)
-        polynomial_list: Vec<Polynomial<Scheme::Scalar, Coeff>>,
-    ) -> Vec<u8>
-    where
-        Scheme::Scalar: WithSmallOrderMulGroup<3>,
-    {
-        // this function, given a list of points x_1,x_2,...,x_n
-        // and polynomials p_1(x),p_2(x),...,p_n(x)
-        // create a witness for the value p_1(x_1), p_2(x_2),...,p_n(x_n)
-        let mut transcript = TW::init(Vec::new());
-
-        let blind = Blind::new(&mut OsRng);
-        // Commit the polynomial p_i(x)
-        for i in 0..polynomial_list.len() {
-            let a = params.commit(&polynomial_list[i], blind).to_affine();
-            transcript.write_point(a).unwrap();
-        }
-        // evaluate the values p_i(x_i) for i=1,2,...,n
-        for i in 0..polynomial_list.len() {
-            let av = eval_polynomial(&polynomial_list[i], points_list[i]);
-            transcript.write_scalar(av).unwrap();
-        }
+#[cfg(test)]
+mod test {
 
-        // this query is used to list all the values p_1(x_1), p_2(x_2),...,p_n(x_n)
-        // in the query list of shplonk prover
-
-        let mut queries: Vec<ProverQuery<'_, <Scheme as CommitmentScheme>::Curve>> = Vec::new();
-        for i in 0..polynomial_list.len() {
-            let mut temp = ProverQuery {
-                point: points_list[i],
-                poly: &polynomial_list[i],
-                blind,
-            };
-            queries.push(temp);
+    use super::*;
+    use crate::{base::B256, machine::AbstractTraceRecord};
+    use halo2_proofs::arithmetic::eval_polynomial;
+
+    #[test]
+    fn test_conversion_fr() {
+        let kzg_scheme = KZGMemoryCommitment::<B256, B256, 32, 32>::new();
+
+        // Create a 32-bytes repr of Base 256
+        let mut chunk = [0u8; 32];
+        for i in 0..32 {
+            chunk[i] = i as u8;
         }
-        // create the proof
-        let prover = P::new(params);
-        prover
-            .create_proof(&mut OsRng, &mut transcript, queries)
-            .unwrap();
 
-        transcript.finalize()
+        // Use my method to convert to Fr
+        let fr = kzg_scheme.be_bytes_to_field(chunk.as_mut_slice());
+
+        // Use Fr's method to convert back to bytes
+        let chunk_fr: [u8; 32] = fr.try_into().unwrap();
+
+        assert_eq!(chunk_fr, chunk);
     }
-}
+
+    #[test]
+    fn test_record_polynomial_conversion() {
+        let kzg_scheme = KZGMemoryCommitment::<B256, B256, 32, 32>::new();
+
+        // Initialize a trace record
+        let trace = TraceRecord::<B256, B256, 32, 32>::new(
+            1u64,
+            2u64,
+            MemoryInstruction::Read,
+            B256::zero(),
+            B256::from(100),
+        );
+
+        // Get the polynomial
+        let poly_trace = kzg_scheme.poly_from_trace(trace);
+
+        // Get only the evals, which is the trace record's elements converted to field elements
+        let poly_evals = kzg_scheme.trace_to_field(trace);
+
+        // Test each eval values
+        let mut base_index = Fr::ONE;
+        for i in 0..8 {
+            assert_eq!(eval_polynomial(&poly_trace, base_index), poly_evals[i]);
+            base_index *= Fr::MULTIPLICATIVE_GENERATOR;
+        }
+    }
+}

zkmemory/src/commitment/merkle_tree.rs

@@ -1,2 +1,2 @@
 /// Merkle tree module
-pub struct MerkleCommitmentScheme;
+pub struct MerkleCommitmentScheme;

zkmemory/src/commitment/mod.rs

@@ -1,6 +1,6 @@
 /// KZG commitment scheme
 pub mod kzg;
+/// Commitment scheme using Merkle Tree
+pub mod merkle_tree;
 /// Commitment scheme using Verkle Tree
 pub mod verkle_tree;
-/// Commitment scheme using Merkle Tree
-pub mod merkle_tree;

zkmemory/src/commitment/verkle_tree.rs

@@ -1,2 +1,2 @@
 /// Verkle tree module
-pub struct VerkleCommitmentScheme;
+pub struct VerkleCommitmentScheme;