fix(vector): Update legacy vector to optimize nth operations on iterators

near-sdk/src/collections/mod.rs

@@ -44,7 +44,7 @@ pub use lookup_map::LookupMap;
 mod lookup_set;
 pub use lookup_set::LookupSet;
 
-mod vector;
+pub mod vector;
 pub use vector::Vector;
 
 mod unordered_map;

near-sdk/src/collections/vector.rs

@@ -1,5 +1,7 @@
 //! A vector implemented on a trie. Unlike standard vector does not support insertion and removal
 //! of an element results in the last element being placed in the empty position.
+use core::ops::Range;
+use std::iter::FusedIterator;
 use std::marker::PhantomData;
 
 use borsh::{BorshDeserialize, BorshSerialize};
@@ -126,11 +128,8 @@ impl<T> Vector<T> {
  }
 
  /// Iterate over raw serialized elements.
- pub fn iter_raw(&self) -> impl Iterator<Item = Vec<u8>> + '_ {
- (0..self.len).map(move |i| {
- let lookup_key = self.index_to_lookup_key(i);
- expect_consistent_state(env::storage_read(&lookup_key))
- })
+ pub fn iter_raw(&self) -> RawIter<T> {
+ RawIter::new(self)
  }
 
  /// Extends vector from the given collection of serialized elements.
@@ -206,8 +205,8 @@ where
  }
 
  /// Iterate over deserialized elements.
- pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
- self.iter_raw().map(|raw_element| Self::deserialize_element(&raw_element))
+ pub fn iter(&self) -> Iter<T> {
+ Iter::new(self)
  }
 
  pub fn to_vec(&self) -> Vec<T> {
@@ -237,6 +236,110 @@ impl<T: std::fmt::Debug + BorshDeserialize> std::fmt::Debug for Vector<T> {
  }
 }
 
+/// An iterator over raw serialized bytes of each element in the [`Vector`].
+pub struct RawIter<'a, T> {
+ vec: &'a Vector<T>,
+ range: Range<u64>,
+}
+
+impl<'a, T> RawIter<'a, T> {
+ fn new(vec: &'a Vector<T>) -> Self {
+ Self { vec, range: Range { start: 0, end: vec.len() } }
+ }
+
+ /// Returns number of elements left to iterate.
+ fn remaining(&self) -> usize {
+ (self.range.end - self.range.start) as usize
+ }
+}
+
+impl<'a, T> Iterator for RawIter<'a, T> {
+ type Item = Vec<u8>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ <Self as Iterator>::nth(self, 0)
+ }
+
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ let remaining = self.remaining();
+ (remaining, Some(remaining))
+ }
+
+ fn count(self) -> usize {
+ self.remaining()
+ }
+
+ fn nth(&mut self, n: usize) -> Option<Self::Item> {
+ let idx = self.range.nth(n)?;
+ self.vec.get_raw(idx)
+ }
+}
+
+impl<'a, T> ExactSizeIterator for RawIter<'a, T> {}
+impl<'a, T> FusedIterator for RawIter<'a, T> {}
+
+impl<'a, T> DoubleEndedIterator for RawIter<'a, T> {
+ fn next_back(&mut self) -> Option<Self::Item> {
+ <Self as DoubleEndedIterator>::nth_back(self, 0)
+ }
+
+ fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+ let idx = self.range.nth_back(n)?;
+ self.vec.get_raw(idx)
+ }
+}
+
+/// An iterator over each element deserialized in the [`Vector`].
+pub struct Iter<'a, T> {
+ inner: RawIter<'a, T>,
+}
+
+impl<'a, T> Iter<'a, T> {
+ fn new(vec: &'a Vector<T>) -> Self {
+ Self { inner: RawIter::new(vec) }
+ }
+}
+
+impl<'a, T> Iterator for Iter<'a, T>
+where
+ T: BorshDeserialize,
+{
+ type Item = T;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ <Self as Iterator>::nth(self, 0)
+ }
+
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ let remaining = self.inner.remaining();
+ (remaining, Some(remaining))
+ }
+
+ fn count(self) -> usize {
+ self.inner.remaining()
+ }
+
+ fn nth(&mut self, n: usize) -> Option<Self::Item> {
+ self.inner.nth(n).map(|raw_element| Vector::deserialize_element(&raw_element))
+ }
+}
+
+impl<'a, T> ExactSizeIterator for Iter<'a, T> where T: BorshDeserialize {}
+impl<'a, T> FusedIterator for Iter<'a, T> where T: BorshDeserialize {}
+
+impl<'a, T> DoubleEndedIterator for Iter<'a, T>
+where
+ T: BorshDeserialize,
+{
+ fn next_back(&mut self) -> Option<Self::Item> {
+ <Self as DoubleEndedIterator>::nth_back(self, 0)
+ }
+
+ fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+ self.inner.nth_back(n).map(|raw_element| Vector::deserialize_element(&raw_element))
+ }
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 #[cfg(test)]
 mod tests {
@@ -393,4 +496,31 @@ mod tests {
  );
  }
  }
+
+ #[test]
+ pub fn iterator_checks() {
+ let mut vec = Vector::new(b"v");
+ let mut baseline = vec![];
+ for i in 0..10 {
+ vec.push(&i);
+ baseline.push(i);
+ }
+
+ let mut vec_iter = vec.iter();
+ let mut bl_iter = baseline.iter();
+ assert_eq!(vec_iter.next(), bl_iter.next().copied());
+ assert_eq!(vec_iter.next_back(), bl_iter.next_back().copied());
+ assert_eq!(vec_iter.nth(3), bl_iter.nth(3).copied());
+ assert_eq!(vec_iter.nth_back(2), bl_iter.nth_back(2).copied());
+
+ // Check to make sure indexing overflow is handled correctly
+ assert!(vec_iter.nth(5).is_none());
+ assert!(bl_iter.nth(5).is_none());
+
+ assert!(vec_iter.next().is_none());
+ assert!(bl_iter.next().is_none());
+
+ // Count check
+ assert_eq!(vec.iter().count(), baseline.len());
+ }
 }