Skip to content

Commit

Permalink
unify etable assignment
Browse files Browse the repository at this point in the history
  • Loading branch information
junyu0312 committed Oct 30, 2023
1 parent df01238 commit 17d46ac
Show file tree
Hide file tree
Showing 8 changed files with 164 additions and 124 deletions.
54 changes: 9 additions & 45 deletions crates/zkwasm/src/circuits/cell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,11 @@ use crate::circuits::utils::bn_to_field;
use crate::circuits::utils::Context;
use crate::nextn;

#[cfg(feature = "continuation")]
pub(crate) type AllocatedStateCell<F> = AllocatedU32Cell<F>;
#[cfg(not(feature = "continuation"))]
pub(crate) type AllocatedStateCell<F> = AllocatedCommonRangeCell<F>;

#[derive(Debug, Clone, Copy)]
pub(crate) struct AllocatedCell<F: FieldExt> {
pub(crate) col: Column<Advice>,
Expand Down Expand Up @@ -48,11 +53,6 @@ pub(crate) trait CellExpression<F: FieldExt> {
) -> Result<AssignedCell<F, F>, Error> {
self.assign(ctx, if value { F::one() } else { F::zero() })
}
fn assign_constant(
&self,
ctx: &mut Context<'_, F>,
value: F,
) -> Result<AssignedCell<F, F>, Error>;
}

impl<F: FieldExt> CellExpression<F> for AllocatedCell<F> {
Expand All @@ -68,19 +68,6 @@ impl<F: FieldExt> CellExpression<F> for AllocatedCell<F> {
|| Ok(value),
)
}

fn assign_constant(
&self,
ctx: &mut Context<'_, F>,
value: F,
) -> Result<AssignedCell<F, F>, Error> {
ctx.region.assign_advice_from_constant(
|| "assign cell",
self.col,
(ctx.offset as i32 + self.rot) as usize,
value,
)
}
}

#[derive(Debug, Clone, Copy)]
Expand All @@ -105,6 +92,10 @@ impl<F: FieldExt> AllocatedU32Cell<F> {
pub(crate) fn expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
self.u32_cell.expr(meta)
}

pub(crate) fn curr_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
self.u32_cell.expr(meta)
}
}

#[derive(Debug, Clone, Copy)]
Expand Down Expand Up @@ -149,21 +140,6 @@ macro_rules! define_cell {

self.0.assign(ctx, value)
}

fn assign_constant(
&self,
ctx: &mut Context<'_, F>,
value: F,
) -> Result<AssignedCell<F, F>, Error> {
assert!(
value <= $limit,
"assigned value {:?} exceeds the limit {:?}",
value,
$limit
);

self.0.assign_constant(ctx, value)
}
}
};
}
Expand All @@ -188,18 +164,6 @@ impl<F: FieldExt> AllocatedU32Cell<F> {
}
self.u32_cell.assign(ctx, (value as u64).into())
}

pub(crate) fn assign_constant(
&self,
ctx: &mut Context<'_, F>,
value: u32,
) -> Result<AssignedCell<F, F>, Error> {
for i in 0..2 {
self.u16_cells_le[i].assign(ctx, (((value >> (i * 16)) & 0xffffu32) as u64).into())?;
}

self.u32_cell.assign_constant(ctx, (value as u64).into())
}
}

impl<F: FieldExt> AllocatedU64Cell<F> {
Expand Down
43 changes: 37 additions & 6 deletions crates/zkwasm/src/circuits/etable/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,24 @@ impl<F: FieldExt> EventTableCellExpression<F> for AllocatedCell<F> {
}
}

impl<F: FieldExt> EventTableCellExpression<F> for AllocatedU32Cell<F> {
fn next_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
nextn!(
meta,
self.u32_cell.0.col,
self.u32_cell.0.rot + EVENT_TABLE_ENTRY_ROWS as i32
)
}

fn prev_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
nextn!(
meta,
self.u32_cell.0.col,
self.u32_cell.0.rot - EVENT_TABLE_ENTRY_ROWS as i32
)
}
}

