Implement iterator on Binary Tree

[d-cap]

Hello I'm trying to implement a Binary Tree using a Vec, I found that I cannot easily implement an iterator for it, following is the code I'm using, am I doing something wrong? I don't understand why self.tree.get_mut(index) cannot outlive the method, should the reference already has the 'a lifetime?

fn main() {
    let mut b = BinaryTreeArray::new(1);


    for d in &mut b.iter_mut() {
        d.data = 2;
    }
}

pub struct BinaryTreeNode<T> {
    pub data: T,
    l: usize,
    r: usize,
}

impl<T> BinaryTreeNode<T> {
    pub fn new(data: T) -> Self {
        Self { data, l: 0, r: 0 }
    }

    #[inline]
    pub fn is_leaf(&self) -> bool {
        self.l == 0 && self.r == 0
    }
}

pub struct BinaryTreeArray<T> {
    nodes: Vec<BinaryTreeNode<T>>,
}

impl<T> BinaryTreeArray<T> {
    pub fn new(data: T) -> Self {
        Self {
            nodes: vec![BinaryTreeNode::new(data)],
        }
    }

    #[inline]
    pub fn get(&self, index: usize) -> Option<&BinaryTreeNode<T>> {
        self.nodes.get(index)
    }

    #[inline]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut BinaryTreeNode<T>> {
        self.nodes.get_mut(index)
    }

    pub fn iter_mut(&mut self) -> BinaryTreeIteratorMut<T> {
        BinaryTreeIteratorMut::new(self)
    }
}

pub struct BinaryTreeIteratorMut<'a, T> {
    tree: &'a mut BinaryTreeArray<T>,
    nodes: Vec<usize>,
}

impl<'a, T> BinaryTreeIteratorMut<'a, T> {
    pub fn new(tree: &'a mut BinaryTreeArray<T>) -> Self {
        Self {
            tree,
            nodes: vec![0],
        }
    }
}

impl<'a, T> Iterator for BinaryTreeIteratorMut<'a, T> {
    type Item = &'a mut BinaryTreeNode<T>;

    fn next(&mut self) -> Option<&'a mut BinaryTreeNode<T>> {
        let Some(index) = self.nodes.pop();
        let some_node: Option<&'a mut BinaryTreeNode<T>> = self.tree.get_mut(index);
        if let Some(node) = some_node {
            if node.r != 0 {
                self.nodes.push(node.r);
            }
            if node.l != 0 {
                self.nodes.push(node.l);
            }
            return Some(node);
        }
        None
    }
}

Thank you

[matklad]

I think this iterator is impossible to implement using only safe code. The reason is that, to do so, we'll need to prove that references are indeed disjoint.

That is, with this iterator it's possible to write code like this:

let nodes: Vec<&mut BinaryTreeNode<_>> = my_tree.iter_mut().collect();

This code materializes unique references to all elements of the tree, and, to be safe, we need to make sure that this elements are indeed distinct. The "distinctness" is a non-trivial property that relies on the fact that nodes indeed form a tree without cycles, because l and r fields are set up in this case.

Often it is possible to prove such disjointedness by leveraging various split_mut methods of slices. For example, here's how a slice iterator can be implemented:

struct MySliceIter<'a, T> {
    slice: &'a mut [T],
}

impl<'a, T> Iterator for MySliceIter<'a, T> {
    type Item = &'a mut T;
    fn next(&mut self) -> Option<&'a mut T> {
        // steal the slice to appease the borrow checker
        let slice = std::mem::replace(&mut self.slice, &mut []);

        // use `split_first_mut` to "prove" that head and tail are disjoint.
        // `split_first_mut` uses `unsafe` inside. 
        let (head, tail) = slice.split_first_mut()?;
        self.slice = tail;

        Some(head)
    }
}

this won't work in this case though, because the nodes are scattered throughout the array.

[d-cap]

Thanks for the answer, if I would like to implement something like that with unsafe rust how can I do it? Do you know some articles I can read to understand the entire process? Should I look at the BTree implementation inside rust? Can that be helpful?

[matklad]

Reading sources of the standard library is tremendously helpful for understanding how unsafe code works. I suggest starting from:

• slice
• Vec
• RawVec

