Searching in tree data structure along path of keys

I have the following code in Python for which I am attempting to write a Rust equivalent:

from dataclasses import dataclass
from typing import Dict, List, Union
import unittest


@dataclass
class Node:
    children: Dict[str, Union['Node', str]]


def get_parent_with_prefix(name_path: List[str], root: Node) -> Node:
    """Returns the node whose root-to-parent-node keys match the longest prefix in name_path.
    Removes the common prefix from name_path."""
    if name_path:
        name_path_first = name_path[0]
        if name_path_first in root.children:
            parent = root.children[name_path_first]
            if isinstance(parent, Node):
                del name_path[0]
                if name_path:
                    return get_parent_with_prefix(name_path, parent)
    return root


class TestNode(unittest.TestCase):
    def helper(self, node: Node, name_path: List[str], name_path_expect: List[str], node_expect: Node):
        self.assertEqual(node_expect, get_parent_with_prefix(name_path, node))
        self.assertEqual(name_path_expect, name_path)

    def test_get_parent_with_prefix(self):
        node_b = Node({'c': Node({}), 'd': 'X'})
        node_a = Node({'b': node_b})
        root = Node({'a': node_a})

        self.helper(root, ['a'], [], root)
        self.helper(root, ['a', 'b'], [], node_a)
        self.helper(root, ['a', 'b', '1'], ['1'], node_b)
        self.helper(root, ['a', 'b', 'c'], [], node_b)
        self.helper(root, ['a', 'b', 'd'], ['d'], node_b)
        self.helper(root, ['b'], ['b'], root)

In Rust, I have so far (the return type needs to be mutable):

use std::collections::HashMap;

#[derive(Debug, Eq, PartialEq)]
pub enum NodeVal {
    A(Node),
    B(String),
}

#[derive(Debug, Default, Eq, PartialEq)]
pub struct Node {
    pub children: HashMap<String, NodeVal>,
}

fn get_parent_with_prefix<'a>(name_path: &mut Vec<String>, root: &'a mut Node) -> &'a mut Node {
    match name_path.first() {
        Some(name_path_first) => match root.children.get_mut(name_path_first) {
            Some(NodeVal::A(parent)) => {
                name_path.remove(0);
                if name_path.is_empty() {
                    root
                } else {
                    get_parent_with_prefix(name_path, parent)
                }
            }
            _ => root,
        },
        _ => root,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn helper(
        node: &mut Node,
        name_path: &mut Vec<String>,
        name_path_expect: &mut Vec<String>,
        node_expect: &mut Node,
    ) {
        assert_eq!(node_expect, get_parent_with_prefix(name_path, node));
        assert_eq!(name_path_expect, name_path);
    }

    #[test]
    fn test_get_parent_with_prefix() {
        let mut node_b = Node {
            children: HashMap::from([
                ("c".to_string(), NodeVal::A(Node::default())),
                ("d".to_string(), NodeVal::B("X".to_string())),
            ]),
        };
        let mut node_a = Node {
            children: HashMap::from([("b".to_string(), NodeVal::A(node_b))]),
        };
        let mut root = Node {
            children: HashMap::from([("a".to_string(), NodeVal::A(node_a))]),
        };

        helper(
            &mut root,
            &mut vec!["a".to_string()],
            &mut Vec::new(),
            &mut root,
        );
        helper(
            &mut root,
            &mut vec!["a".to_string(), "b".to_string()],
            &mut Vec::new(),
            &mut node_a,
        );
        helper(
            &mut root,
            &mut vec!["a".to_string(), "b".to_string(), "1".to_string()],
            &mut vec!["1".to_string()],
            &mut node_b,
        );
        helper(
            &mut root,
            &mut vec!["a".to_string(), "b".to_string(), "c".to_string()],
            &mut Vec::new(),
            &mut node_b,
        );
        helper(
            &mut root,
            &mut vec!["a".to_string(), "b".to_string(), "d".to_string()],
            &mut vec!["d".to_string()],
            &mut node_b,
        );
        helper(
            &mut root,
            &mut vec!["b".to_string()],
            &mut vec!["b".to_string()],
            &mut root,
        );
    }
}

However, I am getting the errors

  • [E0499]: cannot borrow *root as mutable more than once at a time
  • [E0505]: cannot move out of root because it is borrowed

I am not sure how to address these. Any helpful suggestions are appreciated.

1 Like

If you are just starting with Rust, start by returning owned types rather than references. Once you have more experience with the language, then level-up your game.

I haven't worked out your code in detail (tip: use ```rust or ```python to get syntax highlighting), but Rust's ownership model doesn't always play with graph data structures very well. You can normally use it when traversing a tree depth-first, by making sure you're finished with one child branch before doing anything in another.

One solution is to store the graph as a list of nodes and a separate list of edges, which is what the default graph in the excellent petgraph crate uses. Another option is to use shared ownership (Rc) interior mutability (RefCell ).

But if you're new to the language this is a really hard problem. The advantage of Rust over C/C++ with graphs is that you're allowed to gloss over the complexity in C/++ by using raw pointers, but then it's sooo easy to end up with an invalid pointer or a data race or similar.

1 Like

One reason trees are trickier to work with in Rust than they should be for now is that NLL Problem Case #3 is outstanding: conditional return of a borrow conflicts with code which lexically follows, but cannot logically be reached. A workaround is to recreate the conflicting borrow within the arm that returns, so that the borrow created outside that arm can end sooner; that is often enough for the current borrow checker to figure things out.

So here's a version that does that. I made it a method, too.


However, I also had to modify your tests a bit before they would run, mainly:

  • You can't take two &mut root at once
  • You can't move something like node_a into another struct and still use it later unless you clone it when moving into the other struct

So I agree with the others that you would benefit from learning more about Rust's ownership and borrows model. Languages like Python are doing shared ownership and clones and the like all over the place as part of the base language, but you have to do such things explicitly in Rust.

4 Likes