Stack of references in unsafe Rust, but ensuring that the unsafeness does not leak out of the stack?

Disclaimer: this answer originally used traits, and it was a nightmare; Francis Gagne pointed out rightly that using an Option for the tail was a much better alternative, hence the answer was much simplified.

Given the structure of your usage, with the stack in RefStack following the usage of the stack frames, you can simply put elements on the stack frames and build a stack from that.

The main advantage of such an approach is that it is entirely safe. You can review the whole code here, or follow for the blow by blow description which follows.

The key is idea is to build a so-called cons-list.

#[derive(Debug)]
struct Stack<'a, T> {
    element: &'a T,
    tail: Option<&'a Stack<'a, T>>,
}

impl<'a, T> Stack<'a, T> {
    fn new(element: &'a T) -> Self { Stack { element, tail: None } }

    fn top(&self) -> &T { self.element }

    fn get(&self, index: usize) -> Option<&T> {
        if index == 0 {
            Some(self.element)
        } else {
            self.tail.and_then(|tail| tail.get(index - 1))
        }
    }

    fn tail(&self) -> Option<&'a Stack<'a, T>> { self.tail }

    fn push<'b>(&'b self, element: &'b T) -> Stack<'b, T> { Stack { element, tail: Some(self) } }
}

A simple example of usage is:

fn immediate() {
    let (a, b, c) = (0, 1, 2);

    let root = Stack::new(&a);
    let middle = root.push(&b);
    let top = middle.push(&c);
    
    println!("{:?}", top);
}

Which just prints the stack, yielding:

Stack { element: 2, tail: Some(Stack { element: 1, tail: Some(Stack { element: 0, tail: None }) }) }

And a more elaborate recursive version:

fn recursive(n: usize) {
    fn inner(n: usize, stack: &Stack<'_, i32>) {
        if n == 0 {
            print!("{:?}", stack);
            return;
        }

        let element = n as i32;
        let stacked = stack.push(&element);
        inner(n - 1, &stacked);
    }

    if n == 0 {
        println!("()");
        return;
    }

    let element = n as i32;
    let root = Stack::new(&element);
    inner(n - 1, &root);
}

Which prints:

Stack { element: 1, tail: Some(Stack { element: 2, tail: Some(Stack { element: 3, tail: None }) }) }

The one downside is that get performance may not be so good; it has linear complexity. On the other hand, cache-wise sticking to the stack frames is pretty nice. If you mostly access the first few elements, I expect it'll be good enough.

I think storing raw pointer is the way to go. You need a PhantomData to store the lifetime and get proper covariance:

use std::marker::PhantomData;

struct RefStack<'a, T: ?Sized> {
    content: Vec<*const T>,
    _pd: PhantomData<&'a T>,
}

impl<'a, T: ?Sized> RefStack<'a, T> {
    fn new() -> Self {
        RefStack {
            content: Vec::new(),_pd: PhantomData
        }
    }
    fn get(&self, index: usize) -> &'a T {
        unsafe { &*self.content[index] }
    }
    fn len(&self) -> usize {
        self.content.len()
    }
    fn with_element<'t, F: FnOnce(&mut RefStack<'t, T>)>(&mut self, el: &'t T, f: F)
        where 'a: 't,
    {
        self.content.push(el);
        let mut tmp = RefStack {
            content: std::mem::take(&mut self.content),
            _pd: PhantomData,
        };
        f(&mut tmp);
        self.content = tmp.content;
        self.content.pop();
    }
}

(Playground)

The only unsafe code is in converting the pointer back into a reference.

The tricky part is getting the with_element right. I think that the were 'a: 't is implicit, because the whole impl depends on it (but better safe than sorry).

The last problem is how to convert a RefStack<'a, T> into a RefStack<'t, T>. I'm pretty sure I could just std::transmute it. But that would an extra unsafe to pay attention to, and creating a new temporary stack one is quite trivial.

About the 't lifetime

You may think that this 't lifetime is not actually needed, but not adding it may cause subtle unsoundness, as the callback could call get() and get values with a lifetime 'a that is actually longer than the inserted value.

For example this code should not compile. With the 't it correctly fails, but without it it compiles and causes undefined behavior:

fn breaking<'a, 's, 'x>(st: &'s mut RefStack<'a, i32>, v: &'x mut Vec<&'a i32>) {
    v.push(st.get(0));
}
fn main() {
    let mut st = RefStack::<i32>::new();
    let mut y = Vec::new();
    {
        let i = 42;
        st.with_element(&i, |stack| breaking(stack, &mut y));
    }
    println!("{:?}", y);
}

About panic!.

When doing these kind of unsafe things, particularly when you are calling user code, as we are doing now in with_element, we have to consider what would happen if it panics. In the OP code, the last object will not be popped, and when the stack is unwound, any drop implementation could see the now dangling reference. My code is ok in case of panics because, if f(&mut tmp); the dangling references die in the local temporary tmp while self.content is empty.

Disclaimer: A different answer; with a different trade-off.

