From 3e234baec02592600e5ef0e7752f0fe148f3890b Mon Sep 17 00:00:00 2001 From: Zhang Junyu Date: Tue, 12 Dec 2023 10:09:17 +0000 Subject: [PATCH] refine image table --- crates/specs/src/brtable.rs | 2 +- crates/specs/src/lib.rs | 4 +- crates/zkwasm/src/checksum/mod.rs | 2 +- .../zkwasm/src/circuits/image_table/assign.rs | 32 +-- crates/zkwasm/src/circuits/image_table/mod.rs | 223 +----------------- crates/zkwasm/src/circuits/jtable/assign.rs | 8 +- .../circuits/post_image_table/continuation.rs | 29 +-- .../src/circuits/post_image_table/mod.rs | 10 +- .../src/circuits/post_image_table/trivial.rs | 10 +- .../zkwasm/src/circuits/test_circuit/mod.rs | 24 +- .../zkwasm/src/circuits/utils/image_table.rs | 209 ++++++++++++++-- crates/zkwasm/src/continuation/slice.rs | 1 + .../zkwasm/src/runtime/wasmi_interpreter.rs | 28 ++- 13 files changed, 252 insertions(+), 330 deletions(-) diff --git a/crates/specs/src/brtable.rs b/crates/specs/src/brtable.rs index c6c3b3211..b23953c0b 100644 --- a/crates/specs/src/brtable.rs +++ b/crates/specs/src/brtable.rs @@ -12,7 +12,7 @@ pub struct BrTableEntry { pub dst_pc: u32, } -#[derive(Debug)] +#[derive(Default, Serialize, Debug, Clone)] pub struct BrTable(Vec); impl BrTable { diff --git a/crates/specs/src/lib.rs b/crates/specs/src/lib.rs index ca6a6b02d..149335ff0 100644 --- a/crates/specs/src/lib.rs +++ b/crates/specs/src/lib.rs @@ -8,6 +8,7 @@ use std::io::Write; use std::path::PathBuf; use std::sync::Arc; +use brtable::BrTable; use brtable::ElemTable; use configure_table::ConfigureTable; use etable::EventTable; @@ -42,10 +43,11 @@ pub mod state; pub mod step; pub mod types; -#[derive(Default, Serialize, Debug, Clone)] +#[derive(Serialize, Debug, Clone)] pub struct CompilationTable { pub itable: Arc, pub imtable: InitMemoryTable, + pub br_table: Arc, pub elem_table: Arc, pub configure_table: Arc, pub static_jtable: Arc>, diff --git a/crates/zkwasm/src/checksum/mod.rs b/crates/zkwasm/src/checksum/mod.rs index 5d03b8743..3ec68d439 100644 --- a/crates/zkwasm/src/checksum/mod.rs +++ b/crates/zkwasm/src/checksum/mod.rs @@ -3,7 +3,7 @@ use halo2_proofs::arithmetic::CurveAffine; use halo2_proofs::poly::commitment::Params; use specs::CompilationTable; -use crate::circuits::image_table::EncodeCompilationTableValues; +use crate::circuits::utils::image_table::EncodeCompilationTableValues; pub trait ImageCheckSum { fn checksum(&self) -> Output; diff --git a/crates/zkwasm/src/circuits/image_table/assign.rs b/crates/zkwasm/src/circuits/image_table/assign.rs index 70c245076..b0eb61b6f 100644 --- a/crates/zkwasm/src/circuits/image_table/assign.rs +++ b/crates/zkwasm/src/circuits/image_table/assign.rs @@ -7,25 +7,15 @@ use halo2_proofs::circuit::Layouter; use halo2_proofs::plonk::Error; use super::ImageTableChip; -use super::ImageTableLayouter; use crate::circuits::utils::image_table::ImageTableAssigner; +use crate::circuits::utils::image_table::ImageTableLayouter; use crate::circuits::utils::Context; -impl< - const INIT_MEMORY_ENTRIES_OFFSET: usize, - const STACK_LIMIT: usize, - const GLOBAL_LIMIT: usize, - F: FieldExt, - > ImageTableChip -{ +impl ImageTableChip { pub(crate) fn assign( self, layouter: &mut impl Layouter, - image_table_assigner: &mut ImageTableAssigner< - INIT_MEMORY_ENTRIES_OFFSET, - STACK_LIMIT, - GLOBAL_LIMIT, - >, + image_table_assigner: &mut ImageTableAssigner, image_table: ImageTableLayouter, permutation_cells: ImageTableLayouter>, ) -> Result>, Error> { @@ -97,8 +87,6 @@ impl< image_table .instructions - .as_ref() - .unwrap() .iter() .map(|entry| { let offset = ctx.borrow().offset; @@ -123,9 +111,7 @@ impl< ctx.borrow_mut().offset = base_offset; image_table - .br_table - .as_ref() - .unwrap() + .br_table_entires .iter() .map(|entry| { let offset = ctx.borrow().offset; @@ -166,8 +152,6 @@ impl< image_table .init_memory_entries - .as_ref() - .unwrap() .iter() .map(|entry| { let offset = ctx.borrow().offset; @@ -199,10 +183,10 @@ impl< Ok(ImageTableLayouter { initialization_state: result.initialization_state, static_frame_entries: result.static_frame_entries, - instructions: Some(result.instructions), - br_table: Some(result.br_table_entires), - padding: Some(result.padding_entires), - init_memory_entries: None, + instructions: result.instructions, + br_table_entires: result.br_table_entires, + padding_entires: result.padding_entires, + init_memory_entries: result.init_memory_entries, }) }, ) diff --git a/crates/zkwasm/src/circuits/image_table/mod.rs b/crates/zkwasm/src/circuits/image_table/mod.rs index 3a40534e8..b54234233 100644 --- a/crates/zkwasm/src/circuits/image_table/mod.rs +++ b/crates/zkwasm/src/circuits/image_table/mod.rs @@ -4,32 +4,18 @@ use halo2_proofs::plonk::Column; use halo2_proofs::plonk::Expression; use halo2_proofs::plonk::Fixed; use halo2_proofs::plonk::VirtualCells; -use num_bigint::BigUint; -use specs::brtable::BrTable; -use specs::brtable::ElemTable; -use specs::encode::image_table::ImageTableEncoder; -use specs::imtable::InitMemoryTable; -use specs::imtable::InitMemoryTableEntry; -use specs::itable::InstructionTable; -use specs::jtable::StaticFrameEntry; -use specs::mtable::LocationType; -use specs::mtable::VarType; -use specs::state::InitializationState; -use specs::CompilationTable; use std::marker::PhantomData; use wasmi::DEFAULT_VALUE_STACK_LIMIT; -use crate::circuits::config::zkwasm_k; -use crate::circuits::utils::bn_to_field; use crate::curr; use super::test_circuit::RESERVE_ROWS; +use super::utils::image_table::INIT_MEMORY_ENTRIES_OFFSET; mod assign; mod configure; pub const IMAGE_COL_NAME: &str = "img_col"; -pub const INIT_MEMORY_ENTRIES_OFFSET: usize = 40960; /* * 8192: 64 * 1024 / 8 * A page is 64KB, an entry is 8B @@ -50,198 +36,6 @@ pub fn compute_maximal_pages(k: u32) -> u32 { pages } -pub(crate) struct InitMemoryLayouter { - pub(crate) stack: u32, - pub(crate) global: u32, - pub(crate) pages: u32, -} - -impl InitMemoryLayouter { - fn for_each(self, mut f: impl FnMut((LocationType, u32))) { - for offset in 0..self.stack { - f((LocationType::Stack, offset)) - } - - for offset in 0..self.global { - f((LocationType::Global, offset)) - } - - for offset in 0..(self.pages * PAGE_ENTRIES) { - f((LocationType::Heap, offset)) - } - } -} - -pub struct ImageTableLayouter { - pub initialization_state: InitializationState, - pub static_frame_entries: Vec<(T, T)>, - pub instructions: Option>, - pub br_table: Option>, - pub padding: Option>, - pub init_memory_entries: Option>, -} - -impl ImageTableLayouter { - pub fn plain(&self) -> Vec { - let mut buf = vec![]; - - buf.append(&mut self.initialization_state.plain()); - buf.append( - &mut self - .static_frame_entries - .clone() - .to_vec() - .into_iter() - .map(|(enable, fid)| vec![enable, fid]) - .collect::>>() - .concat(), - ); - buf.append(&mut self.instructions.clone().unwrap()); - buf.append(&mut self.br_table.clone().unwrap()); - buf.append(&mut vec![F::zero(); INIT_MEMORY_ENTRIES_OFFSET - buf.len()]); - buf.append(&mut self.init_memory_entries.clone().unwrap()); - - buf - } -} - -pub trait EncodeCompilationTableValues { - fn encode_compilation_table_values(&self) -> ImageTableLayouter; -} - -impl EncodeCompilationTableValues for CompilationTable { - fn encode_compilation_table_values(&self) -> ImageTableLayouter { - fn msg_of_initialization_state( - initialization_state: &InitializationState, - ) -> InitializationState { - initialization_state.map(|field| F::from(*field as u64)) - } - - fn msg_of_instruction_table(instruction_table: &InstructionTable) -> Vec { - let mut cells = vec![]; - - cells.push(bn_to_field( - &ImageTableEncoder::Instruction.encode(BigUint::from(0u64)), - )); - - for e in instruction_table.entries() { - cells.push(bn_to_field( - &ImageTableEncoder::Instruction.encode(e.encode()), - )); - } - - cells - } - - fn msg_of_br_table(br_table: &BrTable, elem_table: &ElemTable) -> Vec { - let mut cells = vec![]; - - cells.push(bn_to_field( - &ImageTableEncoder::BrTable.encode(BigUint::from(0u64)), - )); - - for e in br_table.entries() { - cells.push(bn_to_field(&ImageTableEncoder::BrTable.encode(e.encode()))); - } - - for e in elem_table.entries() { - cells.push(bn_to_field(&ImageTableEncoder::BrTable.encode(e.encode()))); - } - - cells - } - - fn msg_of_init_memory_table(init_memory_table: &InitMemoryTable) -> Vec { - let mut cells = vec![]; - - cells.push(bn_to_field( - &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), - )); - - let layouter = InitMemoryLayouter { - stack: DEFAULT_VALUE_STACK_LIMIT as u32, - global: DEFAULT_VALUE_STACK_LIMIT as u32, - pages: compute_maximal_pages(zkwasm_k()), - }; - - layouter.for_each(|(ltype, offset)| { - if let Some(entry) = init_memory_table.try_find(ltype, offset) { - cells.push(bn_to_field::( - &ImageTableEncoder::InitMemory.encode(entry.encode()), - )); - } else if ltype == LocationType::Heap { - let entry = InitMemoryTableEntry { - ltype, - is_mutable: true, - offset, - vtype: VarType::I64, - value: 0, - eid: 0, - }; - - cells.push(bn_to_field::( - &ImageTableEncoder::InitMemory.encode(entry.encode()), - )); - } else { - cells.push(bn_to_field::( - &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), - )); - } - }); - - cells - } - - fn msg_of_static_frame_table( - static_frame_table: &Vec, - ) -> Vec<(F, F)> { - let mut cells = static_frame_table - .into_iter() - .map(|entry| (F::one(), bn_to_field(&entry.encode()))) - .collect::>(); - - cells.resize( - 2, - ( - F::zero(), - bn_to_field( - &StaticFrameEntry { - enable: false, - frame_id: 0, - next_frame_id: 0, - callee_fid: 0, - fid: 0, - iid: 0, - } - .encode(), - ), - ), - ); - - cells - } - - let initialization_state = msg_of_initialization_state(&self.initialization_state); - let static_frame_entries = msg_of_static_frame_table(&self.static_jtable); - - let instructions = Some(msg_of_instruction_table(&self.itable)); - let br_table = Some(msg_of_br_table( - &self.itable.create_brtable(), - &self.elem_table, - )); - let init_memory_entries = Some(msg_of_init_memory_table(&self.imtable)); - - ImageTableLayouter { - initialization_state, - static_frame_entries, - instructions, - br_table, - padding: None, - init_memory_entries, - } - } -} - #[derive(Clone)] pub struct ImageTableConfig { _memory_addr_sel: Column, @@ -256,22 +50,11 @@ impl ImageTableConfig { } #[derive(Clone)] -pub struct ImageTableChip< - const INIT_MEMORY_ENTRIES_OFFSET: usize, - const STACK_LIMIT: usize, - const GLOBAL_LIMIT: usize, - F: FieldExt, -> { +pub struct ImageTableChip { config: ImageTableConfig, } -impl< - const INIT_MEMORY_ENTRIES_OFFSET: usize, - const STACK_LIMIT: usize, - const GLOBAL_LIMIT: usize, - F: FieldExt, - > ImageTableChip -{ +impl ImageTableChip { pub fn new(config: ImageTableConfig) -> Self { ImageTableChip { config } } diff --git a/crates/zkwasm/src/circuits/jtable/assign.rs b/crates/zkwasm/src/circuits/jtable/assign.rs index 8510d52de..07a7d143b 100644 --- a/crates/zkwasm/src/circuits/jtable/assign.rs +++ b/crates/zkwasm/src/circuits/jtable/assign.rs @@ -206,7 +206,13 @@ impl JumpTableChip { self.init(ctx)?; ctx.reset(); - let mut rest_jops = jtable.entries().len() as u64 * 2 + static_entries.len() as u64; + let mut rest_jops = jtable.entries().len() as u64 * 2; + + for entry in static_entries { + if entry.enable { + rest_jops += 1; + } + } let frame_table_start_jump_cells = self.assign_static_entries(ctx, &mut rest_jops, static_entries)?; diff --git a/crates/zkwasm/src/circuits/post_image_table/continuation.rs b/crates/zkwasm/src/circuits/post_image_table/continuation.rs index d2a78c2cf..7d7b3699c 100644 --- a/crates/zkwasm/src/circuits/post_image_table/continuation.rs +++ b/crates/zkwasm/src/circuits/post_image_table/continuation.rs @@ -16,11 +16,10 @@ use specs::mtable::LocationType; use wasmi::DEFAULT_VALUE_STACK_LIMIT; use crate::circuits::image_table::ImageTableConfig; -use crate::circuits::image_table::ImageTableLayouter; -use crate::circuits::image_table::INIT_MEMORY_ENTRIES_OFFSET; use crate::circuits::mtable::MemoryTableConfig; use crate::circuits::utils::bn_to_field; use crate::circuits::utils::image_table::ImageTableAssigner; +use crate::circuits::utils::image_table::ImageTableLayouter; use crate::circuits::utils::Context; use crate::constant_from; use crate::curr; @@ -106,11 +105,7 @@ impl PostImageTableChipTrait fn assign( self, layouter: &mut impl Layouter, - image_table_assigner: &mut ImageTableAssigner< - INIT_MEMORY_ENTRIES_OFFSET, - DEFAULT_VALUE_STACK_LIMIT, - DEFAULT_VALUE_STACK_LIMIT, - >, + image_table_assigner: &mut ImageTableAssigner, pre_image_table: ImageTableLayouter, post_image_table: ImageTableLayouter, permutation_cells: ImageTableLayouter>, @@ -185,8 +180,6 @@ impl PostImageTableChipTrait permutation_cells .instructions - .as_ref() - .unwrap() .iter() .map(|entry| { let offset = ctx.borrow().offset; @@ -211,9 +204,7 @@ impl PostImageTableChipTrait ctx.borrow_mut().offset = base_offset; permutation_cells - .br_table - .as_ref() - .unwrap() + .br_table_entires .iter() .map(|entry| { let offset = ctx.borrow().offset; @@ -238,9 +229,7 @@ impl PostImageTableChipTrait ctx.borrow_mut().offset = start_offset; permutation_cells - .padding - .as_ref() - .unwrap() + .padding_entires .iter() .map(|entry| { let offset = ctx.borrow().offset; @@ -335,16 +324,8 @@ impl PostImageTableChipTrait pre_image_table .init_memory_entries - .as_ref() - .unwrap() .iter() - .zip( - post_image_table - .init_memory_entries - .as_ref() - .unwrap() - .iter(), - ) + .zip(post_image_table.init_memory_entries.iter()) .map(|(pre, post)| { let entry = ctx.borrow_mut().region.assign_advice( || "post image table: init memory", diff --git a/crates/zkwasm/src/circuits/post_image_table/mod.rs b/crates/zkwasm/src/circuits/post_image_table/mod.rs index 1f1ca6192..e444451c9 100644 --- a/crates/zkwasm/src/circuits/post_image_table/mod.rs +++ b/crates/zkwasm/src/circuits/post_image_table/mod.rs @@ -5,13 +5,11 @@ use halo2_proofs::plonk::Column; use halo2_proofs::plonk::ConstraintSystem; use halo2_proofs::plonk::Error; use halo2_proofs::plonk::Fixed; -use wasmi::DEFAULT_VALUE_STACK_LIMIT; use super::image_table::ImageTableConfig; -use super::image_table::ImageTableLayouter; -use super::image_table::INIT_MEMORY_ENTRIES_OFFSET; use super::mtable::MemoryTableConfig; use super::utils::image_table::ImageTableAssigner; +use super::utils::image_table::ImageTableLayouter; pub(self) mod continuation; pub(self) mod trivial; @@ -34,11 +32,7 @@ pub(in crate::circuits) trait PostImageTableChipTrait< fn assign( self, layouter: &mut impl Layouter, - image_table_assigner: &mut ImageTableAssigner< - INIT_MEMORY_ENTRIES_OFFSET, - DEFAULT_VALUE_STACK_LIMIT, - DEFAULT_VALUE_STACK_LIMIT, - >, + image_table_assigner: &mut ImageTableAssigner, pre_image_table: ImageTableLayouter, post_image_table: ImageTableLayouter, permutation_cells: ImageTableLayouter>, diff --git a/crates/zkwasm/src/circuits/post_image_table/trivial.rs b/crates/zkwasm/src/circuits/post_image_table/trivial.rs index a146b0c0e..d73416d45 100644 --- a/crates/zkwasm/src/circuits/post_image_table/trivial.rs +++ b/crates/zkwasm/src/circuits/post_image_table/trivial.rs @@ -7,13 +7,11 @@ use halo2_proofs::plonk::Column; use halo2_proofs::plonk::ConstraintSystem; use halo2_proofs::plonk::Error; use halo2_proofs::plonk::Fixed; -use wasmi::DEFAULT_VALUE_STACK_LIMIT; use crate::circuits::image_table::ImageTableConfig; -use crate::circuits::image_table::ImageTableLayouter; -use crate::circuits::image_table::INIT_MEMORY_ENTRIES_OFFSET; use crate::circuits::mtable::MemoryTableConfig; use crate::circuits::utils::image_table::ImageTableAssigner; +use crate::circuits::utils::image_table::ImageTableLayouter; use super::PostImageTableChipTrait; use super::PostImageTableConfigTrait; @@ -48,11 +46,7 @@ impl PostImageTableChipTrait> fn assign( self, _layouter: &mut impl Layouter, - _image_table_assigner: &mut ImageTableAssigner< - INIT_MEMORY_ENTRIES_OFFSET, - DEFAULT_VALUE_STACK_LIMIT, - DEFAULT_VALUE_STACK_LIMIT, - >, + _image_table_assigner: &mut ImageTableAssigner, _pre_image_table: ImageTableLayouter, _post_image_table: ImageTableLayouter, _permutation_cells: ImageTableLayouter>, diff --git a/crates/zkwasm/src/circuits/test_circuit/mod.rs b/crates/zkwasm/src/circuits/test_circuit/mod.rs index 0ce571025..bafe47c6c 100644 --- a/crates/zkwasm/src/circuits/test_circuit/mod.rs +++ b/crates/zkwasm/src/circuits/test_circuit/mod.rs @@ -14,7 +14,6 @@ use log::debug; use log::info; use specs::ExecutionTable; use specs::Tables; -use wasmi::DEFAULT_VALUE_STACK_LIMIT; use crate::circuits::bit_table::BitTableChip; use crate::circuits::bit_table::BitTableConfig; @@ -23,10 +22,7 @@ use crate::circuits::etable::EventTableConfig; use crate::circuits::external_host_call_table::ExternalHostCallChip; use crate::circuits::external_host_call_table::ExternalHostCallTableConfig; use crate::circuits::image_table::compute_maximal_pages; -use crate::circuits::image_table::EncodeCompilationTableValues; use crate::circuits::image_table::ImageTableChip; -use crate::circuits::image_table::ImageTableLayouter; -use crate::circuits::image_table::INIT_MEMORY_ENTRIES_OFFSET; use crate::circuits::jtable::JumpTableChip; use crate::circuits::jtable::JumpTableConfig; use crate::circuits::mtable::MemoryTableChip; @@ -36,7 +32,9 @@ use crate::circuits::post_image_table::PostImageTableChipTrait; use crate::circuits::post_image_table::PostImageTableConfigTrait; use crate::circuits::rtable::RangeTableChip; use crate::circuits::rtable::RangeTableConfig; +use crate::circuits::utils::image_table::EncodeCompilationTableValues; use crate::circuits::utils::image_table::ImageTableAssigner; +use crate::circuits::utils::image_table::ImageTableLayouter; use crate::circuits::utils::table_entry::EventTableWithMemoryInfo; use crate::circuits::utils::table_entry::MemoryWritingTable; use crate::circuits::utils::Context; @@ -316,11 +314,7 @@ impl Circuit for TestCircuit { )? ); - let mut image_table_assigner = ImageTableAssigner::< - INIT_MEMORY_ENTRIES_OFFSET, - DEFAULT_VALUE_STACK_LIMIT, - DEFAULT_VALUE_STACK_LIMIT, - >::new( + let mut image_table_assigner = ImageTableAssigner::new( // Add one for default lookup value self.tables.compilation_tables.itable.entries().len() + 1, // FIXME: avoid compute @@ -346,10 +340,10 @@ impl Circuit for TestCircuit { ImageTableLayouter { initialization_state: etable_permutation_cells.pre_initialization_state, static_frame_entries, - instructions: None, - br_table: None, - padding: None, - init_memory_entries: None, + instructions: vec![], + br_table_entires: vec![], + padding_entires: vec![], + init_memory_entries: vec![], } )? ); @@ -369,8 +363,8 @@ impl Circuit for TestCircuit { initialization_state: etable_permutation_cells.post_initialization_state, static_frame_entries: pre_image_table_cells.static_frame_entries, instructions: pre_image_table_cells.instructions, - br_table: pre_image_table_cells.br_table, - padding: pre_image_table_cells.padding, + br_table_entires: pre_image_table_cells.br_table_entires, + padding_entires: pre_image_table_cells.padding_entires, init_memory_entries: pre_image_table_cells.init_memory_entries, }, rest_memory_writing_ops_cell, diff --git a/crates/zkwasm/src/circuits/utils/image_table.rs b/crates/zkwasm/src/circuits/utils/image_table.rs index f483c8339..5689e01e0 100644 --- a/crates/zkwasm/src/circuits/utils/image_table.rs +++ b/crates/zkwasm/src/circuits/utils/image_table.rs @@ -1,7 +1,44 @@ +use anyhow::Error; +use halo2_proofs::arithmetic::FieldExt; +use num_bigint::BigUint; +use specs::encode::image_table::ImageTableEncoder; +use specs::imtable::InitMemoryTableEntry; +use specs::mtable::LocationType; +use specs::mtable::VarType; use specs::state::InitializationState; +use specs::CompilationTable; +use wasmi::DEFAULT_VALUE_STACK_LIMIT; +use crate::circuits::config::zkwasm_k; +use crate::circuits::image_table::compute_maximal_pages; use crate::circuits::image_table::PAGE_ENTRIES; use crate::circuits::jtable::STATIC_FRAME_ENTRY_IMAGE_TABLE_ENTRY; +use crate::circuits::jtable::STATIC_FRAME_ENTRY_NUMBER; +use crate::circuits::utils::bn_to_field; + +pub const STACK_CAPABILITY: usize = DEFAULT_VALUE_STACK_LIMIT; +pub const GLOBAL_CAPABILITY: usize = DEFAULT_VALUE_STACK_LIMIT; +pub const INIT_MEMORY_ENTRIES_OFFSET: usize = 40960; + +pub(crate) struct InitMemoryLayouter { + pub(crate) pages: u32, +} + +impl InitMemoryLayouter { + fn for_each(self, mut f: impl FnMut((LocationType, u32))) { + for offset in 0..STACK_CAPABILITY { + f((LocationType::Stack, offset as u32)) + } + + for offset in 0..GLOBAL_CAPABILITY { + f((LocationType::Global, offset as u32)) + } + + for offset in 0..(self.pages * PAGE_ENTRIES) { + f((LocationType::Heap, offset)) + } + } +} /* * -------------------- @@ -23,21 +60,17 @@ use crate::circuits::jtable::STATIC_FRAME_ENTRY_IMAGE_TABLE_ENTRY; * -------------------- */ #[allow(dead_code)] -pub(crate) struct Layouter { +pub(crate) struct ImageTableLayouter { pub(crate) initialization_state: InitializationState, pub(crate) static_frame_entries: Vec<(T, T)>, pub(crate) instructions: Vec, pub(crate) br_table_entires: Vec, // NOTE: unused instructions and br_table entries. pub(crate) padding_entires: Vec, - pub(crate) init_memory_entires: Vec, + pub(crate) init_memory_entries: Vec, } -pub(crate) struct ImageTableAssigner< - const INIT_MEMORY_OFFSET: usize, - const STACK_CAPABILITY: usize, - const GLOBAL_CAPABILITY: usize, -> { +pub(crate) struct ImageTableAssigner { pub(crate) heap_capability: u32, initialization_state_offset: usize, @@ -48,12 +81,7 @@ pub(crate) struct ImageTableAssigner< init_memory_offset: usize, } -impl< - const INIT_MEMORY_OFFSET: usize, - const STACK_CAPABILITY: usize, - const GLOBAL_CAPABILITY: usize, - > ImageTableAssigner -{ +impl ImageTableAssigner { pub fn new(instruction_number: usize, br_table_number: usize, pages_capability: u32) -> Self { let initialization_state_offset = 0; let static_frame_entries_offset = @@ -61,7 +89,7 @@ impl< let instruction_offset = static_frame_entries_offset + STATIC_FRAME_ENTRY_IMAGE_TABLE_ENTRY; let br_table_offset = instruction_offset + instruction_number; let padding_offset = br_table_offset + br_table_number; - let init_memory_offset = INIT_MEMORY_OFFSET; + let init_memory_offset = INIT_MEMORY_ENTRIES_OFFSET; assert!(padding_offset <= init_memory_offset); @@ -112,7 +140,7 @@ impl< padding_handler(self.padding_offset, self.padding_offset) } - pub fn exec_init_memory_entires( + pub fn exec_init_memory_entries( &mut self, mut init_memory_entries_handler: impl FnMut(usize) -> Result, Error>, ) -> Result, Error> { @@ -127,21 +155,164 @@ impl< br_table_handler: impl FnMut(usize) -> Result, Error>, padding_handler: impl FnMut(usize, usize) -> Result, Error>, init_memory_entries_handler: impl FnMut(usize) -> Result, Error>, - ) -> Result, Error> { + ) -> Result, Error> { let initialization_state = self.exec_initialization_state(initialization_state_handler)?; let static_frame_entries = self.exec_static_frame_entries(static_frame_entries_handler)?; let instructions = self.exec_instruction(instruction_handler)?; let br_table_entires = self.exec_br_table_entires(br_table_handler)?; let padding_entires = self.exec_padding_entires(padding_handler)?; - let init_memory_entires = self.exec_init_memory_entires(init_memory_entries_handler)?; + let init_memory_entries = self.exec_init_memory_entries(init_memory_entries_handler)?; - Ok(Layouter { + Ok(ImageTableLayouter { initialization_state, static_frame_entries, instructions, br_table_entires, padding_entires, - init_memory_entires, + init_memory_entries, }) } } + +pub(crate) trait EncodeCompilationTableValues { + fn encode_compilation_table_values(&self) -> ImageTableLayouter; +} + +impl EncodeCompilationTableValues for CompilationTable { + fn encode_compilation_table_values(&self) -> ImageTableLayouter { + // FIXME: ugly + let pages_capability = compute_maximal_pages(zkwasm_k()); + + let initialization_state_handler = + |_| Ok(self.initialization_state.map(|v| F::from((*v) as u64))); + + let static_frame_entries_handler = |_| { + // Encode disabled static frame entry in image table + assert_eq!(self.static_jtable.len(), STATIC_FRAME_ENTRY_NUMBER); + + Ok(self + .static_jtable + .iter() + .map(|entry| (F::from(entry.enable as u64), bn_to_field(&entry.encode()))) + .collect()) + }; + + let instruction_handler = |_| { + let mut cells = vec![]; + + cells.push(bn_to_field( + &ImageTableEncoder::Instruction.encode(BigUint::from(0u64)), + )); + + for e in self.itable.entries() { + cells.push(bn_to_field( + &ImageTableEncoder::Instruction.encode(e.encode()), + )); + } + + Ok(cells) + }; + + let br_table_handler = |_| { + let mut cells = vec![]; + + cells.push(bn_to_field( + &ImageTableEncoder::BrTable.encode(BigUint::from(0u64)), + )); + + for e in self.br_table.entries() { + cells.push(bn_to_field(&ImageTableEncoder::BrTable.encode(e.encode()))); + } + + for e in self.elem_table.entries() { + cells.push(bn_to_field(&ImageTableEncoder::BrTable.encode(e.encode()))); + } + + Ok(cells) + }; + + let padding_handler = |start, end| Ok(vec![F::zero(); end - start]); + + let init_memory_entries_handler = |_| { + let mut cells = vec![]; + + cells.push(bn_to_field( + &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), + )); + + let layouter = InitMemoryLayouter { + pages: pages_capability, + }; + + layouter.for_each(|(ltype, offset)| { + if let Some(entry) = self.imtable.try_find(ltype, offset) { + cells.push(bn_to_field::( + &ImageTableEncoder::InitMemory.encode(entry.encode()), + )); + } else if ltype == LocationType::Heap { + let entry = InitMemoryTableEntry { + ltype, + is_mutable: true, + offset, + vtype: VarType::I64, + value: 0, + eid: 0, + }; + + cells.push(bn_to_field::( + &ImageTableEncoder::InitMemory.encode(entry.encode()), + )); + } else { + cells.push(bn_to_field::( + &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), + )); + } + }); + + Ok(cells) + }; + + let mut assigner = ImageTableAssigner::new( + self.itable.entries().len(), + self.br_table.entries().len() + self.elem_table.entries().len(), + pages_capability, + ); + + let layouter = assigner + .exec::<_, Error>( + initialization_state_handler, + static_frame_entries_handler, + instruction_handler, + br_table_handler, + padding_handler, + init_memory_entries_handler, + ) + .unwrap(); + + layouter + } +} + +impl ImageTableLayouter { + pub fn plain(&self) -> Vec { + let mut buf = vec![]; + + buf.append(&mut self.initialization_state.plain()); + buf.append( + &mut self + .static_frame_entries + .clone() + .to_vec() + .into_iter() + .map(|(enable, fid)| vec![enable, fid]) + .collect::>>() + .concat(), + ); + buf.append(&mut self.instructions.clone()); + buf.append(&mut self.br_table_entires.clone()); + buf.append(&mut self.padding_entires.clone()); + buf.append(&mut self.init_memory_entries.clone()); + + buf + } +} diff --git a/crates/zkwasm/src/continuation/slice.rs b/crates/zkwasm/src/continuation/slice.rs index 9ec90bd68..85a6e2117 100644 --- a/crates/zkwasm/src/continuation/slice.rs +++ b/crates/zkwasm/src/continuation/slice.rs @@ -103,6 +103,7 @@ impl Iterator for Slices { let post_image_table = CompilationTable { itable: self.origin_table.compilation_tables.itable.clone(), imtable: updated_init_memory_table, + br_table: self.origin_table.compilation_tables.br_table.clone(), elem_table: self.origin_table.compilation_tables.elem_table.clone(), configure_table: self.origin_table.compilation_tables.configure_table.clone(), static_jtable: self.origin_table.compilation_tables.static_jtable.clone(), diff --git a/crates/zkwasm/src/runtime/wasmi_interpreter.rs b/crates/zkwasm/src/runtime/wasmi_interpreter.rs index 7153f317a..3214b0367 100644 --- a/crates/zkwasm/src/runtime/wasmi_interpreter.rs +++ b/crates/zkwasm/src/runtime/wasmi_interpreter.rs @@ -85,6 +85,7 @@ impl Execution CompilationTable { itable: self.tables.itable.clone(), imtable: updated_init_memory_table, + br_table: self.tables.br_table.clone(), elem_table: self.tables.elem_table.clone(), configure_table: self.tables.configure_table.clone(), static_jtable: self.tables.static_jtable.clone(), @@ -144,20 +145,29 @@ impl WasmiRuntime { iid: 0, }); - if instance.has_start() { - tracer - .clone() - .borrow_mut() - .static_jtable_entries - .push(StaticFrameEntry { + tracer + .clone() + .borrow_mut() + .static_jtable_entries + .push(if instance.has_start() { + StaticFrameEntry { enable: true, frame_id: 0, next_frame_id: 0, callee_fid: 0, // the fid of start function is always 0 fid: idx_of_entry, iid: 0, - }); - } + } + } else { + StaticFrameEntry { + enable: false, + frame_id: 0, + next_frame_id: 0, + callee_fid: 0, + fid: 0, + iid: 0, + } + }); if instance.has_start() { 0 @@ -168,6 +178,7 @@ impl WasmiRuntime { let itable = Arc::new(tracer.borrow().itable.clone()); let imtable = tracer.borrow().imtable.finalized(); + let br_table = Arc::new(itable.create_brtable()); let elem_table = Arc::new(tracer.borrow().elem_table.clone()); let configure_table = Arc::new(tracer.borrow().configure_table.clone()); let static_jtable = Arc::new(tracer.borrow().static_jtable_entries.clone()); @@ -195,6 +206,7 @@ impl WasmiRuntime { tables: CompilationTable { itable, imtable, + br_table, elem_table, configure_table, static_jtable,