Skip to main content

min_stack/
min_stack.rs

1/* Here's a nice example of using Hegel to catch a tricky bug hidden in a relatively simple data
2 * structure. See if you can spot the problem.
3 *
4 * A "min stack" is a stack extended with the ability to query the minimum element in constant
5 * time. We implement one here by using an extra stack to keep track of minimums.
6 */
7
8#![allow(dead_code)]
9
10use hegel::TestCase;
11use hegel::generators as gs;
12use std::cmp::Ord;
13use std::marker::Copy;
14
15struct MinStack<T> {
16    stack: Vec<T>,
17    minimums: Vec<T>,
18}
19
20impl<T: Copy + Ord> MinStack<T> {
21    fn new() -> Self {
22        MinStack {
23            stack: Vec::new(),
24            minimums: Vec::new(),
25        }
26    }
27
28    fn push(&mut self, element: T) {
29        self.stack.push(element);
30        if self.minimums.is_empty() || self.minimum().is_some_and(|m| element < m) {
31            self.minimums.push(element);
32        }
33    }
34
35    fn pop(&mut self) -> Option<T> {
36        let element = self.stack.pop();
37        if element.is_some_and(|e| self.minimum().unwrap() == e) {
38            self.minimums.pop();
39        }
40        element
41    }
42
43    fn minimum(&self) -> Option<T> {
44        let length = self.minimums.len();
45        if length > 0 {
46            Some(self.minimums[length - 1])
47        } else {
48            None
49        }
50    }
51
52    // The linear time minimum implementation, used for reference.
53    fn reference(&self) -> Option<T> {
54        let mut minimum = None;
55        for element in &self.stack {
56            if minimum.is_none_or(|m| m > *element) {
57                minimum = Some(*element);
58            }
59        }
60        minimum
61    }
62}
63
64struct MinStackTest {
65    stack: MinStack<i32>,
66}
67
68#[hegel::state_machine]
69impl MinStackTest {
70    #[rule]
71    fn push(&mut self, tc: TestCase) {
72        let element = tc.draw(gs::integers::<i32>());
73        self.stack.push(element);
74    }
75
76    #[rule]
77    fn pop(&mut self, tc: TestCase) {
78        let element = self.stack.pop();
79        match element {
80            Some(element) => {
81                tc.note(&format!("pop {}", element));
82            }
83            _ => {
84                tc.note("pop nothing");
85            }
86        }
87    }
88
89    #[invariant]
90    fn minimums_agree(&mut self, _: TestCase) {
91        assert_eq!(self.stack.minimum(), self.stack.reference());
92    }
93}
94
95#[hegel::test]
96fn test_min_stack(tc: TestCase) {
97    let test = MinStackTest {
98        stack: MinStack::new(),
99    };
100    hegel::stateful::run(test, tc);
101}
102
103fn main() {}