templar_common/
withdrawal_queue.rs

1use std::num::NonZeroU32;
2
3use near_sdk::{collections::LookupMap, near, AccountId, BorshStorageKey, IntoStorageKey};
4
5use crate::asset::BorrowAssetAmount;
6
7#[derive(Debug)]
8#[near(serializers = [borsh])]
9pub struct QueueNode {
10    account_id: AccountId,
11    amount: BorrowAssetAmount,
12    prev: Option<NonZeroU32>,
13    next: Option<NonZeroU32>,
14}
15
16#[derive(Debug)]
17#[near(serializers = [borsh])]
18pub struct WithdrawalQueue {
19    prefix: Vec<u8>,
20    length: u32,
21    next_queue_node_id: NonZeroU32,
22    queue: LookupMap<NonZeroU32, QueueNode>,
23    queue_head: Option<NonZeroU32>,
24    queue_tail: Option<NonZeroU32>,
25    entries: LookupMap<AccountId, NonZeroU32>,
26}
27
28#[derive(BorshStorageKey)]
29#[near(serializers = [borsh])]
30enum StorageKey {
31    Queue,
32    Entries,
33}
34
35fn inconsistent_state<T>() -> T {
36    crate::panic_with_message("Inconsistent state")
37}
38
39impl WithdrawalQueue {
40    pub fn new(prefix: impl IntoStorageKey) -> Self {
41        let prefix = prefix.into_storage_key();
42        macro_rules! key {
43            ($k:ident) => {
44                [prefix.clone(), StorageKey::$k.into_storage_key()].concat()
45            };
46        }
47        Self {
48            prefix: prefix.clone(),
49            length: 0,
50            next_queue_node_id: NonZeroU32::MIN,
51            queue: LookupMap::new(key!(Queue)),
52            queue_head: None,
53            queue_tail: None,
54            entries: LookupMap::new(key!(Entries)),
55        }
56    }
57
58    #[inline]
59    pub fn len(&self) -> u32 {
60        self.length
61    }
62
63    #[inline]
64    pub fn is_empty(&self) -> bool {
65        self.length == 0
66    }
67
68    pub fn get(&self, account_id: &AccountId) -> Option<BorrowAssetAmount> {
69        self.entries
70            .get(account_id)
71            .and_then(|node_id| self.queue.get(&node_id))
72            .map(|queue_node| queue_node.amount)
73    }
74
75    pub fn contains(&self, account_id: &AccountId) -> bool {
76        self.entries.contains_key(account_id)
77    }
78
79    fn mut_existing_node<T>(
80        &mut self,
81        node_id: NonZeroU32,
82        f: impl FnOnce(&mut QueueNode) -> T,
83    ) -> T {
84        let mut node = self.queue.get(&node_id).unwrap_or_else(inconsistent_state);
85        let r = f(&mut node);
86        self.queue.insert(&node_id, &node);
87        r
88    }
89
90    fn set_existing_node_next(&mut self, node_id: NonZeroU32, next: Option<NonZeroU32>) {
91        let mut node = self.queue.get(&node_id).unwrap_or_else(inconsistent_state);
92        node.next = next;
93        self.queue.insert(&node_id, &node);
94    }
95
96    fn set_existing_node_prev(&mut self, node_id: NonZeroU32, prev: Option<NonZeroU32>) {
97        let mut node = self.queue.get(&node_id).unwrap_or_else(inconsistent_state);
98        node.prev = prev;
99        self.queue.insert(&node_id, &node);
100    }
101
102    pub fn peek(&self) -> Option<(AccountId, BorrowAssetAmount)> {
103        if let Some(node_id) = self.queue_head {
104            let QueueNode {
105                account_id, amount, ..
106            } = self.queue.get(&node_id).unwrap_or_else(inconsistent_state);
107            Some((account_id, amount))
108        } else {
109            None
110        }
111    }
112
113    pub fn mut_head<T>(&mut self, f: impl FnOnce(&mut BorrowAssetAmount) -> T) -> Option<T> {
114        self.queue_head
115            .map(|node_id| self.mut_existing_node(node_id, |n| f(&mut n.amount)))
116    }
117
118    /// Only pops if queue is non-empty.
119    pub fn pop(&mut self) -> Option<(AccountId, BorrowAssetAmount)> {
120        if let Some(node_id) = self.queue_head {
121            let QueueNode {
122                account_id,
123                amount,
124                next,
125                ..
126            } = self
127                .queue
128                .remove(&node_id)
129                .unwrap_or_else(inconsistent_state);
130            self.queue_head = next;
131            if let Some(next_id) = next {
132                self.set_existing_node_prev(next_id, None);
133            } else {
134                self.queue_tail = None;
135            }
136            self.entries.remove(&account_id);
137            self.length -= 1;
138            Some((account_id, amount))
139        } else {
140            None
141        }
142    }
143
144    /// If the queue is locked, accounts can only be removed if they are not
145    /// at the head of the queue.
146    pub fn remove(&mut self, account_id: &AccountId) -> Option<BorrowAssetAmount> {
147        if let Some(node_id) = self.entries.remove(account_id) {
148            let node = self
149                .queue
150                .remove(&node_id)
151                .unwrap_or_else(inconsistent_state);
152
153            if let Some(next_id) = node.next {
154                self.set_existing_node_prev(next_id, node.prev);
155            } else {
156                self.queue_tail = node.prev;
157            }
158
159            if let Some(prev_id) = node.prev {
160                self.set_existing_node_next(prev_id, node.next);
161            } else {
162                self.queue_head = node.next;
163            }
164
165            self.length -= 1;
166
167            Some(node.amount)
168        } else {
169            None
170        }
171    }
172
173    pub fn insert_or_update(&mut self, account_id: &AccountId, amount: BorrowAssetAmount) {
174        if let Some(node_id) = self.entries.get(account_id) {
175            // update existing
176            self.mut_existing_node(node_id, |node| node.amount = amount);
177        } else {
178            // add new
179            let node_id = self.next_queue_node_id;
180            {
181                #![allow(clippy::unwrap_used)]
182                // assume the collection never processes more than u32::MAX items
183                self.next_queue_node_id = self.next_queue_node_id.checked_add(1).unwrap();
184            }
185
186            if let Some(tail_id) = self.queue_tail {
187                self.set_existing_node_next(tail_id, Some(node_id));
188            }
189            let node = QueueNode {
190                account_id: account_id.clone(),
191                amount,
192                prev: self.queue_tail,
193                next: None,
194            };
195            if self.queue_head.is_none() {
196                self.queue_head = Some(node_id);
197            }
198            self.queue_tail = Some(node_id);
199            self.queue.insert(&node_id, &node);
200            self.entries.insert(account_id, &node_id);
201            self.length += 1;
202        }
203    }
204
205    pub fn iter(&self) -> WithdrawalQueueIter {
206        WithdrawalQueueIter {
207            withdrawal_queue: self,
208            next_node_id: self.queue_head,
209        }
210    }
211
212    pub fn get_status(&self) -> WithdrawalQueueStatus {
213        WithdrawalQueueStatus {
214            depth: self
215                .iter()
216                .map(|(_, amount)| u128::from(amount))
217                .sum::<u128>()
218                .into(),
219            length: self.len(),
220        }
221    }
222
223    pub fn get_request_status(&self, account_id: &AccountId) -> Option<WithdrawalRequestStatus> {
224        if !self.contains(account_id) {
225            return None;
226        }
227
228        let mut depth = 0.into();
229        for (index, (current_account, amount)) in self.iter().enumerate() {
230            if &current_account == account_id {
231                return Some(WithdrawalRequestStatus {
232                    #[allow(
233                        clippy::cast_possible_truncation,
234                        reason = "Queue length is u32, so this will never truncate"
235                    )]
236                    index: index as u32,
237                    depth,
238                    amount,
239                });
240            }
241
242            depth += amount;
243        }
244
245        unreachable!()
246    }
247}
248
249impl<'a> IntoIterator for &'a WithdrawalQueue {
250    type IntoIter = WithdrawalQueueIter<'a>;
251    type Item = (AccountId, BorrowAssetAmount);
252
253    fn into_iter(self) -> Self::IntoIter {
254        self.iter()
255    }
256}
257
258pub struct WithdrawalQueueIter<'a> {
259    withdrawal_queue: &'a WithdrawalQueue,
260    next_node_id: Option<NonZeroU32>,
261}
262
263impl Iterator for WithdrawalQueueIter<'_> {
264    type Item = (AccountId, BorrowAssetAmount);
265
266    fn next(&mut self) -> Option<Self::Item> {
267        let next_node_id = self.next_node_id?;
268        let r = self
269            .withdrawal_queue
270            .queue
271            .get(&next_node_id)
272            .unwrap_or_else(inconsistent_state);
273        self.next_node_id = r.next;
274        Some((r.account_id, r.amount))
275    }
276}
277
278/// Status of a single account in the withdrawal queue.
279#[derive(Clone, Debug, PartialEq, Eq)]
280#[near(serializers = [json])]
281pub struct WithdrawalRequestStatus {
282    /// What index is this account in the queue?
283    /// That is, how many other withdrawal requests are ahead of this account
284    /// in the queue?
285    pub index: u32,
286    /// Sum of requested amounts of the requests ahead of this account in the
287    /// queue.
288    pub depth: BorrowAssetAmount,
289    /// The amount that this account has requested to withdraw from the
290    /// contract.
291    pub amount: BorrowAssetAmount,
292}
293
294/// Status of the withdrawal queue.
295#[derive(Clone, Debug, PartialEq, Eq)]
296#[near(serializers = [json])]
297pub struct WithdrawalQueueStatus {
298    /// Sum of all amounts of requests in the queue.
299    pub depth: BorrowAssetAmount,
300    /// Number of requests in the queue.
301    pub length: u32,
302}
303
304/// Return value after executing requests from the withdrawal queue.
305#[derive(Clone, Debug, PartialEq, Eq)]
306#[near(serializers = [json])]
307pub struct WithdrawalQueueExecutionResult {
308    /// What is the total value of the requests that were cleared from the queue?
309    pub depth: BorrowAssetAmount,
310    /// How many requests were cleared from the queue?
311    pub length: u32,
312}
313
314pub mod error {
315    use thiserror::Error;
316
317    #[derive(Error, Debug)]
318    #[error("The withdrawal queue is empty")]
319    pub struct EmptyError;
320}
321
322#[cfg(test)]
323mod tests {
324    use near_sdk::AccountId;
325
326    use super::WithdrawalQueue;
327
328    #[test]
329    fn mut_head() {
330        let mut wq = WithdrawalQueue::new(b"w");
331
332        let alice: AccountId = "alice".parse().unwrap();
333        let bob: AccountId = "bob".parse().unwrap();
334        let charlie: AccountId = "charlie".parse().unwrap();
335
336        wq.insert_or_update(&alice, 1.into());
337        wq.insert_or_update(&bob, 2.into());
338        wq.insert_or_update(&charlie, 3.into());
339
340        wq.mut_head(|a| *a += 10).unwrap();
341        assert_eq!(wq.len(), 3);
342
343        assert_eq!(wq.get(&alice).unwrap(), 11.into());
344        assert_eq!(wq.get(&bob).unwrap(), 2.into());
345        assert_eq!(wq.get(&charlie).unwrap(), 3.into());
346        assert_eq!(wq.remove(&alice).unwrap(), 11.into());
347        assert_eq!(wq.len(), 2);
348
349        wq.mut_head(|a| *a += 20).unwrap();
350        assert_eq!(wq.get(&alice), None);
351        assert_eq!(wq.get(&bob).unwrap(), 22.into());
352        assert_eq!(wq.get(&charlie).unwrap(), 3.into());
353        assert_eq!(wq.remove(&bob).unwrap(), 22.into());
354        assert_eq!(wq.len(), 1);
355
356        wq.mut_head(|a| *a += 30).unwrap();
357        assert_eq!(wq.get(&alice), None);
358        assert_eq!(wq.get(&bob), None);
359        assert_eq!(wq.get(&charlie).unwrap(), 33.into());
360        assert_eq!(wq.remove(&charlie).unwrap(), 33.into());
361        assert_eq!(wq.len(), 0);
362    }
363
364    #[test]
365    fn withdrawal_remove() {
366        let mut wq = WithdrawalQueue::new(b"w");
367
368        let alice: AccountId = "alice".parse().unwrap();
369        let bob: AccountId = "bob".parse().unwrap();
370        let charlie: AccountId = "charlie".parse().unwrap();
371
372        wq.insert_or_update(&alice, 1.into());
373        wq.insert_or_update(&bob, 2.into());
374        wq.insert_or_update(&charlie, 3.into());
375        assert_eq!(wq.len(), 3);
376        assert_eq!(wq.remove(&bob), Some(2.into()));
377        assert_eq!(wq.len(), 2);
378        assert_eq!(wq.remove(&charlie), Some(3.into()));
379        assert_eq!(wq.len(), 1);
380        assert_eq!(wq.remove(&alice), Some(1.into()));
381        assert_eq!(wq.len(), 0);
382    }
383
384    #[test]
385    fn withdrawal_queueing() {
386        let mut wq = WithdrawalQueue::new(b"w");
387
388        let alice: AccountId = "alice".parse().unwrap();
389        let bob: AccountId = "bob".parse().unwrap();
390        let charlie: AccountId = "charlie".parse().unwrap();
391
392        assert_eq!(wq.len(), 0);
393        assert_eq!(wq.peek(), None);
394        wq.insert_or_update(&alice, 1.into());
395        assert_eq!(wq.len(), 1);
396        assert_eq!(wq.peek(), Some((alice.clone(), 1.into())));
397        wq.insert_or_update(&alice, 99.into());
398        assert_eq!(wq.len(), 1);
399        assert_eq!(wq.peek(), Some((alice.clone(), 99.into())));
400        wq.insert_or_update(&bob, 123.into());
401        assert_eq!(wq.len(), 2);
402
403        assert_eq!(wq.pop(), Some((alice.clone(), 99.into())));
404        assert_eq!(wq.len(), 1);
405        assert_eq!(wq.peek(), Some((bob.clone(), 123.into())));
406
407        wq.insert_or_update(&charlie, 8080.into());
408        assert_eq!(wq.len(), 2);
409        assert_eq!(wq.peek(), Some((bob.clone(), 123.into())));
410
411        assert_eq!(wq.pop(), Some((bob.clone(), 123.into())));
412        assert_eq!(wq.len(), 1);
413        assert_eq!(wq.peek(), Some((charlie.clone(), 8080.into())));
414
415        assert_eq!(wq.pop(), Some((charlie.clone(), 8080.into())));
416        assert_eq!(wq.len(), 0);
417        assert_eq!(wq.peek(), None);
418    }
419}
templar_common/withdrawal_queue.rs

templar_common/
withdrawal_queue.rs