BTreeMap would also be interesting, but it uses parent pointers, so it won't be similar to this case. Off the top of my head, I can't remember any code that does similar "interesting" disjoint indexing, but it should be possible.

Finally, running the unsafe code under miri is helpful in finding things which are definitely unsound.

[d-cap]

Thank you very much!

[Yandros]

Using a streaming Iterator API such as ::streaming-iterator's should fix you issue.

Another way to get around the restriction is to use internal iteration: Playground.

Finally, you can create a iter() version instead, which would allow overlapping references to the vec, and use Interior Mutability to mutate the elements of the tree (this should match other programming languages style, specially with integer types and Cell for the interior mutability):

fn main ()
{
    use ::std::cell::Cell;

    let mut b = BinaryTreeArray::new(Cell::new(1));


    for d in b.iter() {
        d.data.set(2);
    }
}

pub struct BinaryTreeNode<T> {
    pub data: T,
    l: usize,
    r: usize,
}

impl<T> BinaryTreeNode<T> {
    pub fn new(data: T) -> Self {
        Self { data, l: 0, r: 0 }
    }

    #[inline]
    pub fn is_leaf(&self) -> bool {
        self.l == 0 && self.r == 0
    }
}

pub struct BinaryTreeArray<T> {
    nodes: Vec<BinaryTreeNode<T>>,
}

impl<T> BinaryTreeArray<T> {
    pub fn new(data: T) -> Self {
        Self {
            nodes: vec![BinaryTreeNode::new(data)],
        }
    }

    #[inline]
    pub fn get(&self, index: usize) -> Option<&BinaryTreeNode<T>> {
        self.nodes.get(index)
    }

    #[inline]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut BinaryTreeNode<T>> {
        self.nodes.get_mut(index)
    }

    pub fn iter (&'_ self) -> BinaryTreeIterator<'_, T>
    {
        BinaryTreeIteratorMut::new(self)
    }
}

pub struct BinaryTreeIterator<'a, T> {
    tree: &'a BinaryTreeArray<T>,
    nodes: Vec<usize>,
}

impl<'a, T> BinaryTreeIterator<'a, T> {
    pub fn new(tree: &'a BinaryTreeArray<T>) -> Self {
        Self {
            tree,
            nodes: vec![0],
        }
    }
}

impl<'a, T> Iterator for BinaryTreeIteratorMut<'a, T> {
    type Item = &'a BinaryTreeNode<T>;

    fn next (&mut self) -> Option<&'a BinaryTreeNode<T>> {
        let Some(index) = self.nodes.pop();
        let some_node: Option<&'a BinaryTreeNode<T>> = self.tree.get(index);
        if let Some(node) = some_node {
            if node.r != 0 {
                self.nodes.push(node.r);
            }
            if node.l != 0 {
                self.nodes.push(node.l);
            }
            return Some(node);
        }
        None
    }
}

[d-cap]

Thank you for the answer, yes that could probably fix my problem but I would like to avoid Cell and RefCell just for learning purpose.

[d-cap]

Maybe I found a solution, I'm not sure about the correctness, for now seems to work as expected.

impl<T> BinaryTreeArray<T> {
    pub fn iter_mut(&mut self) -> IterMut<T> {
        let ptr = self.nodes.as_mut_ptr();
        let end = unsafe { ptr.add(self.nodes.len()) };
        debug_assert!(!end.is_null());

        IterMut {
            ptr,
            end,
            _marker: marker::PhantomData,
        }
    }
}
pub struct IterMut<'a, T> {
    ptr: *mut BinaryTreeNode<T>,
    end: *mut BinaryTreeNode<T>,
    _marker: marker::PhantomData<&'a BinaryTreeNode<T>>,
}

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = &'a mut BinaryTreeNode<T>;

    fn next(&mut self) -> Option<&'a mut BinaryTreeNode<T>> {
        unsafe {
            if self.ptr == self.end {
                None
            } else {
                let val = self.ptr.as_mut();
                self.ptr = self.ptr.offset(1);
                val
            }
        }
    }
}

Thank you guys for the help

[smarnach]

Doesn't this just iterate the vector in index order, rather than walking the tree? If you only want to iterate the vector, you don't need to write your own iterator at all – you can simply use self.nodes.iter_mut().

[d-cap]

Yes I need now to implement the logic to put the index that is inside the tree node instead of the simple +1 index.