1use alloc::vec;
11use alloc::vec::Vec;
12use std::ops::{Index, Range};
13
14#[derive(Debug)]
17pub struct Stack<T: Clone> {
18 ops: Vec<StackOp<T>>,
19 cache: Vec<T>,
20 snapshots: Vec<usize>,
21}
22
23impl<T: Clone> Stack<T> {
24 pub fn new() -> Self {
26 Stack {
27 ops: vec![],
28 cache: vec![],
29 snapshots: vec![],
30 }
31 }
32
33 #[allow(dead_code)]
35 pub fn is_empty(&self) -> bool {
36 self.cache.is_empty()
37 }
38
39 pub fn peek(&self) -> Option<&T> {
41 self.cache.last()
42 }
43
44 pub fn push(&mut self, elem: T) {
46 self.ops.push(StackOp::Push(elem.clone()));
47 self.cache.push(elem);
48 }
49
50 pub fn pop(&mut self) -> Option<T> {
52 let popped = self.cache.pop();
53 if let Some(ref val) = popped {
54 self.ops.push(StackOp::Pop(val.clone()));
55 }
56 popped
57 }
58
59 pub fn len(&self) -> usize {
61 self.cache.len()
62 }
63
64 pub fn snapshot(&mut self) {
66 self.snapshots.push(self.ops.len());
67 }
68
69 pub fn clear_snapshot(&mut self) {
71 self.snapshots.pop();
72 }
73
74 pub fn restore(&mut self) {
77 match self.snapshots.pop() {
78 Some(ops_index) => {
79 self.rewind_to(ops_index);
80 self.ops.truncate(ops_index);
81 }
82 None => {
83 self.cache.clear();
84 self.ops.clear();
85 }
86 }
87 }
88
89 fn rewind_to(&mut self, index: usize) {
91 let ops_to_rewind = &self.ops[index..];
92 for op in ops_to_rewind.iter().rev() {
93 match *op {
94 StackOp::Push(_) => {
95 self.cache.pop();
96 }
97 StackOp::Pop(ref elem) => {
98 self.cache.push(elem.clone());
99 }
100 }
101 }
102 }
103}
104
105impl<T: Clone> Index<Range<usize>> for Stack<T> {
106 type Output = [T];
107
108 fn index(&self, range: Range<usize>) -> &[T] {
109 self.cache.index(range)
110 }
111}
112
113#[derive(Debug)]
114enum StackOp<T> {
115 Push(T),
116 Pop(T),
117}
118
119#[cfg(test)]
120mod test {
121 use super::Stack;
122
123 #[test]
124 fn snapshot_with_empty() {
125 let mut stack = Stack::new();
126
127 stack.snapshot();
128 assert!(stack.is_empty());
130 stack.push(0);
132 stack.restore();
133 assert!(stack.is_empty());
134 }
135
136 #[test]
137 fn snapshot_twice() {
138 let mut stack = Stack::new();
139
140 stack.push(0);
141
142 stack.snapshot();
143 stack.snapshot();
144 stack.restore();
145 stack.restore();
146
147 assert_eq!(stack[0..stack.len()], [0]);
148 }
149
150 #[test]
151 fn stack_ops() {
152 let mut stack = Stack::new();
153
154 assert!(stack.is_empty());
156 assert_eq!(stack.peek(), None);
157 assert_eq!(stack.pop(), None);
158
159 stack.push(0);
161 assert!(!stack.is_empty());
162 assert_eq!(stack.peek(), Some(&0));
163
164 stack.push(1);
166 assert!(!stack.is_empty());
167 assert_eq!(stack.peek(), Some(&1));
168
169 assert_eq!(stack.pop(), Some(1));
171 assert!(!stack.is_empty());
172 assert_eq!(stack.peek(), Some(&0));
173
174 stack.push(2);
176 assert!(!stack.is_empty());
177 assert_eq!(stack.peek(), Some(&2));
178
179 stack.push(3);
181 assert!(!stack.is_empty());
182 assert_eq!(stack.peek(), Some(&3));
183
184 stack.snapshot();
187
188 assert_eq!(stack.pop(), Some(3));
190 assert!(!stack.is_empty());
191 assert_eq!(stack.peek(), Some(&2));
192
193 stack.snapshot();
196
197 assert_eq!(stack.pop(), Some(2));
199 assert!(!stack.is_empty());
200 assert_eq!(stack.peek(), Some(&0));
201
202 assert_eq!(stack.pop(), Some(0));
204 assert!(stack.is_empty());
205
206 stack.restore();
209 assert_eq!(stack.pop(), Some(2));
210 assert_eq!(stack.pop(), Some(0));
211 assert_eq!(stack.pop(), None);
212
213 stack.restore();
216 assert_eq!(stack.pop(), Some(3));
217 assert_eq!(stack.pop(), Some(2));
218 assert_eq!(stack.pop(), Some(0));
219 assert_eq!(stack.pop(), None);
220 }
221}