nstack/
lib.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4//
5// Copyright (c) DUSK NETWORK. All rights reserved.
6
7//! NStack
8//!
9//! A stack data structure with indexed lookup.
10#![no_std]
11
12pub mod annotation;
13
14extern crate alloc;
15use alloc::boxed::Box;
16
17use core::mem;
18
19use microkelvin::{Child, ChildMut, Compound, MutableLeaves};
20use ranno::{Annotated, Annotation};
21
22const N: usize = 4;
23
24type NStackRef<T, A> = Box<NStack<T, A>>;
25
26#[derive(Debug)]
27pub enum NStack<T, A> {
28    Leaf([Option<T>; N]),
29    Node([Option<Annotated<NStackRef<T, A>, A>>; N]),
30}
31
32impl<T, A> NStack<T, A> {
33    /// Creates a new empty NStack
34    pub const fn new() -> Self {
35        NStack::Leaf([None, None, None, None])
36    }
37}
38
39enum Push<T> {
40    Ok,
41    NoRoom { t: T, depth: usize },
42}
43
44enum Pop<T> {
45    Ok(T),
46    Last(T),
47    None,
48}
49
50impl<T, A> Compound<A> for NStack<T, A> {
51    type Leaf = T;
52
53    fn child(&self, index: usize) -> Child<Self, A> {
54        match (index, self) {
55            (0, NStack::Node([Some(a), _, _, _])) => Child::Node(a),
56            (1, NStack::Node([_, Some(b), _, _])) => Child::Node(b),
57            (2, NStack::Node([_, _, Some(c), _])) => Child::Node(c),
58            (3, NStack::Node([_, _, _, Some(d)])) => Child::Node(d),
59            (0, NStack::Leaf([Some(a), _, _, _])) => Child::Leaf(a),
60            (1, NStack::Leaf([_, Some(b), _, _])) => Child::Leaf(b),
61            (2, NStack::Leaf([_, _, Some(c), _])) => Child::Leaf(c),
62            (3, NStack::Leaf([_, _, _, Some(d)])) => Child::Leaf(d),
63            _ => Child::EndOfNode,
64        }
65    }
66
67    fn child_mut(&mut self, index: usize) -> ChildMut<Self, A> {
68        match (index, self) {
69            (0, NStack::Node([Some(a), _, _, _])) => ChildMut::Node(a),
70            (1, NStack::Node([_, Some(b), _, _])) => ChildMut::Node(b),
71            (2, NStack::Node([_, _, Some(c), _])) => ChildMut::Node(c),
72            (3, NStack::Node([_, _, _, Some(d)])) => ChildMut::Node(d),
73            (0, NStack::Leaf([Some(a), _, _, _])) => ChildMut::Leaf(a),
74            (1, NStack::Leaf([_, Some(b), _, _])) => ChildMut::Leaf(b),
75            (2, NStack::Leaf([_, _, Some(c), _])) => ChildMut::Leaf(c),
76            (3, NStack::Leaf([_, _, _, Some(d)])) => ChildMut::Leaf(d),
77            _ => ChildMut::EndOfNode,
78        }
79    }
80}
81
82impl<T, A> NStack<T, A>
83where
84    A: Annotation<Self>,
85{
86    /// Pushes a new element onto the stack
87    pub fn push(&mut self, t: T) {
88        match self._push(t) {
89            Push::Ok => (),
90            Push::NoRoom { t, .. } => {
91                let old_root = mem::take(self);
92
93                let mut new_node = [None, None, None, None];
94                new_node[0] = Some(Annotated::new(Box::new(old_root)));
95
96                *self = NStack::Node(new_node);
97
98                // the first child of our new root will be our old root
99                self.push(t)
100            }
101        }
102    }
103
104    fn _push(&mut self, t: T) -> Push<T> {
105        match self {
106            NStack::Leaf(leaf) => {
107                for mut item in leaf.iter_mut() {
108                    match item {
109                        ref mut empty @ None => {
110                            **empty = Some(t);
111                            return Push::Ok;
112                        }
113                        Some(_) => (),
114                    }
115                }
116                Push::NoRoom { t, depth: 0 }
117            }
118            NStack::Node(node) => {
119                let mut insert_node = None;
120
121                // find last node, searching from reverse
122                for i in 0..N {
123                    let i = N - i - 1;
124
125                    match &mut node[i] {
126                        None => (),
127                        Some(anno) => {
128                            match anno.child_mut()._push(t) {
129                                Push::Ok => return Push::Ok,
130                                Push::NoRoom { t, depth } => {
131                                    // Are we in the last node
132                                    if i == N - 1 {
133                                        return Push::NoRoom {
134                                            t,
135                                            depth: depth + 1,
136                                        };
137                                    } else {
138                                        // create a new node
139                                        let mut new_node = NStack::Leaf([
140                                            Some(t),
141                                            None,
142                                            None,
143                                            None,
144                                        ]);
145
146                                        // give it enough depth
147                                        for _ in 0..depth {
148                                            let old_root = mem::replace(
149                                                &mut new_node,
150                                                NStack::new(),
151                                            );
152                                            new_node = NStack::Node([
153                                                Some(Annotated::new(Box::new(
154                                                    old_root,
155                                                ))),
156                                                None,
157                                                None,
158                                                None,
159                                            ]);
160                                        }
161
162                                        // Insert node
163                                        insert_node = Some((new_node, i + 1));
164                                        break;
165                                    }
166                                }
167                            }
168                        }
169                    }
170                }
171                // break out and insert
172                if let Some((new_node, index)) = insert_node {
173                    node[index] = Some(Annotated::new(Box::new(new_node)));
174                } else {
175                    unreachable!()
176                }
177
178                Push::Ok
179            }
180        }
181    }
182
183    /// Pop an element off the stack.
184    ///
185    /// Returns the popped element, if any.
186    pub fn pop(&mut self) -> Option<T> {
187        match self._pop() {
188            Pop::Ok(t) | Pop::Last(t) => Some(t),
189            Pop::None => None,
190        }
191    }
192
193    fn _pop(&mut self) -> Pop<T> {
194        let mut clear_node = None;
195
196        match self {
197            NStack::Leaf(leaf) => {
198                for i in 0..N {
199                    // reverse
200                    let i = N - i - 1;
201                    if let Some(leaf) = leaf[i].take() {
202                        return if i > 0 {
203                            Pop::Ok(leaf)
204                        } else {
205                            Pop::Last(leaf)
206                        };
207                    }
208                }
209                Pop::None
210            }
211            NStack::Node(node) => {
212                for i in 0..N {
213                    // reverse
214                    let i = N - i - 1;
215                    if let Some(ref mut anno) = node[i] {
216                        match anno.child_mut()._pop() {
217                            Pop::Ok(t) => return Pop::Ok(t),
218                            Pop::Last(t) => {
219                                if i == 0 {
220                                    return Pop::Last(t);
221                                } else {
222                                    clear_node = Some((t, i));
223                                    break;
224                                }
225                            }
226                            Pop::None => return Pop::None,
227                        }
228                    }
229                }
230                if let Some((popped, clear_index)) = clear_node {
231                    node[clear_index] = None;
232                    Pop::Ok(popped)
233                } else {
234                    unreachable!()
235                }
236            }
237        }
238    }
239}
240
241impl<T, A> MutableLeaves for NStack<T, A> {}
242
243impl<T, A> Default for NStack<T, A> {
244    fn default() -> Self {
245        Self::new()
246    }
247}
248
249impl<T, A> Clone for NStack<T, A>
250where
251    T: Clone,
252    A: Annotation<Self>,
253{
254    fn clone(&self) -> Self {
255        match self {
256            NStack::Leaf(anno) => NStack::Leaf(anno.clone()),
257            NStack::Node(node) => NStack::Node(node.clone()),
258        }
259    }
260}