Iteration in the multi-node tree

Hi,

I am trying to implement a tree-like automata structure. It accepts vectors of chars, and accumulates them to store strings in the leaf node. So, for example, after I store a data like this diagram,

'a' -- 'b' -- ('a': "the_value1")
    \      \
      \     -- ('c': "the_value3")
        \
          ------('c': "the_value2")

I query 'a' -> 'b' -> 'c', and I will get "the_value3."

I am a totally newbie for rust. So first I tried C-like iteration. But even in constructing the tree, I met tons of the lifetime hell... I cannot store pointers to point the current node. After meeting tons of compilation error, I cannot even figure out what can I do more..T.T

use std::collections::BTreeMap;

#[derive(Clone)]
pub enum Node {
    Internal { children: BTreeMap<char, Box<Node>> },
    Leaf(String)
}

pub struct Tree {
    head: Node,
}

impl Tree {
    pub fn new() -> Tree {
        Tree {
            head: Node::Internal { children: BTreeMap::new() }
        }
    }

    pub fn add(&mut self, name: String, keys: Vec<char>) {
        let mut cur = &mut self.head;
        for (i, key) in keys.iter().enumerate() {
            let mut next = if let Node::Internal { children: ref mut children } = *cur {
                if let Some(&mut ref n) = children.get_mut(&key) {
                    if i == keys.len() - 1 {
                        // Leaf node is already occupied.
                        panic!("Internal Error.");
                    }
                    &mut n
                } else {
                    let mut node: &mut Node = if i == keys.len() - 1 {
                        &mut Node::Leaf(name.clone())
                    } else {
                        &mut Node::Internal { children: BTreeMap::new() }
                    };
                    children.insert(*key, Box::new(node.clone()));
                    children.get_mut(&key).unwrap()
                }
            } else {
                // Leaf node cannot has children. so this must be exception.
                panic!("Internal Error.")
            };
            cur = next;
        }
    }
}

fn main() {
    let mut tree = Tree::new();
    tree.add(String::from("the_value1"), vec!['a', 'b', 'a']);
    tree.add(String::from("the_value2"), vec!['a', 'c']);
    tree.add(String::from("the_value3"), vec!['a', 'b', 'c']);
}

Can I have some hint? Sorry for stealing your time!

There's indeed a ton of errors there, because there are a ton of references!

First of all, is the definition of Node itself right? How would you store values when one key is a prefix of another. For example, if there is ('a', "val1") and ('aa', "val2")?

I would try to implement this in smaller steps. I'd start with

impl Node {
  fn insert(&mut self, key: char, value: Node) { .. }
}

After I've got that compiling, I'd proceed to

  fn insert_rec(&mut self, key: &[char], value: Node) { .. }

That way, you should be able to tell, where exactly the problems start!

Hope it's useful! :slight_smile:

1 Like

Actually I want to prevent such cases. So, if there are already ('a', "val1"), then ('aa', "val2") must be rejected, and vice versa.

Thank you for giving a good idea! I will try again!

Thank you very much! from the point from your help, I achived the goal!
And I am trying to make an non-recursive version of this!

use std::collections::BTreeMap;
use std::ops::DerefMut;

#[derive(Clone, Debug)]
pub enum Node {
    Internal { children: BTreeMap<char, Box<Node>> },
    Leaf(String)
}

impl Node {
    fn insert(&mut self, keys: Vec<char>, idx: usize, value: String) {
        if let &mut Node::Internal { ref mut children } = self {
            if idx == keys.len() - 1 {
                // Leaf node
                let leaf = Node::Leaf(value.clone());
                children.insert(keys[idx], Box::new(leaf));
            } else {
                // Inner node
                if !children.contains_key(&keys[idx]) {
                    let inner = Node::Internal { children: BTreeMap::new() };
                    children.insert(keys[idx], Box::new(inner));
                }
                if let Some(inner) = children.get_mut(&keys[idx]) {
                    inner.deref_mut().insert(keys, idx + 1, value);
                } else {
                    panic!("Internal Error.");
                }
            }
        } else {
            panic!("Internal error.");
        }
    }
}

#[derive(Debug)]
pub struct Tree {
    head: Node,
}

impl Tree {
    pub fn new() -> Tree {
        Tree {
            head: Node::Internal { children: BTreeMap::new() }
        }
    }

    pub fn insert_rec(&mut self, keys: Vec<char>, value: String) {
        self.head.insert(keys, 0, value);
    }
}

fn main() {
    let mut tree = Tree::new();
    tree.insert_rec(vec!['a', 'b', 'a'], String::from("the_value1"));
    tree.insert_rec(vec!['a', 'c'], String::from("the_value2"));
    tree.insert_rec(vec!['a', 'b', 'c'], String::from("the_value3"));

    println!("{:?}", tree);
}