1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
/// An "entry" object corresponding to the top element of the stack.
///
/// Existence of this object guarantees that the stack is not empty.
///
/// In a vector, this is an object representing the last element.
pub struct LIFOEntry<'a, C: ?Sized>(&'a mut C);

impl<'a, C: ?Sized + Stack> LIFOEntry<'a, C> {
    /// Creates a new "entry" object from the mutable reference to the container.
    ///
    /// ## Safety
    ///
    /// The stack must not be empty.
    pub unsafe fn new(stack: &'a mut C) -> Self {
        // SAFETY: The stack is not empty, so the call is safe.
        Self(stack)
    }

    /// Pops the LIFO element from the stack.
    pub fn pop_pointee(self) -> C::Item {
        let LIFOEntry(stack) = self;
        // SAFETY: The stack is not empty by the virtue of
        // existence of the LIFOEntry object, so the call is safe.
        unsafe { stack.s_pop_unchecked() }
    }
}

impl<'a, C: ?Sized + Stack> std::ops::Deref for LIFOEntry<'a, C> {
    type Target = C::Item;

    fn deref(&self) -> &Self::Target {
        let LIFOEntry(stack) = self;
        // SAFETY: The stack is not empty, so the call is safe.
        unsafe { stack.lifo_ref_unchecked() }
    }
}

impl<'a, C: ?Sized + Stack> std::ops::DerefMut for LIFOEntry<'a, C> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        let LIFOEntry(stack) = self;
        // SAFETY: The stack is not empty, so the call is safe.
        unsafe { stack.lifo_mut_unchecked() }
    }
}

/// Implementors of this method can be used as a [stack] in DSA terms.
///
/// [stack]: https://www.geeksforgeeks.org/stack-data-structure/
pub trait Stack {
    /// The type of the items stored in the stack.
    type Item;

    /// Returns `true` if the stack is empty.
    fn s_is_empty(&self) -> bool;

    /// Pushes an item to the stack.
    ///
    /// For vector, use [`Vec::push`] instead.
    ///
    /// ## Also see
    ///
    /// * [`Stack::s_push_checked`]
    // s_ prefix prevents name collision with Vec::push
    fn s_push(&mut self, item: Self::Item);

    /// Pushes an item to the stack.
    ///
    /// ## Notes
    ///
    /// For vector, use [`Vec::push`] instead. It is meant primarily for the `heapless::Vec`.
    ///
    /// ## Also see
    ///
    /// * [`Stack::s_push`]
    fn s_push_checked(&mut self, item: Self::Item) -> Option<()> {
        self.s_push(item);
        Some(())
    }

    // we don't create chain push because Extend::extend_one API will be better

    /// Pushes an item to the stack and returns an the "entry" object
    /// corresponding to the pushed element.
    fn lifo_push(&mut self, item: Self::Item) -> LIFOEntry<Self>;

    /// Pops an item from the stack.
    ///
    /// ## Notes
    ///
    /// For vector, use [`Vec::pop`] instead.
    ///
    /// ## Also see
    ///
    /// * [`Stack::s_pop_unchecked`].
    fn s_pop(&mut self) -> Option<Self::Item>;

    /// Pops an item from the stack without checking if the stack is empty.
    ///
    /// ## Safety
    ///
    /// The stack must not be empty.
    ///
    /// ## Also see
    ///
    /// * [`Stack::s_pop`]
    #[inline]
    unsafe fn s_pop_unchecked(&mut self) -> Self::Item {
        self.s_pop().unwrap_unchecked()
    }

    /// Returns an "entry" object corresponding to the top element of the stack.
    ///
    /// ## Notes
    ///
    /// This is useful when the entry might need to be used to be dereferenced to
    /// `&T` and/or `&mut T` and/or be promoted to `T` by popping the item from the stack.
    ///
    /// ## Also see
    ///
    /// * [`Stack::lifo_unchecked`]
    #[inline]
    fn lifo(&mut self) -> Option<LIFOEntry<Self>> {
        if self.s_is_empty() {
            None
        } else {
            Some(unsafe { LIFOEntry::new(self) })
        }
    }

    /// Returns an "entry" object corresponding to the top element of the stack
    /// without checking if the stack is empty.
    ///
    /// ## Safety
    ///
    /// The stack must not be empty.
    ///
    /// ## Also see
    ///
    /// * [`Stack::lifo`]
    #[inline]
    unsafe fn lifo_unchecked(&mut self) -> LIFOEntry<Self> {
        self.lifo().unwrap_unchecked()
    }

    /// Returns a shared reference to the top element of the stack.
    ///
    /// ## Also see
    ///
    /// * [`Stack::lifo_ref_unchecked`]
    /// * [`Stack::lifo`]
    /// * [`Stack::s_pop`]
    /// * [`Stack::lifo_mut`]
    fn lifo_ref(&self) -> Option<&Self::Item>;

    /// Returns a shared reference to the top element of the stack without checking if the stack is empty.
    ///
    /// ## Safety
    ///
    /// The stack must not be empty.
    ///
    /// ## Also see
    ///
    /// * [`Stack::lifo_ref`]
    #[inline]
    unsafe fn lifo_ref_unchecked(&self) -> &Self::Item {
        self.lifo_ref().unwrap_unchecked()
    }

    /// Returns a mutable reference to the top element of the stack.
    ///
    /// ## Also see
    ///
    /// * [`Stack::lifo_ref`]
    /// * [`Stack::lifo_mut_unchecked`]
    /// * [`Stack::lifo`]
    /// * [`Stack::s_pop`]
    fn lifo_mut(&mut self) -> Option<&mut Self::Item>;

    /// Returns a mutable reference to the top element of the stack without checking if the stack is empty.
    ///
    /// ## Safety
    ///
    /// The stack must not be empty.
    #[inline]
    unsafe fn lifo_mut_unchecked(&mut self) -> &mut Self::Item {
        self.lifo_mut().unwrap_unchecked()
    }
}

impl<T> Stack for Vec<T> {
    type Item = T;

    #[inline]
    fn s_is_empty(&self) -> bool {
        self.is_empty()
    }

    #[inline]
    fn s_push(&mut self, item: Self::Item) {
        self.push(item);
    }

    #[inline]
    fn lifo_push(&mut self, item: Self::Item) -> LIFOEntry<Self> {
        self.push(item);
        // We just pushed to the vector, so the vector is not empty.
        unsafe { self.lifo_unchecked() }
    }

    #[inline]
    fn s_pop(&mut self) -> Option<Self::Item> {
        self.pop()
    }

    #[inline]
    fn lifo_ref(&self) -> Option<&Self::Item> {
        self.last()
    }

    #[inline]
    fn lifo_mut(&mut self) -> Option<&mut Self::Item> {
        self.last_mut()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn get_or_insert() {
        let mut stack = vec![1, 2, 3];
        let mut entry = stack.lifo_push(4);
        assert_eq!(*entry, 4);
        *entry = 5;
        assert_eq!(*entry, 5);
        drop(entry);
        assert_eq!(stack, vec![1, 2, 3, 5]);
        let entry = stack.lifo().unwrap();
        assert_eq!(entry.pop_pointee(), 5);
        assert_eq!(stack, vec![1, 2, 3]);
    }
}