From 9c15030ad485bb9e058164c8c97c82349d7d4617 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=82=A2=E3=83=AC=E3=83=BC=E3=82=AF?= Date: Mon, 5 Aug 2024 15:33:28 +0100 Subject: [PATCH] feat: Allow store collection iterators to be cloned (this enables standard Iterator methods like `cycle()`) (#1224) --- near-sdk/src/store/free_list/iter.rs | 1 + near-sdk/src/store/free_list/mod.rs | 2 +- near-sdk/src/store/iterable_map/iter.rs | 1 + near-sdk/src/store/iterable_set/iter.rs | 1 + near-sdk/src/store/tree_map/iter.rs | 3 ++ near-sdk/src/store/unordered_map/iter.rs | 47 ++++++++++++++++++++++++ near-sdk/src/store/unordered_set/iter.rs | 1 + near-sdk/src/store/vec/iter.rs | 2 +- 8 files changed, 56 insertions(+), 2 deletions(-) diff --git a/near-sdk/src/store/free_list/iter.rs b/near-sdk/src/store/free_list/iter.rs index c8208bfa8..844ad73f3 100644 --- a/near-sdk/src/store/free_list/iter.rs +++ b/near-sdk/src/store/free_list/iter.rs @@ -34,6 +34,7 @@ fn decrement_count(count: &mut u32) { } /// An iterator over elements in the storage bucket. This only yields the occupied entries. +#[derive(Clone)] pub struct Iter<'a, T> where T: BorshDeserialize + BorshSerialize, diff --git a/near-sdk/src/store/free_list/mod.rs b/near-sdk/src/store/free_list/mod.rs index 3c5513d04..ccffe758a 100644 --- a/near-sdk/src/store/free_list/mod.rs +++ b/near-sdk/src/store/free_list/mod.rs @@ -34,7 +34,7 @@ where } #[near(inside_nearsdk)] -#[derive(Debug)] +#[derive(Debug, Clone)] enum Slot { /// Represents a filled cell of a value in the collection. Occupied(T), diff --git a/near-sdk/src/store/iterable_map/iter.rs b/near-sdk/src/store/iterable_map/iter.rs index 549e8df2b..8eb784c4b 100644 --- a/near-sdk/src/store/iterable_map/iter.rs +++ b/near-sdk/src/store/iterable_map/iter.rs @@ -37,6 +37,7 @@ where /// An iterator over elements of a [`IterableMap`]. /// /// This `struct` is created by the `iter` method on [`IterableMap`]. +#[derive(Clone)] pub struct Iter<'a, K, V, H> where K: BorshSerialize + Ord + BorshDeserialize, diff --git a/near-sdk/src/store/iterable_set/iter.rs b/near-sdk/src/store/iterable_set/iter.rs index 6e551570b..db45d5460 100644 --- a/near-sdk/src/store/iterable_set/iter.rs +++ b/near-sdk/src/store/iterable_set/iter.rs @@ -24,6 +24,7 @@ where /// See its documentation for more. /// /// [`iter`]: IterableSet::iter +#[derive(Clone)] pub struct Iter<'a, T> where T: BorshSerialize + Ord + BorshDeserialize, diff --git a/near-sdk/src/store/tree_map/iter.rs b/near-sdk/src/store/tree_map/iter.rs index 05778b82a..629cdd3da 100644 --- a/near-sdk/src/store/tree_map/iter.rs +++ b/near-sdk/src/store/tree_map/iter.rs @@ -40,6 +40,7 @@ where /// An iterator over elements of a [`TreeMap`], in sorted order. /// /// This `struct` is created by the `iter` method on [`TreeMap`]. +#[derive(Clone)] pub struct Iter<'a, K, V, H> where K: BorshSerialize + Ord + BorshDeserialize, @@ -408,6 +409,7 @@ where /// An iterator over the keys of a [`TreeMap`], in sorted order. /// /// This `struct` is created by the `keys` method on [`TreeMap`]. +#[derive(Clone)] pub struct Keys<'a, K: 'a> where K: BorshSerialize + BorshDeserialize + Ord, @@ -676,6 +678,7 @@ where /// An iterator over the values of a [`TreeMap`], in order by key. /// /// This `struct` is created by the `values` method on [`TreeMap`]. +#[derive(Clone)] pub struct Values<'a, K, V, H> where K: BorshSerialize + Ord + BorshDeserialize, diff --git a/near-sdk/src/store/unordered_map/iter.rs b/near-sdk/src/store/unordered_map/iter.rs index 42199ea5d..7d28c6c42 100644 --- a/near-sdk/src/store/unordered_map/iter.rs +++ b/near-sdk/src/store/unordered_map/iter.rs @@ -36,6 +36,7 @@ where /// An iterator over elements of a [`UnorderedMap`]. /// /// This `struct` is created by the `iter` method on [`UnorderedMap`]. +#[derive(Clone)] pub struct Iter<'a, K, V, H> where K: BorshSerialize + Ord + BorshDeserialize, @@ -222,6 +223,7 @@ where /// An iterator over the keys of a [`UnorderedMap`]. /// /// This `struct` is created by the `keys` method on [`UnorderedMap`]. +#[derive(Clone)] pub struct Keys<'a, K: 'a> where K: BorshSerialize + BorshDeserialize, @@ -277,6 +279,7 @@ where /// An iterator over the values of a [`UnorderedMap`]. /// /// This `struct` is created by the `values` method on [`UnorderedMap`]. +#[derive(Clone)] pub struct Values<'a, K, V, H> where K: BorshSerialize + Ord + BorshDeserialize, @@ -521,3 +524,47 @@ where Some(self.remove_value(key)) } } + +#[cfg(test)] +mod tests { + use super::*; + use borsh::{BorshDeserialize, BorshSerialize}; + + #[derive(BorshSerialize, BorshDeserialize, Ord, PartialOrd, Eq, PartialEq, Debug, Clone)] + struct Key(i32); + + #[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq)] + struct Value(String); + + #[test] + fn test_unordered_map_iter_clone() { + let mut store = UnorderedMap::new(b'a'); + + store.insert(Key(1), Value("one".to_string())); + store.insert(Key(2), Value("two".to_string())); + store.insert(Key(3), Value("three".to_string())); + + let mut iter = store.iter().cycle(); + + let mut collected = vec![]; + for _ in 0..9 { + if let Some((key, value)) = iter.next() { + collected.push((key.clone(), value.clone())); + } + } + + let expected = vec![ + (Key(1), Value("one".to_string())), + (Key(2), Value("two".to_string())), + (Key(3), Value("three".to_string())), + (Key(1), Value("one".to_string())), + (Key(2), Value("two".to_string())), + (Key(3), Value("three".to_string())), + (Key(1), Value("one".to_string())), + (Key(2), Value("two".to_string())), + (Key(3), Value("three".to_string())), + ]; + + assert_eq!(collected, expected); + } +} diff --git a/near-sdk/src/store/unordered_set/iter.rs b/near-sdk/src/store/unordered_set/iter.rs index 321df6d96..a31dab95d 100644 --- a/near-sdk/src/store/unordered_set/iter.rs +++ b/near-sdk/src/store/unordered_set/iter.rs @@ -24,6 +24,7 @@ where /// See its documentation for more. /// /// [`iter`]: UnorderedSet::iter +#[derive(Clone)] pub struct Iter<'a, T> where T: BorshSerialize + Ord + BorshDeserialize, diff --git a/near-sdk/src/store/vec/iter.rs b/near-sdk/src/store/vec/iter.rs index ba39a1de2..cbb7da783 100644 --- a/near-sdk/src/store/vec/iter.rs +++ b/near-sdk/src/store/vec/iter.rs @@ -5,7 +5,7 @@ use super::{Vector, ERR_INDEX_OUT_OF_BOUNDS}; use crate::env; /// An iterator over references to each element in the stored vector. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Iter<'a, T> where T: BorshSerialize + BorshDeserialize,