diff --git a/src/log.rs b/src/log.rs index 19541366..03b8f1b3 100644 --- a/src/log.rs +++ b/src/log.rs @@ -56,7 +56,9 @@ //! ``` use crate::{read_u64, safe_write, write_u64, Address, GrowFailed, Memory, Storable}; use std::borrow::Cow; +use std::cell::RefCell; use std::marker::PhantomData; +use std::thread::LocalKey; #[cfg(test)] mod tests; @@ -425,3 +427,65 @@ where self.next() } } + +/// Returns an iterator over entries in the log stored in a thread-local variable. +pub fn iter_thread_local( + local_key: &'static LocalKey>>, +) -> ThreadLocalRefIterator +where + T: Storable, + I: Memory, + D: Memory, +{ + ThreadLocalRefIterator { + log: local_key, + buf: vec![], + pos: 0, + } +} + +pub struct ThreadLocalRefIterator +where + T: Storable + 'static, + I: Memory + 'static, + D: Memory + 'static, +{ + log: &'static LocalKey>>, + buf: Vec, + pos: u64, +} + +impl Iterator for ThreadLocalRefIterator +where + T: Storable, + I: Memory, + D: Memory, +{ + type Item = T; + + fn next(&mut self) -> Option { + self.log.with( + |log| match log.borrow().read_entry(self.pos, &mut self.buf) { + Ok(()) => { + self.pos = self.pos.saturating_add(1); + Some(T::from_bytes(Cow::Borrowed(&self.buf))) + } + Err(NoSuchEntry) => None, + }, + ) + } + + fn size_hint(&self) -> (usize, Option) { + let count = self.log.with(|cell| cell.borrow().len()); + (count.saturating_sub(self.pos) as usize, None) + } + + fn count(self) -> usize { + self.size_hint().0 + } + + fn nth(&mut self, n: usize) -> Option { + self.pos = self.pos.saturating_add(n as u64); + self.next() + } +} diff --git a/src/log/tests.rs b/src/log/tests.rs index 6859b3ee..01ec32d2 100644 --- a/src/log/tests.rs +++ b/src/log/tests.rs @@ -1,6 +1,16 @@ -use crate::log::{InitError, Log, WriteError}; +use crate::log::{iter_thread_local, InitError, Log, WriteError}; use crate::vec_mem::VectorMemory; use crate::{Memory, RestrictedMemory, WASM_PAGE_SIZE}; +use std::cell::RefCell; + +thread_local! { + static STRING_LOG: RefCell> = RefCell::new(new_string_log()); +} + +fn new_string_log() -> Log { + Log::init(VectorMemory::default(), VectorMemory::default()) + .expect("failed to initialize stable log") +} #[test] fn test_log_construct() { @@ -227,3 +237,44 @@ fn test_iter() { assert_eq!(log.iter().skip(4).count(), 0); assert_eq!(log.iter().skip(usize::MAX).count(), 0); } + +#[allow(clippy::iter_nth_zero)] +#[test] +fn test_thread_local_iter() { + let new_iter = || iter_thread_local(&STRING_LOG); + + assert_eq!(new_iter().next(), None); + + STRING_LOG.with(|cell| { + *cell.borrow_mut() = new_string_log(); + let log = cell.borrow(); + log.append(&"apple".to_string()).unwrap(); + log.append(&"banana".to_string()).unwrap(); + log.append(&"cider".to_string()).unwrap(); + }); + + let mut iter = new_iter(); + assert_eq!(iter.size_hint(), (3, None)); + assert_eq!(iter.next(), Some("apple".to_string())); + assert_eq!(iter.size_hint(), (2, None)); + assert_eq!(iter.next(), Some("banana".to_string())); + assert_eq!(iter.size_hint(), (1, None)); + assert_eq!(iter.next(), Some("cider".to_string())); + assert_eq!(iter.size_hint(), (0, None)); + assert_eq!(iter.next(), None); + + assert_eq!(new_iter().nth(0), Some("apple".to_string())); + assert_eq!(new_iter().nth(1), Some("banana".to_string())); + assert_eq!(new_iter().nth(2), Some("cider".to_string())); + assert_eq!(new_iter().nth(3), None); + assert_eq!(new_iter().nth(4), None); + assert_eq!(new_iter().nth(usize::MAX), None); + + assert_eq!(new_iter().count(), 3); + assert_eq!(new_iter().skip(0).count(), 3); + assert_eq!(new_iter().skip(1).count(), 2); + assert_eq!(new_iter().skip(2).count(), 1); + assert_eq!(new_iter().skip(3).count(), 0); + assert_eq!(new_iter().skip(4).count(), 0); + assert_eq!(new_iter().skip(usize::MAX).count(), 0); +}