Is it a safe implementation of mutable iterator for a tree?

code review
1

Hi, I wrote a DFS mutable iterator for trees. To satisfy the borrow checker, I transmute the lifetime of &mut Node<T> in Drop. Is there any potential safety issue? Glad to see any ideas.
The code is on playground and also here:

use std::{
    collections::VecDeque,
    ops::{Deref, DerefMut},
};

/// Tree node containing data and child nodes
#[derive(Debug)]
struct Node<T> {
    data: T,
    children: Vec<Node<T>>,
}
impl<T> Node<T> {
    /// Creates a depth-first iterator over the node and its descendants
    fn iter_mut(&mut self) -> NodeMutIter<'_, T> {
        NodeMutIter {
            nodes: VecDeque::from([self]),
        }
    }
}

/// Guard that provides safe access to a node during iteration
/// Uses lifetime parameters to ensure iterator outlives the guard
struct NodeGuard<'a: 'b, 'b, T> {
    iter: &'b mut NodeMutIter<'a, T>,
    node: &'a mut Node<T>,
}
impl<'a, 'b, T> Drop for NodeGuard<'a, 'b, T> {
    /// When guard is dropped, add all children to the front of the queue
    /// This implements depth-first traversal
    fn drop(&mut self) {
        for n in self.node.children.iter_mut().rev() {
            // Lifetime transmutation to match the iterator's lifetime
            let n: &'a mut Node<T> = unsafe { &mut *(&mut *n as *mut _) };
            self.iter.nodes.push_front(n);
        }
    }
}
impl<'a, 'b, T> Deref for NodeGuard<'a, 'b, T> {
    type Target = Node<T>;
    fn deref(&self) -> &Self::Target {
        &self.node
    }
}
impl<'a, 'b, T> DerefMut for NodeGuard<'a, 'b, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.node
    }
}

/// Iterator that yields mutable references to nodes
#[derive(Debug)]
struct NodeMutIter<'a, T> {
    nodes: VecDeque<&'a mut Node<T>>,
}

/// Generic Associated Types (GAT) iterator trait
/// Allows returning references with lifetimes tied to the iterator
trait LendingIterator {
    type Item<'a>
    where
        Self: 'a;
    fn next<'a>(&'a mut self) -> Option<Self::Item<'a>>;
}

impl<'a, T> LendingIterator for NodeMutIter<'a, T> {
    /// Return type is a NodeGuard that borrows from both the node and iterator
    type Item<'b>
        = NodeGuard<'a, 'b, T>
    where
        Self: 'b;
    
    /// Returns the next node wrapped in a guard
    /// The borrow checker prevents getting multiple guards simultaneously
    fn next<'b>(&'b mut self) -> Option<Self::Item<'b>> {
        if let Some(node) = self.nodes.pop_front() {
            return Some(NodeGuard { iter: self, node });
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn basic_test() {
        let mut node = Node {
            data: 0,
            children: vec![
                Node {
                    data: 1,
                    children: vec![Node {
                        data: 2,
                        children: vec![],
                    }],
                },
                Node {
                    data: 3,
                    children: vec![],
                },
            ],
        };
        let mut iter = node.iter_mut();
        while let Some(mut guard) = iter.next() {
            guard.data += 1;
            println!("{}", guard.data);
        }
    }
    #[test]
    fn test_move_iter() {
        let mut node = Node {
            data: 0,
            children: vec![Node {
                data: 1,
                children: vec![],
            }],
        };
        let mut iter = node.iter_mut();
        let guard = iter.next().unwrap();
        drop(guard); // if comment this line, the following code will not compile
        let iter_moved = Box::new(iter);
        println!("iter_moved: {:?}", iter_moved);
    }
    #[test]
    fn test_double_guard() {
        let mut node = Node {
            data: 0,
            children: vec![],
        };
        let mut iter = node.iter_mut();
        let guard1 = iter.next().unwrap();
        // let guard2 = iter.next().unwrap(); will not compile
    }
}
2

You don’t need unsafe here. Put the node in an Option, so that you can take() it when it is time to access its children, and get references with the 'a lifetime:

struct NodeGuard<'a: 'b, 'b, T> {
    iter: &'b mut NodeMutIter<'a, T>,
    node: Option<&'a mut Node<T>>,
}
impl<'a, 'b, T> Drop for NodeGuard<'a, 'b, T> {
    fn drop(&mut self) {
        for n in self.node.take().unwrap().children.iter_mut().rev() {
            //            ^^^^^^^^^^^^^^^^
            self.iter.nodes.push_front(n);
        }
    }
}
impl<'a, 'b, T> Deref for NodeGuard<'a, 'b, T> {
    type Target = Node<T>;
    fn deref(&self) -> &Self::Target {
        self.node.as_ref().unwrap()
    }
}
impl<'a, 'b, T> DerefMut for NodeGuard<'a, 'b, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.node.as_mut().unwrap()
    }
}

This is essentially a workaround for the fact that Drop::drop() is given a mutable reference rather than ownership of the fields.

3 Likes