perf: minor optimiazation for assignment

crates/specs/src/slice.rs

@@ -78,6 +78,7 @@ impl Slice {
             .collect::<Vec<Vec<_>>>()
             .concat();
 
+        // Use a set to deduplicate
         let mut set = HashSet::<MemoryTableEntry>::default();
 
         memory_entries.iter().for_each(|entry| {
@@ -109,7 +110,7 @@ impl Slice {
 
         memory_entries.append(&mut set.into_iter().collect());
 
-        memory_entries.sort_by_key(|item| (item.ltype, item.offset, item.eid));
+        memory_entries.sort_unstable_by_key(|item| (item.ltype, item.offset, item.eid));
 
         MTable::new(memory_entries)
     }

crates/zkwasm/src/circuits/utils/table_entry.rs

@@ -1,3 +1,5 @@
+use rayon::iter::IntoParallelRefIterator;
+use rayon::iter::ParallelIterator;
 use serde::Serialize;
 use specs::etable::EventTable;
 use specs::etable::EventTableEntry;
@@ -6,7 +8,7 @@ use specs::mtable::LocationType;
 use specs::mtable::MTable;
 use specs::mtable::MemoryTableEntry;
 use std::cmp::Ordering;
-use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::collections::HashSet;
 use std::env;
 use std::io::Write;
@@ -101,20 +103,19 @@ impl MemoryWritingTable {
 
 impl MemoryWritingTable {
     // (location, offset) |-> Vec<(start_eid, end_eid)>
-    fn build_lookup_mapping(&self) -> BTreeMap<(LocationType, u32), Vec<(u32, u32)>> {
-        let mut mapping = BTreeMap::<_, Vec<(u32, u32)>>::new();
+    fn build_lookup_mapping(&self) -> HashMap<(LocationType, u32), Vec<(u32, u32)>> {
+        let mut mapping = HashMap::new();
 
         for entry in &self.0 {
             let ltype = entry.entry.ltype;
             let offset = entry.entry.offset;
             let start_eid = entry.entry.eid;
             let end_eid = entry.end_eid;
 
-            if let Some(entries) = mapping.get_mut(&(ltype, offset)) {
-                entries.push((start_eid, end_eid));
-            } else {
-                mapping.insert((ltype, offset), vec![(start_eid, end_eid)]);
-            }
+            mapping
+                .entry((ltype, offset))
+                .and_modify(|v: &mut Vec<(u32, u32)>| v.push((start_eid, end_eid)))
+                .or_insert(vec![(start_eid, end_eid)]);
         }
 
         mapping
@@ -190,7 +191,7 @@ impl EventTableWithMemoryInfo {
         EventTableWithMemoryInfo(
             event_table
                 .entries()
-                .iter()
+                .par_iter()
                 .map(|eentry| EventTableEntryWithMemoryInfo {
                     eentry: eentry.clone(),
                     memory_rw_entires: memory_event_of_step(eentry)

crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs

@@ -224,7 +224,7 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
     }
 
     fn synthesize(&self, config: Self::Config, layouter: impl Layouter<F>) -> Result<(), Error> {
-        let assign_timer = start_timer!(|| "Assign");
+        let timer = start_timer!(|| "Prepare assignment");
 
         let rchip = RangeTableChip::new(config.rtable);
         let image_chip = ImageTableChip::new(config.image_table);
@@ -241,16 +241,22 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
             ExternalHostCallChip::new(config.external_host_call_table, config.max_available_rows);
         let context_chip = ContextContHelperTableChip::new(config.context_helper_table);
 
-        let image_table_assigner = ImageTableAssigner::new(
-            // Add one for default lookup value
-            self.slice.itable.len() + 1,
-            self.slice.br_table.entries().len() + self.slice.elem_table.entries().len() + 1,
-            config.circuit_maximal_pages,
+        let image_table_assigner = exec_with_profile!(
+            || "Prepare image table assigner",
+            ImageTableAssigner::new(
+                // Add one for default lookup value
+                self.slice.itable.len() + 1,
+                self.slice.br_table.entries().len() + self.slice.elem_table.entries().len() + 1,
+                config.circuit_maximal_pages,
+            )
         );
 
-        let memory_writing_table: MemoryWritingTable = MemoryWritingTable::from(
-            config.k,
-            self.slice.create_memory_table(memory_event_of_step),
+        let memory_writing_table: MemoryWritingTable = exec_with_profile!(
+            || "Prepare mtable",
+            MemoryWritingTable::from(
+                config.k,
+                self.slice.create_memory_table(memory_event_of_step),
+            )
         );
 
         let etable = exec_with_profile!(
@@ -262,7 +268,9 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
 
         let layouter_cloned = layouter.clone();
         let assigned_cells_cloned = assigned_cells.clone();
+        end_timer!(timer);
 
+        let timer = start_timer!(|| "Assign");
         rayon::scope(move |s| {
             let memory_writing_table = Arc::new(memory_writing_table);
             let etable = Arc::new(etable);
@@ -433,6 +441,7 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
                 });
             });
         });
+        end_timer!(timer);
 
         macro_rules! into_inner {
             ($arc:ident) => {
@@ -451,10 +460,12 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
         into_inner!(pre_image_table_cells);
         into_inner!(post_image_table_cells);
         into_inner!(rest_jops_cell_in_frame_table);
+
         /*
          * Permutation between chips
          *
          */
+        let permutation = start_timer!(|| "permutation");
         layouter_cloned.assign_region(
             || "permutation between tables",
             |region| {
@@ -568,8 +579,7 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
                 Ok(())
             },
         )?;
-
-        end_timer!(assign_timer);
+        end_timer!(permutation);
 
         Ok(())
     }