Compared to my other answer, this answer presents a solution that is:

• unsafe: it's encapsulated, but subtle.
• simpler to use.
• simpler code, likely faster.

The idea is to still use the stack to bind the lifetimes of the references, yet storing all lifetimes in a single Vec for O(1) random access. So we're building a stack on the stack, but not storing the references themselves on the stack. Alright?

The full code is available here.

The stack itself is very easily defined:

struct StackRoot<T: ?Sized>(Vec<*const T>);

struct Stack<'a, T: ?Sized>{
    len: usize,
    stack: &'a mut Vec<*const T>,
}

impl<T: ?Sized> StackRoot<T> {
    fn new() -> Self { Self(vec!()) }

    fn stack(&mut self) -> Stack<'_, T> { Stack { len: 0, stack: &mut self.0 } }
}

The implementation of Stack is trickier, as always when unsafe is involved:

impl<'a, T: ?Sized> Stack<'a, T> {
    fn len(&self) -> usize { self.len }

    fn get(&self, index: usize) -> Option<&'a T> {
        if index < self.len {
            //  Safety:
            //  -   Index is bounds as per above branch.
            //  -   Lifetime of reference is guaranteed to be at least 'a (see push).
            Some(unsafe { &**self.stack.get_unchecked(index) })
        } else {
            None
        }
    }

    fn push<'b>(&'b mut self, element: &'b T) -> Stack<'b, T>
        where
            'a: 'b
    {
        //  Stacks could have been built and forgotten, resulting in `self.stack`
        //  containing references to further elements, so that the newly pushed
        //  element would not be at index `self.len`, as expected.
        //
        //  Note that on top of being functionally important, it's also a safety
        //  requirement: `self` should never be able to access elements that are
        //  not guaranteed to have a lifetime longer than `'a`.
        self.stack.truncate(self.len);

        self.stack.push(element as *const _);
        Stack { len: self.len + 1, stack: &mut *self.stack }
    }
}

impl<'a, T: ?Sized> Drop for Stack<'a, T> {
    fn drop(&mut self) {
        self.stack.truncate(self.len);
    }
}

Do note the unsafe here; the invariant is that 'a parameter is always stricter that the lifetimes of the elements pushed into the stack so far.

By refusing to access elements pushed by other members, we thus guarantee that the lifetime of the returned reference is valid.

It does require a generic definition of do_recursion, however generic lifetime parameters are erased at code generation, so there's no code bloat involved:

fn do_recursion<'a, 'b>(nodes: &[&'a str], stack: &mut Stack<'b, str>) 
    where
        'a: 'b
{
    let tmp: &str = stack.get(stack.len() - 1).expect("Not empty");
    println!("{:?}", tmp);

    if let [head, tail @ ..] = nodes {
        let mut new = stack.push(head);
        do_recursion(tail, &mut new);
    }
}

A simple main to show it off:

fn main() {
    let nodes = ["Hello", ",", "World", "!"];
    let mut root = StackRoot::new();
    let mut stack = root.stack();
    let mut stack = stack.push(nodes[0]);

    do_recursion(&nodes[1..], &mut stack);
}

Resulting in:

"Hello"
","
"World"
"!"

Based on rodrigo's answer, I implemented this slightly simpler version:

struct RefStack<'a, T: ?Sized + 'static> {
    content: Vec<&'a T>,
}

impl<'a, T: ?Sized + 'static> RefStack<'a, T> {
    fn new() -> Self {
        RefStack {
            content: Vec::new(),
        }
    }

    fn get(&self, index: usize) -> &'a T {
        self.content[index]
    }

    fn len(&self) -> usize {
        self.content.len()
    }

    fn with_element<'t, F: >(&mut self, el: &'t T, f: F)
    where
        F: FnOnce(&mut RefStack<'t, T>),
        'a: 't,
    {
        let mut st = RefStack {
            content: std::mem::take(&mut self.content),
        };
        st.content.push(el);
        f(&mut st);
        st.content.pop();
        self.content = unsafe { std::mem::transmute(st.content) };
    }
}

The only difference to rodrigo's solution is that the vector is represented as vector of references instead of pointers, so we don't need the PhantomData and the unsafe code to access an element.

When a new element is pushed to the stack in with_element(), we require that it has a shorter lifetime than the existing elements with the a': t' bound. We then create a new stack with the shorter lifetime, which is possible in safe code since we know the data the references in the vector are pointing to even lives for the longer lifetime 'a. We then push the new element with lifetime 't to the new vector, again in safe code, and only after we removed that element again we move the vector back in it's original place. This requires unsafe code since we are extending the lifetime of the references in the vector from 't to 'a this time. We know this is safe, since the vector is back to its original state, but the compiler doesn't know this.

I feel this version represents the intent better than rodrigo's almost identical version. The type of the vector always is "correct", in that it desribes that the elements are actually references, not raw pointers, and it always assigns the correct lifetime to the vector. And we use unsafe code exactly in the place where something potentially unsafe happens – when extending the lifetime of the references in the vector.