1#![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 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 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 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 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 if i == N - 1 {
133 return Push::NoRoom {
134 t,
135 depth: depth + 1,
136 };
137 } else {
138 let mut new_node = NStack::Leaf([
140 Some(t),
141 None,
142 None,
143 None,
144 ]);
145
146 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 = Some((new_node, i + 1));
164 break;
165 }
166 }
167 }
168 }
169 }
170 }
171 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 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 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 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}