experience: precheck whether slice meets circuit

DelphinusLab · Apr 12, 2024 · 1558530 · 1558530
1 parent 6f84f5a
commit 1558530
Show file tree

Hide file tree

Showing 14 changed files with 139 additions and 49 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/cli/src/config.rs b/crates/cli/src/config.rs
@@ -260,14 +260,6 @@ impl Config {
             result
         };
 
-        if cfg!(not(feature = "continuation")) {
-            if result.tables.execution_tables.etable.len() != 1 {
-                return Err(anyhow::anyhow!(
-                    "Only support single slice for non-continuation mode.\nYou could increase K or enable continuation feature."
-                ));
-            }
-        }
-
         {
             if let Some(context_output_filename) = context_output_filename {
                 let context_output_path = output_dir.join(context_output_filename);
@@ -318,8 +310,10 @@ impl Config {
 
         let progress_bar = ProgressBar::new(result.tables.execution_tables.etable.len() as u64);
 
-        let mut slices = loader.slice(result).into_iter().enumerate().peekable();
+        let mut slices = loader.slice(result)?.into_iter().enumerate().peekable();
         while let Some((index, circuit)) = slices.next() {
+            let circuit = circuit?;
+
             let _is_finalized_circuit = slices.peek().is_none();
 
             if mock_test {

diff --git a/crates/zkwasm/Cargo.toml b/crates/zkwasm/Cargo.toml
@@ -22,6 +22,7 @@ strum = "0.24.1"
 strum_macros = "0.24.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+thiserror = "1.0.58"
 ff = "0.12"
 sha2 = "0.10.6"
 anyhow.workspace = true

diff --git a/crates/zkwasm/src/circuits/image_table/mod.rs b/crates/zkwasm/src/circuits/image_table/mod.rs
@@ -7,10 +7,10 @@ use halo2_proofs::plonk::Fixed;
 use halo2_proofs::plonk::VirtualCells;
 use std::marker::PhantomData;
 
-use super::zkwasm_circuit::RESERVE_ROWS;
 use super::utils::image_table::GLOBAL_CAPABILITY;
 use super::utils::image_table::INIT_MEMORY_ENTRIES_OFFSET;
 use super::utils::image_table::STACK_CAPABILITY;
+use super::zkwasm_circuit::RESERVE_ROWS;
 
 mod assign;
 mod configure;

diff --git a/crates/zkwasm/src/circuits/mod.rs b/crates/zkwasm/src/circuits/mod.rs
@@ -1,4 +1,5 @@
 use crate::circuits::utils::Context;
+use crate::error::BuildingCircuitError;
 
 use halo2_proofs::arithmetic::FieldExt;
 use halo2_proofs::plonk::ConstraintSystem;
@@ -8,6 +9,10 @@ use num_bigint::BigUint;
 use specs::slice::Slice;
 use std::marker::PhantomData;
 
+use self::etable::EVENT_TABLE_ENTRY_ROWS;
+use self::image_table::compute_maximal_pages;
+use self::zkwasm_circuit::RESERVE_ROWS;
+
 pub(crate) mod cell;
 pub(crate) mod etable;
 
@@ -34,19 +39,55 @@ pub mod zkwasm_circuit;
 pub type CompilationTable = specs::CompilationTable;
 pub type ExecutionTable = specs::ExecutionTable;
 
+pub(crate) fn compute_slice_capability(k: u32) -> u32 {
+    ((1 << k) - RESERVE_ROWS as u32 - 1024) / EVENT_TABLE_ENTRY_ROWS as u32
+}
+
 pub struct ZkWasmCircuit<F: FieldExt> {
+    pub k: u32,
     pub slice: Slice,
-    pub slice_capability: usize,
     _data: PhantomData<F>,
 }
 
 impl<F: FieldExt> ZkWasmCircuit<F> {
-    pub fn new(slice: Slice, slice_capability: usize) -> Self {
-        ZkWasmCircuit {
+    pub fn new(k: u32, slice: Slice) -> Result<Self, BuildingCircuitError> {
+        {
+            // entries is empty when called by without_witness
+            let allocated_memory_pages = slice
+                .etable
+                .entries()
+                .last()
+                .map(|entry| entry.allocated_memory_pages);
+            let maximal_pages = compute_maximal_pages(k);
+            if let Some(allocated_memory_pages) = allocated_memory_pages {
+                if allocated_memory_pages > maximal_pages {
+                    return Err(BuildingCircuitError::PagesExceedLimit(
+                        allocated_memory_pages,
+                        maximal_pages,
+                        k,
+                    ));
+                }
+            }
+        }
+
+        {
+            let etable_entires = slice.etable.entries().len() as u32;
+            let etable_capacity = compute_slice_capability(k);
+
+            if etable_entires > etable_capacity {
+                return Err(BuildingCircuitError::EtableEntriesExceedLimit(
+                    etable_entires as u32,
+                    etable_capacity as u32,
+                    k,
+                ));
+            }
+        }
+
+        Ok(ZkWasmCircuit {
+            k,
             slice,
-            slice_capability,
             _data: PhantomData,
-        }
+        })
     }
 }
 

diff --git a/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs b/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs
@@ -23,6 +23,7 @@ use specs::slice::Slice;
 use crate::circuits::bit_table::BitTableChip;
 use crate::circuits::bit_table::BitTableConfig;
 use crate::circuits::bit_table::BitTableTrait;
+use crate::circuits::compute_slice_capability;
 use crate::circuits::etable::EventTableChip;
 use crate::circuits::etable::EventTableConfig;
 use crate::circuits::external_host_call_table::ExternalHostCallChip;
@@ -108,6 +109,7 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
 
     fn without_witnesses(&self) -> Self {
         ZkWasmCircuit::new(
+            self.k,
             // fill slice like circuit_without_witness
             Slice {
                 itable: self.slice.itable.clone(),
@@ -127,8 +129,8 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
 
                 is_last_slice: self.slice.is_last_slice,
             },
-            self.slice_capability,
         )
+        .unwrap()
     }
 
     fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
@@ -231,7 +233,7 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
         let frame_table_chip = JumpTableChip::new(config.jtable, config.max_available_rows);
         let echip = EventTableChip::new(
             config.etable,
-            self.slice_capability,
+            compute_slice_capability(self.k) as usize,
             config.max_available_rows,
         );
         let bit_chip = BitTableChip::new(config.bit_table, config.max_available_rows);

diff --git a/crates/zkwasm/src/error.rs b/crates/zkwasm/src/error.rs
@@ -0,0 +1,17 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum CompilationError {}
+
+#[derive(Debug, Error)]
+pub enum ExecutionError {}
+
+#[derive(Debug, Error)]
+pub enum BuildingCircuitError {
+    #[error("Only support single slice for non-continuation mode but {0} provided. You could increase K or enable continuation feature.")]
+    MultiSlicesNotSupport(usize),
+    #[error("Allocated pages({0}) exceed the limit({1}). Current K is {2}, consider increase the circuit size K.")]
+    PagesExceedLimit(u32, u32, u32),
+    #[error("Etable entries({0}) exceed the limit({1}). Current K is {2}, consider increase the circuit size K.")]
+    EtableEntriesExceedLimit(u32, u32, u32),
+}
diff --git a/crates/zkwasm/src/lib.rs b/crates/zkwasm/src/lib.rs
@@ -4,6 +4,7 @@
 
 pub mod checksum;
 pub mod circuits;
+pub mod error;
 pub mod foreign;
 pub mod loader;
 pub mod runtime;

diff --git a/crates/zkwasm/src/loader/mod.rs b/crates/zkwasm/src/loader/mod.rs
@@ -26,10 +26,11 @@ use wasmi::RuntimeValue;
 
 use crate::checksum::CompilationTableWithParams;
 use crate::checksum::ImageCheckSum;
+use crate::circuits::compute_slice_capability;
 use crate::circuits::config::init_zkwasm_runtime;
-use crate::circuits::etable::EVENT_TABLE_ENTRY_ROWS;
 use crate::circuits::image_table::compute_maximal_pages;
 use crate::circuits::ZkWasmCircuit;
+use crate::error::BuildingCircuitError;
 use crate::loader::err::Error;
 use crate::loader::err::PreCheckErr;
 #[cfg(feature = "profile")]
@@ -107,20 +108,21 @@ impl<E: MultiMillerLoop, T, EnvBuilder: HostEnvBuilder<Arg = T>> ZkWasmLoader<E,
             dryrun,
             &self.phantom_functions,
             backend,
-            self.compute_slice_capability() as u32,
+            compute_slice_capability(self.k),
         )
     }
 
     pub fn circuit_without_witness(
         &self,
         envconfig: EnvBuilder::HostConfig,
         is_last_slice: bool,
-    ) -> Result<ZkWasmCircuit<E::Scalar>> {
+    ) -> Result<ZkWasmCircuit<E::Scalar>, BuildingCircuitError> {
         let (env, _wasm_runtime_io) = EnvBuilder::create_env_without_value(self.k, envconfig);
 
-        let compiled_module = self.compile(&env, false, TraceBackend::Memory)?;
+        let compiled_module = self.compile(&env, false, TraceBackend::Memory).unwrap();
 
-        Ok(ZkWasmCircuit::new(
+        ZkWasmCircuit::new(
+            self.k,
             Slice {
                 itable: compiled_module.tables.itable.clone(),
                 br_table: compiled_module.tables.br_table.clone(),
@@ -137,8 +139,7 @@ impl<E: MultiMillerLoop, T, EnvBuilder: HostEnvBuilder<Arg = T>> ZkWasmLoader<E,
 
                 is_last_slice,
             },
-            self.compute_slice_capability(),
-        ))
+        )
     }
 
     /// Create a ZkWasm Loader
@@ -201,11 +202,6 @@ impl<E: MultiMillerLoop, T, EnvBuilder: HostEnvBuilder<Arg = T>> ZkWasmLoader<E,
 }
 
 impl<E: MultiMillerLoop, T, EnvBuilder: HostEnvBuilder<Arg = T>> ZkWasmLoader<E, T, EnvBuilder> {
-    pub(crate) fn compute_slice_capability(&self) -> usize {
-        // FIXME
-        ((1 << self.k) - 200) / EVENT_TABLE_ENTRY_ROWS as usize
-    }
-
     pub fn run(
         &self,
         arg: T,
@@ -224,8 +220,11 @@ impl<E: MultiMillerLoop, T, EnvBuilder: HostEnvBuilder<Arg = T>> ZkWasmLoader<E,
         Ok(result)
     }
 
-    pub fn slice(&self, execution_result: ExecutionResult<RuntimeValue>) -> Slices<E::Scalar> {
-        Slices::new(execution_result.tables, self.compute_slice_capability())
+    pub fn slice(
+        &self,
+        execution_result: ExecutionResult<RuntimeValue>,
+    ) -> Result<Slices<E::Scalar>, BuildingCircuitError> {
+        Slices::new(self.k, execution_result.tables)
     }
 
     pub fn mock_test(

diff --git a/crates/zkwasm/src/loader/slice.rs b/crates/zkwasm/src/loader/slice.rs
@@ -16,10 +16,13 @@ use std::collections::VecDeque;
 use std::sync::Arc;
 
 use crate::circuits::ZkWasmCircuit;
+use crate::error::BuildingCircuitError;
 use crate::runtime::state::UpdateInitMemoryTable;
 use crate::runtime::state::UpdateInitializationState;
 
 pub struct Slices<F: FieldExt> {
+    k: u32,
+
     itable: Arc<InstructionTable>,
     br_table: Arc<BrTable>,
     elem_table: Arc<ElemTable>,
@@ -31,15 +34,22 @@ pub struct Slices<F: FieldExt> {
     initialization_state: Arc<InitializationState<u32>>,
     etables: VecDeque<EventTableBackend>,
 
-    // the length of etable entries
-    capability: usize,
-
     _marker: std::marker::PhantomData<F>,
 }
 
 impl<F: FieldExt> Slices<F> {
-    pub fn new(tables: Tables, capability: usize) -> Self {
-        Self {
+    pub fn new(k: u32, tables: Tables) -> Result<Self, BuildingCircuitError> {
+        if cfg!(not(feature = "continuation")) {
+            let slices = tables.execution_tables.etable.len();
+
+            if slices != 1 {
+                return Err(BuildingCircuitError::MultiSlicesNotSupport(slices));
+            }
+        }
+
+        Ok(Self {
+            k,
+
             itable: tables.compilation_tables.itable,
             br_table: tables.compilation_tables.br_table,
             elem_table: tables.compilation_tables.elem_table,
@@ -52,10 +62,8 @@ impl<F: FieldExt> Slices<F> {
 
             etables: tables.execution_tables.etable.into(),
 
-            capability,
-
             _marker: std::marker::PhantomData,
-        }
+        })
     }
 
     pub fn mock_test_all(self, k: u32, instances: Vec<F>) -> anyhow::Result<()> {
@@ -64,7 +72,7 @@ impl<F: FieldExt> Slices<F> {
         let mut iter = self.into_iter();
 
         while let Some(slice) = iter.next() {
-            let prover = MockProver::run(k, &slice, vec![instances.clone()])?;
+            let prover = MockProver::run(k, &slice?, vec![instances.clone()])?;
             assert_eq!(prover.verify(), Ok(()));
         }
 
@@ -73,7 +81,7 @@ impl<F: FieldExt> Slices<F> {
 }
 
 impl<F: FieldExt> Iterator for Slices<F> {
-    type Item = ZkWasmCircuit<F>;
+    type Item = Result<ZkWasmCircuit<F>, BuildingCircuitError>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.etables.is_empty() {
@@ -127,6 +135,8 @@ impl<F: FieldExt> Iterator for Slices<F> {
         self.imtable = post_imtable;
         self.initialization_state = post_initialization_state;
 
-        Some(ZkWasmCircuit::new(slice, self.capability))
+        let circuit = ZkWasmCircuit::new(self.k, slice);
+
+        Some(circuit)
     }
 }
diff --git a/crates/zkwasm/src/test/mod.rs b/crates/zkwasm/src/test/mod.rs
@@ -40,7 +40,7 @@ pub fn test_circuit_with_env(
     )?;
     let instances = execution_result.public_inputs_and_outputs();
     loader
-        .slice(execution_result)
+        .slice(execution_result)?
         .mock_test_all(MIN_K, instances)?;
 
     Ok(())

diff --git a/crates/zkwasm/src/test/test_rlp.rs b/crates/zkwasm/src/test/test_rlp.rs
@@ -166,7 +166,7 @@ fn run_test() -> Result<()> {
 
     let instances = result.public_inputs_and_outputs();
 
-    let slices = loader.slice(result).into_iter();
+    let slices = loader.slice(result)?.into_iter();
 
     slices.mock_test_all(K, instances)?;