diff --git a/crates/cli/src/file_backend.rs b/crates/cli/src/file_backend.rs index ebf607b5e..ffc80e2c2 100644 --- a/crates/cli/src/file_backend.rs +++ b/crates/cli/src/file_backend.rs @@ -29,6 +29,8 @@ impl Into for &SlicePath { } pub(crate) struct FileBackend { + peeked: Option, + dir_path: PathBuf, name: String, slices: VecDeque, @@ -37,6 +39,8 @@ pub(crate) struct FileBackend { impl FileBackend { pub(crate) fn new(name: String, dir_path: PathBuf) -> Self { FileBackend { + peeked: None, + dir_path, name, slices: VecDeque::new(), @@ -44,14 +48,6 @@ impl FileBackend { } } -impl Iterator for FileBackend { - type Item = Slice; - - fn next(&mut self) -> Option { - self.slices.pop_front().map(|slice| (&slice).into()) - } -} - impl SliceBackend for FileBackend { fn push(&mut self, slice: Slice) { let index = self.slices.len(); @@ -89,18 +85,40 @@ impl SliceBackend for FileBackend { }); } + fn pop(&mut self) -> Option { + match self.peeked.take() { + Some(v) => Some(v), + None => self.slices.pop_front().map(|slice| (&slice).into()), + } + } + + fn first(&mut self) -> Option<&Slice> { + if self.peeked.is_none() { + self.peeked = self.slices.pop_front().map(|slice| (&slice).into()); + } + + self.peeked.as_ref() + } + fn len(&self) -> usize { - self.slices.len() + self.slices.len() + self.peeked.is_some() as usize } fn is_empty(&self) -> bool { - self.slices.is_empty() + self.slices.is_empty() && self.peeked.is_none() } - fn for_each1<'a>(&'a self, f: Box) { + fn for_each<'a>(&'a self, f: Box) { + let mut offset = 0usize; + + if let Some(slice) = self.peeked.as_ref() { + f((offset, slice)); + offset = offset + 1; + } + self.slices.iter().enumerate().for_each(|(index, slice)| { let slice: Slice = slice.into(); - f((index, &slice)) + f((index + offset, &slice)) }) } } diff --git a/crates/specs/src/lib.rs b/crates/specs/src/lib.rs index 3dbc43f72..26330d63d 100644 --- a/crates/specs/src/lib.rs +++ b/crates/specs/src/lib.rs @@ -87,7 +87,7 @@ impl Tables { self.execution_tables .slice_backend - .for_each1(Box::new(|(index, slice)| { + .for_each(Box::new(|(index, slice)| { if DEBUG { let path = dir.join(name_of_event_table_slice(index)); slice.etable.write(&path).unwrap(); diff --git a/crates/specs/src/slice_backend/memory.rs b/crates/specs/src/slice_backend/memory.rs index 33dac8723..27609183e 100644 --- a/crates/specs/src/slice_backend/memory.rs +++ b/crates/specs/src/slice_backend/memory.rs @@ -8,17 +8,17 @@ pub struct InMemoryBackend { slices: VecDeque, } -impl Iterator for InMemoryBackend { - type Item = Slice; +impl SliceBackend for InMemoryBackend { + fn push(&mut self, slice: Slice) { + self.slices.push_back(slice) + } - fn next(&mut self) -> Option { + fn pop(&mut self) -> Option { self.slices.pop_front() } -} -impl SliceBackend for InMemoryBackend { - fn push(&mut self, slice: Slice) { - self.slices.push_back(slice) + fn first(&mut self) -> Option<&Slice> { + self.slices.front() } fn len(&self) -> usize { @@ -29,7 +29,7 @@ impl SliceBackend for InMemoryBackend { self.slices.is_empty() } - fn for_each1<'a>(&'a self, f: Box) { + fn for_each<'a>(&'a self, f: Box) { self.slices.iter().enumerate().for_each(f) } } diff --git a/crates/specs/src/slice_backend/mod.rs b/crates/specs/src/slice_backend/mod.rs index ebdd60cda..8172bf5ff 100644 --- a/crates/specs/src/slice_backend/mod.rs +++ b/crates/specs/src/slice_backend/mod.rs @@ -10,11 +10,11 @@ pub struct Slice { pub external_host_call_table: ExternalHostCallTable, } -pub trait SliceBackend: Iterator { +pub trait SliceBackend { fn push(&mut self, slice: Slice); + fn pop(&mut self) -> Option; + fn first(&mut self) -> Option<&Slice>; fn len(&self) -> usize; fn is_empty(&self) -> bool; - // An renaming for_each to avoid conflict with Iterator trait and - // to avoid consume self. - fn for_each1<'a>(&'a self, f: Box); + fn for_each<'a>(&'a self, f: Box); } diff --git a/crates/zkwasm/src/loader/slice.rs b/crates/zkwasm/src/loader/slice.rs index da6420332..e31496d53 100644 --- a/crates/zkwasm/src/loader/slice.rs +++ b/crates/zkwasm/src/loader/slice.rs @@ -13,7 +13,6 @@ use specs::slice::Slice; use specs::slice_backend::SliceBackend; use specs::state::InitializationState; use specs::Tables; -use std::iter::Peekable; use std::sync::Arc; use crate::circuits::ZkWasmCircuit; @@ -36,7 +35,7 @@ pub struct Slices { imtable: Arc, initialization_state: Arc>, - slices: Peekable>>, + slices: Box, context_input_table: Arc>, context_output_table: Arc>, @@ -73,7 +72,7 @@ impl Slices { imtable: tables.compilation_tables.imtable, initialization_state: tables.compilation_tables.initialization_state, - slices: tables.execution_tables.slice_backend.into_iter().peekable(), + slices: tables.execution_tables.slice_backend, context_input_table: tables.execution_tables.context_input_table.into(), context_output_table: tables.execution_tables.context_output_table.into(), @@ -146,8 +145,7 @@ impl Iterator for Slices { type Item = Result, BuildingCircuitError>; fn next(&mut self) -> Option { - let slice = self.slices.peek(); - if slice.is_none() { + if self.slices.is_empty() { return None; } @@ -155,14 +153,14 @@ impl Iterator for Slices { return Some(self.trivial_slice()); } - let slice = self.slices.next().unwrap(); + let slice = self.slices.pop().unwrap(); let frame_table = slice.frame_table.into(); let external_host_call_table = slice.external_host_call_table; let etable = slice.etable; let post_imtable = Arc::new(self.imtable.update_init_memory_table(&etable)); let post_initialization_state = Arc::new({ - let next_event_entry = if let Some(slice) = self.slices.peek() { + let next_event_entry = if let Some(slice) = self.slices.first() { slice.etable.entries().first().cloned() } else { None @@ -177,7 +175,7 @@ impl Iterator for Slices { let post_inherited_frame_table = self.slices - .peek() + .first() .map_or(Arc::new(InheritedFrameTable::default()), |slice| { let post_inherited_frame_table = slice.frame_table.inherited.clone(); @@ -205,7 +203,7 @@ impl Iterator for Slices { context_input_table: self.context_input_table.clone(), context_output_table: self.context_output_table.clone(), - is_last_slice: self.slices.peek().is_none(), + is_last_slice: self.slices.is_empty(), }; self.imtable = post_imtable;