chore: extract initialization state (#198)

* feat: preparation for continuation

* fix comments

crates/cli/Cargo.toml

@@ -25,4 +25,5 @@ wasmi.workspace = true
 
 [features]
 default = []
-cuda = ["delphinus-zkwasm/cuda"]
+cuda = ["delphinus-zkwasm/cuda"]
+continuation = ["delphinus-zkwasm/continuation", "specs/continuation"]

crates/specs/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 [dependencies]
 lazy_static = "1.4.0"
 num-bigint = { version = "0.4", features = ["rand"] }
-serde = { version = "1.0", features = ["derive"] }
+serde = { version = "1.0", features = ["derive", "rc"] }
 serde_json = "1.0"
 strum = "0.24.1"
 strum_macros = "0.24.1"
@@ -19,3 +19,4 @@ rayon.workspace = true
 [features]
 default = []
 cuda = ["halo2_proofs/cuda"]
+continuation = []

crates/specs/src/encode/init_memory_table.rs

@@ -12,20 +12,23 @@ pub fn encode_init_memory_table_entry<T: FromBn>(
     is_mutable: T,
     start_offset: T,
     end_offset: T,
+    eid: T,
     value: T,
 ) -> T {
-    const LTYPE_SHIFT: u32 = IS_MUTABLE_SHIFT + COMMON_RANGE_OFFSET;
+    const LTYPE_SHIFT: u32 = IS_MUTABLE_SHIFT + 1;
     const IS_MUTABLE_SHIFT: u32 = START_OFFSET_SHIFT + COMMON_RANGE_OFFSET;
-    const START_OFFSET_SHIFT: u32 = END_OFFSET_SHIFT + 64;
-    const END_OFFSET_SHIFT: u32 = VALUE_SHIFT + 64;
+    const START_OFFSET_SHIFT: u32 = END_OFFSET_SHIFT + COMMON_RANGE_OFFSET;
+    const END_OFFSET_SHIFT: u32 = EID_OFFSET_SHIFT + 32;
+    const EID_OFFSET_SHIFT: u32 = VALUE_SHIFT + 64;
     const VALUE_SHIFT: u32 = 0;
 
-    assert!(LTYPE_SHIFT + COMMON_RANGE_OFFSET <= INIT_MEMORY_ENCODE_BOUNDARY);
+    assert!(LTYPE_SHIFT + 8 <= INIT_MEMORY_ENCODE_BOUNDARY);
 
     ltype * T::from_bn(&(1u64.to_biguint().unwrap() << LTYPE_SHIFT))
         + is_mutable * T::from_bn(&(1u64.to_biguint().unwrap() << IS_MUTABLE_SHIFT))
         + start_offset * T::from_bn(&(1u64.to_biguint().unwrap() << START_OFFSET_SHIFT))
         + end_offset * T::from_bn(&(1u64.to_biguint().unwrap() << END_OFFSET_SHIFT))
+        + eid * T::from_bn(&(1u64.to_biguint().unwrap() << EID_OFFSET_SHIFT))
         + value
 }
 
@@ -36,6 +39,7 @@ impl InitMemoryTableEntry {
             BigUint::from(self.is_mutable as u32),
             BigUint::from(self.start_offset),
             BigUint::from(self.end_offset),
+            BigUint::from(self.eid),
             self.value.to_biguint().unwrap(),
         )
     }

crates/specs/src/imtable.rs

@@ -1,4 +1,5 @@
 use std::cmp::Ordering;
+use std::collections::BTreeMap;
 
 use crate::mtable::LocationType;
 use crate::mtable::VarType;
@@ -13,53 +14,97 @@ pub struct InitMemoryTableEntry {
     pub vtype: VarType,
     /// convert from [u8; 8] via u64::from_le_bytes
     pub value: u64,
+    pub eid: u32,
 }
 
 #[derive(Serialize, Default, Debug, Clone)]
 pub struct InitMemoryTable {
     entries: Vec<InitMemoryTableEntry>,
-    sorted_global_init_entries: Vec<InitMemoryTableEntry>,
+    sorted_global_init_entries: BTreeMap<u32, InitMemoryTableEntry>,
+    sorted_stack_init_entries: BTreeMap<u32, InitMemoryTableEntry>,
     sorted_heap_init_entries: Vec<InitMemoryTableEntry>,
 }
 
 impl InitMemoryTable {
-    pub fn new(entries: Vec<InitMemoryTableEntry>, k: u32) -> Self {
-        let mut imtable = Self {
-            entries: entries
-                .into_iter()
-                .map(|entry| InitMemoryTableEntry {
-                    ltype: entry.ltype,
-                    is_mutable: entry.is_mutable,
-                    start_offset: entry.start_offset,
-                    end_offset: if entry.end_offset == u32::MAX {
-                        (1u32 << (k - 1)) - 1
+    pub fn new(mut entries: Vec<InitMemoryTableEntry>) -> Self {
+        fn sort(entries: &mut Vec<InitMemoryTableEntry>) {
+            entries.sort_by_key(|item| (item.ltype, item.start_offset));
+        }
+
+        fn merge(entries: Vec<InitMemoryTableEntry>) -> Vec<InitMemoryTableEntry> {
+            let mut merged_entries: Vec<_> = entries
+                .iter()
+                .filter(|entry| entry.ltype != LocationType::Heap)
+                .map(|entry| entry.clone())
+                .collect();
+
+            let heap_initial: Vec<_> = entries
+                .iter()
+                .filter(|entry| entry.ltype == LocationType::Heap)
+                .collect();
+
+            if !heap_initial.is_empty() {
+                let mut scan = 0;
+                let mut cursor = scan + 1;
+                while scan < heap_initial.len() && cursor < heap_initial.len() {
+                    if heap_initial[scan].value == heap_initial[cursor].value
+                        && heap_initial[scan].eid == heap_initial[cursor].eid
+                    {
+                        cursor += 1;
                     } else {
-                        entry.end_offset
-                    },
-                    vtype: entry.vtype,
-                    value: entry.value,
-                })
-                .collect(),
-            sorted_global_init_entries: vec![],
-            sorted_heap_init_entries: vec![],
-        };
-        imtable.sort();
-        imtable.merge();
-
-        imtable.sorted_heap_init_entries = imtable
-            .entries
+                        merged_entries.push(InitMemoryTableEntry {
+                            ltype: LocationType::Heap,
+                            is_mutable: true,
+                            start_offset: heap_initial[scan].start_offset,
+                            end_offset: heap_initial[cursor - 1].end_offset,
+                            vtype: VarType::I64,
+                            value: heap_initial[scan].value,
+                            eid: heap_initial[scan].eid,
+                        });
+
+                        scan = cursor;
+                        cursor = scan + 1;
+                    }
+                }
+                merged_entries.push(InitMemoryTableEntry {
+                    ltype: LocationType::Heap,
+                    is_mutable: true,
+                    start_offset: heap_initial[scan].start_offset,
+                    end_offset: heap_initial[cursor - 1].end_offset,
+                    vtype: VarType::I64,
+                    value: heap_initial[scan].value,
+                    eid: heap_initial[scan].eid,
+                });
+            }
+
+            merged_entries
+        }
+
+        sort(&mut entries);
+        let entries = merge(entries);
+
+        let sorted_heap_init_entries = entries
             .iter()
             .filter(|entry| entry.ltype == LocationType::Heap)
             .map(|entry| entry.clone())
             .collect();
-        imtable.sorted_global_init_entries = imtable
-            .entries
+        let sorted_global_init_entries = entries
             .iter()
             .filter(|entry| entry.ltype == LocationType::Global)
-            .map(|entry| entry.clone())
+            .map(|entry| (entry.start_offset, entry.clone()))
+            .collect();
+        let sorted_stack_init_entries = entries
+            .iter()
+            .filter(|entry| entry.ltype == LocationType::Stack)
+            .map(|entry| (entry.start_offset, entry.clone()))
             .collect();
 
-        imtable
+        InitMemoryTable {
+            entries,
+            sorted_global_init_entries,
+            sorted_stack_init_entries,
+            sorted_heap_init_entries,
+        }
     }
 
     pub fn entries(&self) -> &Vec<InitMemoryTableEntry> {
@@ -70,7 +115,7 @@ impl InitMemoryTable {
         serde_json::to_string(&self.entries).unwrap()
     }
 
-    pub fn try_find(&self, ltype: LocationType, offset: u32) -> Option<(u32, u32, u64)> {
+    pub fn try_find(&self, ltype: LocationType, offset: u32) -> Option<(u32, u32, u32, u64)> {
         match ltype {
             LocationType::Heap => {
                 let idx = self
@@ -89,72 +134,23 @@ impl InitMemoryTable {
                 return Some((
                     self.sorted_heap_init_entries[idx].start_offset,
                     self.sorted_heap_init_entries[idx].end_offset,
+                    self.sorted_heap_init_entries[idx].eid,
                     self.sorted_heap_init_entries[idx].value,
                 ));
             }
             LocationType::Global => {
-                for entry in self.sorted_global_init_entries.iter() {
-                    if entry.start_offset == offset {
-                        return Some((offset, offset, entry.value));
-                    }
-                }
+                return self
+                    .sorted_global_init_entries
+                    .get(&offset)
+                    .map(|entry| (entry.start_offset, entry.end_offset, entry.eid, entry.value));
             }
-            LocationType::Stack => unreachable!(),
-        }
-
-        None
-    }
-
-    fn sort(&mut self) {
-        self.entries
-            .sort_by_key(|item| (item.ltype, item.start_offset))
-    }
-
-    fn merge(&mut self) {
-        let mut merged_entries: Vec<_> = self
-            .entries()
-            .iter()
-            .filter(|entry| entry.ltype == LocationType::Global)
-            .map(|entry| entry.clone())
-            .collect();
-
-        let heap_initial: Vec<_> = self
-            .entries()
-            .iter()
-            .filter(|entry| entry.ltype == LocationType::Heap)
-            .collect();
-
-        if !heap_initial.is_empty() {
-            let mut scan = 0;
-            let mut cursor = scan + 1;
-            while scan < heap_initial.len() && cursor < heap_initial.len() {
-                if heap_initial[scan].value == heap_initial[cursor].value {
-                    cursor += 1;
-                } else {
-                    merged_entries.push(InitMemoryTableEntry {
-                        ltype: LocationType::Heap,
-                        is_mutable: true,
-                        start_offset: heap_initial[scan].start_offset,
-                        end_offset: heap_initial[cursor - 1].end_offset,
-                        vtype: VarType::I64,
-                        value: heap_initial[scan].value,
-                    });
-
-                    scan = cursor;
-                    cursor = scan + 1;
-                }
+            LocationType::Stack => {
+                return self
+                    .sorted_stack_init_entries
+                    .get(&offset)
+                    .map(|entry| (entry.start_offset, entry.end_offset, entry.eid, entry.value));
             }
-            merged_entries.push(InitMemoryTableEntry {
-                ltype: LocationType::Heap,
-                is_mutable: true,
-                start_offset: heap_initial[scan].start_offset,
-                end_offset: heap_initial[cursor - 1].end_offset,
-                vtype: VarType::I64,
-                value: heap_initial[scan].value,
-            });
         }
-
-        self.entries = merged_entries;
     }
 
     pub fn filter(&self, ltype: LocationType) -> Vec<&InitMemoryTableEntry> {

crates/specs/src/itable.rs

@@ -296,12 +296,12 @@ pub enum Opcode {
 
 impl Opcode {
     pub fn mops(&self) -> u64 {
-        let opcode_class: OpcodeClass = self.clone().into();
+        let opcode_class: OpcodeClass = self.into();
         opcode_class.mops()
     }
 
     pub fn jops(&self) -> u64 {
-        let opcode_class: OpcodeClass = self.clone().into();
+        let opcode_class: OpcodeClass = self.into();
         opcode_class.jops()
     }
 
@@ -414,7 +414,7 @@ impl Into<BigUint> for Opcode {
             Opcode::InternalHostCall {
                 op_index_in_plugin, ..
             } => {
-                let opcode_class_plain: OpcodeClassPlain = self.into();
+                let opcode_class_plain: OpcodeClassPlain = (&self).into();
 
                 (BigUint::from(opcode_class_plain.0) << OPCODE_CLASS_SHIFT)
                     + (BigUint::from(op_index_in_plugin as u64))
@@ -537,7 +537,7 @@ impl Into<BigUint> for Opcode {
     }
 }
 
-impl Into<OpcodeClass> for Opcode {
+impl Into<OpcodeClass> for &Opcode {
     fn into(self) -> OpcodeClass {
         match self {
             Opcode::LocalGet { .. } => OpcodeClass::LocalGet,
@@ -573,12 +573,12 @@ impl Into<OpcodeClass> for Opcode {
     }
 }
 
-impl Into<OpcodeClassPlain> for Opcode {
+impl Into<OpcodeClassPlain> for &Opcode {
     fn into(self) -> OpcodeClassPlain {
-        let class: OpcodeClass = self.clone().into();
+        let class: OpcodeClass = self.into();
 
         if let Opcode::InternalHostCall { plugin, .. } = self {
-            OpcodeClassPlain(class as usize + plugin as usize)
+            OpcodeClassPlain(class as usize + (*plugin) as usize)
         } else {
             OpcodeClassPlain(class as usize)
         }
@@ -645,7 +645,7 @@ impl InstructionTable {
         let mut opcodeclass: HashSet<OpcodeClassPlain> = HashSet::new();
 
         self.entries().iter().for_each(|entry| {
-            opcodeclass.insert(entry.opcode.clone().into());
+            opcodeclass.insert((&entry.opcode).into());
         });
 
         opcodeclass