macro_rules! impl_cell {
($x: ident) => {
impl<F: FieldExt> EventTableCellExpression<F> for $x<F> {
Expand Down Expand Up @@ -157,11 +175,11 @@ pub(crate) enum EventTableCellType {
const BIT_COLUMNS: usize = 12;
const U8_COLUMNS: usize = 1;
// Should be the multiple of 2.
const U32_CELLS: usize = 2;
const U32_CELLS: usize = if cfg!(feature = "continuation") { 2 } else { 0 };
const U64_CELLS: usize = 5;
const U16_COLUMNS: usize = U64_CELLS + (U32_CELLS / 2);
const COMMON_RANGE_COLUMNS: usize = 7;
const UNLIMITED_COLUMNS: usize = 8;
const UNLIMITED_COLUMNS: usize = if cfg!(feature = "continuation") { 8 } else { 7 };
const MEMORY_TABLE_LOOKUP_COLUMNS: usize = 2;
const JUMP_TABLE_LOOKUP_COLUMNS: usize = 1;

Expand Down Expand Up @@ -231,6 +249,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
* enable(meta),
]
});
meta.enable_equality(u32_cell.0.col);
AllocatedU32Cell {
u16_cells_le,
u32_cell,
Expand Down Expand Up @@ -431,6 +450,18 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
AllocatedCommonRangeCell(self.alloc(&EventTableCellType::CommonRange))
}

pub(crate) fn alloc_state_cell(&mut self) -> AllocatedStateCell<F> {
#[cfg(not(feature = "continuation"))]
{
AllocatedCommonRangeCell(self.alloc(&EventTableCellType::CommonRange))
}

#[cfg(feature = "continuation")]
{
self.alloc_u32_cell()
}
}

pub(crate) fn alloc_u8_cell(&mut self) -> AllocatedU8Cell<F> {
AllocatedU8Cell(self.alloc(&EventTableCellType::U8))
}
Expand All @@ -451,7 +482,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
&mut self,
name: &'static str,
constraint_builder: &mut ConstraintBuilder<F>,
eid: AllocatedU32Cell<F>,
eid: AllocatedStateCell<F>,
location_type: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
offset: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
is_i32: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
Expand Down Expand Up @@ -501,7 +532,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
&mut self,
name: &'static str,
constraint_builder: &mut ConstraintBuilder<F>,
eid: AllocatedU32Cell<F>,
eid: AllocatedStateCell<F>,
location_type: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
offset: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
is_i32: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
Expand Down Expand Up @@ -542,7 +573,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
&mut self,
name: &'static str,
constraint_builder: &mut ConstraintBuilder<F>,
eid: AllocatedU32Cell<F>,
eid: AllocatedStateCell<F>,
location_type: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
offset: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
is_i32: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
Expand Down Expand Up @@ -590,7 +621,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
&mut self,
name: &'static str,
constraint_builder: &mut ConstraintBuilder<F>,
eid: AllocatedU32Cell<F>,
eid: AllocatedStateCell<F>,
location_type: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
offset: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
is_i32: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F> + 'static,
Expand Down
89 changes: 59 additions & 30 deletions crates/zkwasm/src/circuits/etable/assign.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,46 +114,60 @@ impl<F: FieldExt> EventTableChip<F> {
configure_table: &ConfigureTable,
initialization_state: &InitializationState<u32>,
) -> Result<InitializationState<Cell>, Error> {
macro_rules! assign_advice {
#[cfg(feature = "continuation")]
macro_rules! assign_state_advice {
($cell:ident, $value:expr) => {
self.config.common_config.$cell.assign(ctx, $value)?.cell()
};
}

let eid = assign_advice!(eid_cell, initialization_state.eid);
let fid = assign_advice!(fid_cell, F::from(initialization_state.fid as u64));
let iid = assign_advice!(iid_cell, F::from(initialization_state.iid as u64));
let sp = assign_advice!(sp_cell, F::from(initialization_state.sp as u64));
let frame_id = assign_advice!(frame_id_cell, F::from(initialization_state.frame_id as u64));
#[cfg(not(feature = "continuation"))]
macro_rules! assign_state_advice {
($cell:ident, $value:expr) => {
assign_common_range_advice!($cell, $value)
};
}

let host_public_inputs = assign_advice!(
input_index_cell,
F::from(initialization_state.host_public_inputs as u64)
);
let context_in_index = assign_advice!(
macro_rules! assign_common_range_advice {
($cell:ident, $value:expr) => {
self.config
.common_config
.$cell
.assign(ctx, F::from($value as u64))?
.cell()
};
}

let eid = assign_state_advice!(eid_cell, initialization_state.eid);
let fid = assign_common_range_advice!(fid_cell, initialization_state.fid);
let iid = assign_common_range_advice!(iid_cell, initialization_state.iid);
let sp = assign_common_range_advice!(sp_cell, initialization_state.sp);
let frame_id = assign_state_advice!(frame_id_cell, initialization_state.frame_id);

let host_public_inputs =
assign_common_range_advice!(input_index_cell, initialization_state.host_public_inputs);
let context_in_index = assign_common_range_advice!(
context_input_index_cell,
F::from(initialization_state.context_in_index as u64)
initialization_state.context_in_index
);
let context_out_index = assign_advice!(
let context_out_index = assign_common_range_advice!(
context_output_index_cell,
F::from(initialization_state.context_out_index as u64)
initialization_state.context_out_index
);
let external_host_call_call_index = assign_advice!(
let external_host_call_call_index = assign_common_range_advice!(
external_host_call_index_cell,
F::from(initialization_state.external_host_call_call_index as u64)
initialization_state.external_host_call_call_index
);

let initial_memory_pages = assign_advice!(
mpages_cell,
F::from(initialization_state.initial_memory_pages as u64)
);
let maximal_memory_pages = assign_advice!(
let initial_memory_pages =
assign_common_range_advice!(mpages_cell, initialization_state.initial_memory_pages);
let maximal_memory_pages = assign_common_range_advice!(
maximal_memory_pages_cell,
F::from(configure_table.maximal_memory_pages as u64)
configure_table.maximal_memory_pages
);

#[cfg(feature = "continuation")]
let jops = assign_advice!(jops_cell, F::from(initialization_state.jops as u64));
let jops = assign_common_range_advice!(jops_cell, initialization_state.jops);

Ok(InitializationState {
eid,
Expand Down Expand Up @@ -184,6 +198,24 @@ impl<F: FieldExt> EventTableChip<F> {
rest_ops: Vec<(u32, u32)>,
initialization_state: &InitializationState<u32>,
) -> Result<(), Error> {
#[cfg(feature = "continuation")]
macro_rules! assign_state_advice {
($cell:ident, $value:expr) => {
self.config.common_config.$cell.assign(ctx, $value)?.cell()
};
}

#[cfg(not(feature = "continuation"))]
macro_rules! assign_state_advice {
($cell:ident, $value:expr) => {
self.config
.common_config
.$cell
.assign(ctx, F::from($value as u64))?
.cell()
};
}

macro_rules! assign_advice {
($cell:ident, $value:expr) => {
self.config.common_config.$cell.assign(ctx, $value)?
Expand Down Expand Up @@ -292,8 +324,8 @@ impl<F: FieldExt> EventTableChip<F> {
maximal_memory_pages_cell,
F::from(configure_table.maximal_memory_pages as u64)
);
assign_advice!(frame_id_cell, F::from(entry.eentry.last_jump_eid as u64));
assign_advice!(eid_cell, entry.eentry.eid);
assign_state_advice!(frame_id_cell, entry.eentry.last_jump_eid);
assign_state_advice!(eid_cell, entry.eentry.eid);
assign_advice!(fid_cell, F::from(entry.eentry.inst.fid as u64));
assign_advice!(iid_cell, F::from(entry.eentry.inst.iid as u64));
assign_advice!(itable_lookup_cell, bn_to_field(&entry.eentry.inst.encode()));
Expand All @@ -320,14 +352,11 @@ impl<F: FieldExt> EventTableChip<F> {
}

// Assign terminate status
assign_advice!(eid_cell, status.last().unwrap().eid);
assign_state_advice!(eid_cell, status.last().unwrap().eid);
assign_advice!(fid_cell, F::from(status.last().unwrap().fid as u64));
assign_advice!(iid_cell, F::from(status.last().unwrap().iid as u64));
assign_advice!(sp_cell, F::from(status.last().unwrap().sp as u64));
assign_advice!(
frame_id_cell,
F::from(status.last().unwrap().last_jump_eid as u64)
);
assign_state_advice!(frame_id_cell, status.last().unwrap().last_jump_eid);
assign_advice!(
mpages_cell,
F::from(status.last().unwrap().allocated_memory_pages as u64)
Expand Down
Loading

0 comments on commit 17d46ac

Please sign in to comment.