Thursday, July 20, 2023

Rust binary tree with iterator

For few days I scratch my head over Rust. Never tried it before, not going to be a Rust developer but I feel I'm missing new challenges and this looks like one.
While the official docs are rather concise, there are good books, the Programming Rust: Fast, Safe Systems Development 2nd Edition by Blandy/Orendorff/Tindall is really good and it feels like it's written by folks who really like to help the reader. A lot of good examples and the background I missed in the official docs.
Anyway, my first successfull attempt at something not-that-obvious, a binary tree with an iterator. Doing this in C# would be a breeze, Rust had some surprises around every corner. I am lucky to have Łukasz who answers my questions, corrects my mistakes and gives critical hints on specific problems. Thanks!
The code below. For anyone wondering how much time did it take, starting from zero knowledge of the language to this point - few hours reading the book and then like 5-6 hours to get this result.
use std::fmt::Debug;

struct TreeNode<'a, T: 'a> {
    val  : &'a T,
    left : Option<Box<TreeNode<'a, T>>>,
    right: Option<Box<TreeNode<'a, T>>>
}

impl<'a, T> TreeNode<'a, T> {
    fn new(
        val: &'a T, 
        left: Option<Box<TreeNode<'a, T>>>, 
        right: Option<Box<TreeNode<'a, T>>>) -> Self {
        Self { val, left, right }
    }
}

/* C#'s IEnumerable on TreeNode - returns an iterator/enumerator */
impl<'a, T> IntoIterator for &'a TreeNode<'a, T> {
    type Item = &'a TreeNode<'a, T>;
    type IntoIter = TreeNodeIterator<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        TreeNodeIterator::init(&self)    
    }
}

/* C#'s IEnumerator */
struct TreeNodeIterator<'a, T: 'a> {
    stack: Vec<&'a TreeNode<'a, T>>
}

impl<'a, T> TreeNodeIterator<'a, T> {
    fn init(node: &'a TreeNode<'a, T>) -> Self {
        Self { stack: vec![node] }
    }
}

impl<'a, T> Iterator for TreeNodeIterator<'a, T> {
    type Item = &'a TreeNode<'a, T>;

    fn next(&mut self) -> Option<Self::Item> {
        
        if self.stack.len() > 0 {

            let curr = self.stack.pop()?;

            /* one way to tackle the child noe */
            if let Some(right) = &curr.right {
                self.stack.push( right.as_ref() );    
            }
            /* another way of doing the same */
            if curr.left.is_some() {
                self.stack.push(curr.left.as_ref().unwrap());
            }

            return Some(curr);
        }

        None
    }
}

/* auxiliary code to be able to test the enumerator */
struct TreeNodeUtility;
impl TreeNodeUtility  {
    fn do_work<'a, T> ( node: &'a TreeNode<'a, T> ) where &'a T: Debug {
        for n in node {
            println!("{:?}", n.val );
        }
    }
}

fn main() {

    let s1 = String::from("1");
    let s2 = String::from("2");
    let s3 = String::from("3");

    let root = Some( TreeNode::<String>::new(
        &s1, 
        
            Some(Box::new(TreeNode::<String>::new(
                &s2,
                None,
                None
            ))), 
        
            Some(Box::new(TreeNode::<String>::new(
                &s3,
                None,
                None
            )))
    ) ) ;

    if root.is_some() {
        let rootc = root.unwrap();
        TreeNodeUtility::do_work( &rootc );
        TreeNodeUtility::do_work( &rootc );
    }
}
After consulting the clippy and trying to improve the code further, I came up with following changes:
  • changed the type of the iterator's item to &T from &TreeNode<T>
  • changed the TreeNodeIterator to a tupled struct (a struct with just a single field that doesn't even need a name)
  • dropped the init method of the iterator struct
  • changed the "if non empty stack - pop" code into a single liner
  • changed the "if option is some - do" code into a single liner
use std::fmt::Debug;

struct TreeNode<'a, T: 'a> {
    val  : &'a T,
    left : Option<Box<TreeNode<'a, T>>>,
    right: Option<Box<TreeNode<'a, T>>>
}

impl<'a, T> TreeNode<'a, T> {
    fn new(
        val: &'a T, 
        left: Option<Box<TreeNode<'a, T>>>, 
        right: Option<Box<TreeNode<'a, T>>>) -> Self {
        Self { val, left, right }
    }
}

/* C#'s IEnumerable on TreeNode - returns an iterator/enumerator */
impl<'a, T> IntoIterator for &'a TreeNode<'a, T> {
    type Item = &'a T;
    type IntoIter = TreeNodeIterator<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        TreeNodeIterator(vec![self])
    }
}

/* C#'s IEnumerator */
struct TreeNodeIterator<'a, T: 'a>(Vec<&'a TreeNode<'a, T>>);

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

    fn next(&mut self) -> Option<Self::Item> {
        
        if let Some(curr) = self.0.pop() {

            /* one way to tackle the child node */
            if let Some(right) = &curr.right {
                self.0.push( right.as_ref() );    
            }
            /* another way of doing the same */
            if curr.left.is_some() {
                self.0.push(curr.left.as_ref()?);
            }

            return Some(curr.val);
        }

        None
    }
}

/* auxiliary code to be able to test the enumerator */
struct TreeNodeUtility;
impl TreeNodeUtility  {
    fn do_work<'a, T> ( node: &'a TreeNode<'a, T> ) where &'a T: Debug {
        for n in node {
            println!("{:?}", n );
        }
    }
}

fn main() {

    let s1 = String::from("1");
    let s2 = String::from("2");
    let s3 = String::from("3");

    let root = Some( TreeNode::<String>::new(
        &s1, 
        
            Some(Box::new(TreeNode::<String>::new(
                &s2,
                None,
                None
            ))), 
        
            Some(Box::new(TreeNode::<String>::new(
                &s3,
                None,
                None
            )))
    ) ) ;

    if let Some( rootc ) = root  {
        TreeNodeUtility::do_work( &rootc );
        TreeNodeUtility::do_work( &rootc );
    }
}

No comments: