Skip to content

Commit

Permalink
refine image table
Browse files Browse the repository at this point in the history
  • Loading branch information
@junyu0312
junyu0312 committed Dec 12, 2023
1 parent d3d0986 commit 3e234ba
Show file tree
Hide file tree
Showing 13 changed files with 252 additions and 330 deletions.
2 changes: 1 addition & 1 deletion crates/specs/src/brtable.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ pub struct BrTableEntry {
pub dst_pc: u32,
}

#[derive(Debug)]
#[derive(Default, Serialize, Debug, Clone)]
pub struct BrTable(Vec<BrTableEntry>);

impl BrTable {
Expand Down
4 changes: 3 additions & 1 deletion crates/specs/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<InstructionTable>,
pub imtable: InitMemoryTable,
pub br_table: Arc<BrTable>,
pub elem_table: Arc<ElemTable>,
pub configure_table: Arc<ConfigureTable>,
pub static_jtable: Arc<Vec<StaticFrameEntry>>,
Expand Down
2 changes: 1 addition & 1 deletion crates/zkwasm/src/checksum/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Output> {
fn checksum(&self) -> Output;
Expand Down
32 changes: 8 additions & 24 deletions crates/zkwasm/src/circuits/image_table/assign.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<INIT_MEMORY_ENTRIES_OFFSET, STACK_LIMIT, GLOBAL_LIMIT, F>
{
impl<F: FieldExt> ImageTableChip<F> {
pub(crate) fn assign(
self,
layouter: &mut impl Layouter<F>,
image_table_assigner: &mut ImageTableAssigner<
INIT_MEMORY_ENTRIES_OFFSET,
STACK_LIMIT,
GLOBAL_LIMIT,
>,
image_table_assigner: &mut ImageTableAssigner,
image_table: ImageTableLayouter<F>,
permutation_cells: ImageTableLayouter<AssignedCell<F, F>>,
) -> Result<ImageTableLayouter<AssignedCell<F, F>>, Error> {
Expand Down Expand Up @@ -97,8 +87,6 @@ impl<

image_table
.instructions
.as_ref()
.unwrap()
.iter()
.map(|entry| {
let offset = ctx.borrow().offset;
Expand All @@ -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;
Expand Down Expand Up @@ -166,8 +152,6 @@ impl<

image_table
.init_memory_entries
.as_ref()
.unwrap()
.iter()
.map(|entry| {
let offset = ctx.borrow().offset;
Expand Down Expand Up @@ -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,
})
},
)
Expand Down
223 changes: 3 additions & 220 deletions crates/zkwasm/src/circuits/image_table/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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<T: Clone> {
pub initialization_state: InitializationState<T>,
pub static_frame_entries: Vec<(T, T)>,
pub instructions: Option<Vec<T>>,
pub br_table: Option<Vec<T>>,
pub padding: Option<Vec<T>>,
pub init_memory_entries: Option<Vec<T>>,
}

impl<F: FieldExt> ImageTableLayouter<F> {
pub fn plain(&self) -> Vec<F> {
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::<Vec<Vec<_>>>()
.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<F: Clone> {
fn encode_compilation_table_values(&self) -> ImageTableLayouter<F>;
}

impl<F: FieldExt> EncodeCompilationTableValues<F> for CompilationTable {
fn encode_compilation_table_values(&self) -> ImageTableLayouter<F> {
fn msg_of_initialization_state<F: FieldExt>(
initialization_state: &InitializationState<u32>,
) -> InitializationState<F> {
initialization_state.map(|field| F::from(*field as u64))
}

fn msg_of_instruction_table<F: FieldExt>(instruction_table: &InstructionTable) -> Vec<F> {
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<F: FieldExt>(br_table: &BrTable, elem_table: &ElemTable) -> Vec<F> {
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<F: FieldExt>(init_memory_table: &InitMemoryTable) -> Vec<F> {
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::<F>(
&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::<F>(
&ImageTableEncoder::InitMemory.encode(entry.encode()),
));
} else {
cells.push(bn_to_field::<F>(
&ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)),
));
}
});

cells
}

fn msg_of_static_frame_table<F: FieldExt>(
static_frame_table: &Vec<StaticFrameEntry>,
) -> Vec<(F, F)> {
let mut cells = static_frame_table
.into_iter()
.map(|entry| (F::one(), bn_to_field(&entry.encode())))
.collect::<Vec<_>>();

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<F: FieldExt> {
_memory_addr_sel: Column<Fixed>,
Expand All @@ -256,22 +50,11 @@ impl<F: FieldExt> ImageTableConfig<F> {
}

#[derive(Clone)]
pub struct ImageTableChip<
const INIT_MEMORY_ENTRIES_OFFSET: usize,
const STACK_LIMIT: usize,
const GLOBAL_LIMIT: usize,
F: FieldExt,
> {
pub struct ImageTableChip<F: FieldExt> {
config: ImageTableConfig<F>,
}

impl<
const INIT_MEMORY_ENTRIES_OFFSET: usize,
const STACK_LIMIT: usize,
const GLOBAL_LIMIT: usize,
F: FieldExt,
> ImageTableChip<INIT_MEMORY_ENTRIES_OFFSET, STACK_LIMIT, GLOBAL_LIMIT, F>
{
impl<F: FieldExt> ImageTableChip<F> {
pub fn new(config: ImageTableConfig<F>) -> Self {
ImageTableChip { config }
}
Expand Down
8 changes: 7 additions & 1 deletion crates/zkwasm/src/circuits/jtable/assign.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,13 @@ impl<F: FieldExt> JumpTableChip<F> {
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)?;
Expand Down
Loading

0 comments on commit 3e234ba

Please sign in to comment.