use alloc::vec;
use alloc::vec::Vec;
use core::ops::{Index, Range};
#[derive(Debug)]
pub struct Stack<T: Clone> {
cache: Vec<T>,
popped: Vec<T>,
lengths: Vec<(usize, usize)>,
}
impl<T: Clone> Default for Stack<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Clone> Stack<T> {
pub fn new() -> Self {
Stack {
cache: vec![],
popped: vec![],
lengths: vec![],
}
}
#[allow(dead_code)]
pub fn is_empty(&self) -> bool {
self.cache.is_empty()
}
pub fn peek(&self) -> Option<&T> {
self.cache.last()
}
pub fn push(&mut self, elem: T) {
self.cache.push(elem);
}
pub fn pop(&mut self) -> Option<T> {
let len = self.cache.len();
let popped = self.cache.pop();
if let Some(popped) = &popped {
if let Some((_, remained_count)) = self.lengths.last_mut() {
if len == *remained_count {
*remained_count -= 1;
self.popped.push(popped.clone());
}
}
}
popped
}
pub fn len(&self) -> usize {
self.cache.len()
}
pub fn snapshot(&mut self) {
self.lengths.push((self.cache.len(), self.cache.len()))
}
pub fn clear_snapshot(&mut self) {
if let Some((len, unpopped)) = self.lengths.pop() {
self.popped.truncate(self.popped.len() - (len - unpopped));
}
}
pub fn restore(&mut self) {
match self.lengths.pop() {
Some((len_stack, remained)) => {
if remained < self.cache.len() {
self.cache.truncate(remained);
}
if len_stack > remained {
let rewind_count = len_stack - remained;
let new_len = self.popped.len() - rewind_count;
let recovered_elements = self.popped.drain(new_len..);
self.cache.extend(recovered_elements.rev());
debug_assert_eq!(self.popped.len(), new_len);
}
}
None => {
self.cache.clear();
debug_assert!(self.popped.is_empty());
debug_assert!(self.lengths.is_empty());
}
}
}
}
impl<T: Clone> Index<Range<usize>> for Stack<T> {
type Output = [T];
fn index(&self, range: Range<usize>) -> &[T] {
self.cache.index(range)
}
}
#[cfg(test)]
mod test {
use super::Stack;
#[test]
fn snapshot_with_empty() {
let mut stack = Stack::new();
stack.snapshot();
assert!(stack.is_empty());
stack.push(0);
stack.restore();
assert!(stack.is_empty());
}
#[test]
fn snapshot_twice() {
let mut stack = Stack::new();
stack.push(0);
stack.snapshot();
stack.snapshot();
stack.restore();
stack.restore();
assert_eq!(stack[0..stack.len()], [0]);
}
#[test]
fn restore_without_snapshot() {
let mut stack = Stack::new();
stack.push(0);
stack.restore();
assert_eq!(stack[0..stack.len()], [0; 0]);
}
#[test]
fn snapshot_pop_restore() {
let mut stack = Stack::new();
stack.push(0);
stack.snapshot();
stack.pop();
stack.restore();
assert_eq!(stack[0..stack.len()], [0]);
}
#[test]
fn snapshot_pop_push_restore() {
let mut stack = Stack::new();
stack.push(0);
stack.snapshot();
stack.pop();
stack.push(1);
stack.restore();
assert_eq!(stack[0..stack.len()], [0]);
}
#[test]
fn snapshot_push_pop_restore() {
let mut stack = Stack::new();
stack.push(0);
stack.snapshot();
stack.push(1);
stack.push(2);
stack.pop();
stack.restore();
assert_eq!(stack[0..stack.len()], [0]);
}
#[test]
fn snapshot_push_clear() {
let mut stack = Stack::new();
stack.push(0);
stack.snapshot();
stack.push(1);
stack.clear_snapshot();
assert_eq!(stack[0..stack.len()], [0, 1]);
}
#[test]
fn snapshot_pop_clear() {
let mut stack = Stack::new();
stack.push(0);
stack.push(1);
stack.snapshot();
stack.pop();
stack.clear_snapshot();
assert_eq!(stack[0..stack.len()], [0]);
}
#[test]
fn stack_ops() {
let mut stack = Stack::new();
assert!(stack.is_empty());
assert_eq!(stack.peek(), None);
assert_eq!(stack.pop(), None);
stack.push(0);
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&0));
stack.push(1);
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&1));
assert_eq!(stack.pop(), Some(1));
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&0));
stack.push(2);
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&2));
stack.push(3);
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&3));
stack.snapshot();
assert_eq!(stack.pop(), Some(3));
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&2));
stack.snapshot();
assert_eq!(stack.pop(), Some(2));
assert!(!stack.is_empty());
assert_eq!(stack.peek(), Some(&0));
assert_eq!(stack.pop(), Some(0));
assert!(stack.is_empty());
stack.restore();
assert_eq!(stack.pop(), Some(2));
assert_eq!(stack.pop(), Some(0));
assert_eq!(stack.pop(), None);
stack.restore();
assert_eq!(stack.pop(), Some(3));
assert_eq!(stack.pop(), Some(2));
assert_eq!(stack.pop(), Some(0));
assert_eq!(stack.pop(), None);
}
}