templar_common/
chunked_append_only_list.rs

1use borsh::{BorshDeserialize, BorshSerialize};
2use near_sdk::{near, store::Vector, BorshStorageKey, IntoStorageKey};
3
4#[derive(Debug, Clone, Copy, BorshSerialize, BorshStorageKey, PartialEq, Eq, PartialOrd, Ord)]
5enum StorageKey {
6    Inner,
7}
8
9/// Represents an append-only iterable list that stores multiple items per
10/// storage slot to reduce gas cost when reading.
11#[derive(Debug)]
12#[near(serializers = [borsh])]
13pub struct ChunkedAppendOnlyList<T: BorshSerialize + BorshDeserialize, const CHUNK_SIZE: u32> {
14    inner: Vector<Vec<T>>,
15    last_chunk_next_index: u32,
16}
17
18impl<T: BorshSerialize + BorshDeserialize, const CHUNK_SIZE: u32>
19    ChunkedAppendOnlyList<T, CHUNK_SIZE>
20{
21    pub fn new(prefix: impl IntoStorageKey) -> Self {
22        Self {
23            inner: Vector::new(
24                [
25                    prefix.into_storage_key(),
26                    StorageKey::Inner.into_storage_key(),
27                ]
28                .concat(),
29            ),
30            last_chunk_next_index: 0,
31        }
32    }
33
34    pub fn len(&self) -> u32 {
35        let Some(last_index) = self.inner.len().checked_sub(1) else {
36            return 0;
37        };
38        last_index * CHUNK_SIZE
39            + if self.last_chunk_next_index == 0 {
40                CHUNK_SIZE
41            } else {
42                self.last_chunk_next_index
43            }
44    }
45
46    pub fn is_empty(&self) -> bool {
47        self.inner.is_empty()
48    }
49
50    pub fn push(&mut self, item: T) {
51        if self.last_chunk_next_index == 0 {
52            let v = vec![item];
53            self.inner.push(v);
54        } else {
55            let last_inner = self
56                .inner
57                .len()
58                .checked_sub(1)
59                .unwrap_or_else(|| crate::panic_with_message("Inconsistent state: len == 0"));
60            let v = self
61                .inner
62                .get_mut(last_inner)
63                .unwrap_or_else(|| crate::panic_with_message("Inconsistent state: tail dne"));
64            v.push(item);
65        }
66        self.last_chunk_next_index = (self.last_chunk_next_index + 1) % CHUNK_SIZE;
67    }
68
69    pub fn get(&self, index: u32) -> Option<&T> {
70        self.inner
71            .get(index / CHUNK_SIZE)
72            .and_then(|v| v.get((index % CHUNK_SIZE) as usize))
73    }
74
75    pub fn replace_last(&mut self, item: T) {
76        let Some(entry) = self
77            .inner
78            .len()
79            .checked_sub(1)
80            .and_then(|last_index| self.inner.get_mut(last_index))
81            .and_then(|v| v.last_mut())
82        else {
83            crate::panic_with_message("Cannot replace_last in empty list");
84        };
85        *entry = item;
86    }
87
88    pub fn iter(&self) -> Iter<'_, T, CHUNK_SIZE> {
89        Iter {
90            list: self,
91            next_index: 0,
92            until_index: self.len(),
93        }
94    }
95
96    pub fn flush(&mut self) {
97        self.inner.flush();
98    }
99}
100
101impl<'a, T: BorshSerialize + BorshDeserialize, const CHUNK_SIZE: u32> IntoIterator
102    for &'a ChunkedAppendOnlyList<T, CHUNK_SIZE>
103{
104    type Item = &'a T;
105
106    type IntoIter = Iter<'a, T, CHUNK_SIZE>;
107
108    fn into_iter(self) -> Self::IntoIter {
109        self.iter()
110    }
111}
112
113pub struct Iter<'a, T: BorshSerialize + BorshDeserialize, const CHUNK_SIZE: u32> {
114    list: &'a ChunkedAppendOnlyList<T, CHUNK_SIZE>,
115    next_index: u32,
116    until_index: u32,
117}
118
119impl<'a, T: BorshSerialize + BorshDeserialize, const CHUNK_SIZE: u32> Iterator
120    for Iter<'a, T, CHUNK_SIZE>
121{
122    type Item = &'a T;
123
124    fn nth(&mut self, n: usize) -> Option<Self::Item> {
125        if self.until_index <= self.next_index {
126            return None;
127        }
128        #[allow(
129            clippy::unwrap_used,
130            reason = "Assume collection size is never > u32::MAX"
131        )]
132        let n = u32::try_from(n).unwrap();
133        let index = self.next_index + n;
134        let value = self.list.get(index)?;
135        self.next_index = index + 1;
136        Some(value)
137    }
138
139    fn next(&mut self) -> Option<Self::Item> {
140        self.nth(0)
141    }
142}
143
144impl<T: BorshSerialize + BorshDeserialize, const CHUNK_SIZE: u32> DoubleEndedIterator
145    for Iter<'_, T, CHUNK_SIZE>
146{
147    fn next_back(&mut self) -> Option<Self::Item> {
148        if self.until_index <= self.next_index {
149            return None;
150        }
151        let (index, value) = self
152            .until_index
153            .checked_sub(1)
154            .and_then(|index| self.list.get(index).map(|x| (index, x)))?;
155        self.until_index = index;
156        Some(value)
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn basic() {
166        let mut list = ChunkedAppendOnlyList::<_, 47>::new(b"l");
167        assert_eq!(list.len(), 0);
168        assert!(list.is_empty());
169
170        for i in 0..10_000usize {
171            list.push(i);
172            assert_eq!(list.len() as usize, i + 1);
173            assert!(!list.is_empty());
174        }
175
176        let mut count = 0;
177        for (i, v) in list.iter().enumerate() {
178            assert_eq!(i, *v);
179            count += 1;
180        }
181
182        assert_eq!(count, 10_000);
183    }
184
185    #[test]
186    fn replace_last() {
187        let mut list = ChunkedAppendOnlyList::<_, 47>::new(b"l");
188        for i in 0..10_000u32 {
189            list.push(i);
190            list.replace_last(i * 2);
191            assert_eq!(list.len(), i + 1);
192            assert!(!list.is_empty());
193        }
194
195        for i in 0..10_000u32 {
196            let x = list.get(i).unwrap();
197            assert_eq!(*x, i * 2);
198        }
199
200        assert_eq!(list.len(), 10_000);
201    }
202
203    #[test]
204    fn next_back() {
205        let mut list = ChunkedAppendOnlyList::<_, 47>::new(b"l");
206        for i in 0..10_000u32 {
207            list.push(i);
208        }
209
210        let mut it = list.iter();
211
212        let mut i = 10_000;
213        while let Some(x) = it.next_back() {
214            i -= 1;
215            assert_eq!(*x, i);
216        }
217
218        assert_eq!(i, 0);
219    }
220}
templar_common/chunked_append_only_list.rs

templar_common/
chunked_append_only_list.rs
