4
\$\begingroup\$

I have a piece of working code that takes some strings from Consul configuration exports, sanitizes them a little and converts them into valid json and/or hocon structures.

While I am iterating over a vector of Nodes in the process function, I couldn't come up with a better solution than having to clone them. After inspecting some related APIs and looking for possible solutions, I understand (hopefully correctly) that it's a baked-in feature of Rust that you cannot move the ownership of the 'iteratee' to the result of the iteration.

Probably, I have to re-write everything in order to optimize the performance, but I need to research more to understand what would be the most idiomatic way to write this in Rust.

In the process I'm cloning the Vec's of Node's. I want to collect them as Node's, and not as &Node's, and this doesn't play well with the ownership transfer.

I apologize for my shallow Rust knowledge in advance.

pub struct ConsulKV {
    pub Key: String,
    pub Value: Option<String>
}

pub struct RawKV {
    pub key: String,
    pub value: String
}

struct Location {
    ns: Vec<String>,
}

impl Location {
    fn path(&self) -> String {
        self.ns.join("/")
    }

    fn name(&self) -> String {
        self.ns.last().get_or_insert(&String::from(".")).clone()
    }

    fn base(&self) -> Location {
        match self.ns.first() {
            None => Location { ns: vec![] },
            Some(head) => Location {
                ns: vec![head.clone()],
            },
        }
    }

    fn drop_base_mut(&mut self) -> &Location {
        self.ns.remove(0);
        self
    }
}

impl PartialEq for Location {
    fn eq(&self, other: &Location) -> bool {
        self.ns == other.ns
    }
}

pub enum Node {
    KeyValue {
        path: Location,
        name: String,
        value: String,
    },
    Directory {
        path: Location,
        nodes: Vec<Node>,
    },
}

impl Node {
    pub fn new(r1: Vec<RawKV>) -> ResulT<Vec<Node>> {
        let raws = Node::preprocess(r1)?;
        Ok(Node::process(raws))
    }

    fn preprocess(raw_kvs: Vec<RawKV>) -> ResulT<Vec<Node>> {
        let nodes = raw_kvs.into_iter().map(|rkv| {
            let path: Vec<String> = rkv.key.split("/").map(|s| s.to_string()).collect();
            let name = path.last().ok_or("Can't construct a node")?;
            let loc = Location { ns: path.clone() };
            if !rkv.key.ends_with("/") && !rkv.value.is_empty() {
                let n = Node::KeyValue {
                    path: loc,
                    name: name.clone(),
                    value: rkv.value.clone(),
                };
                Ok(n)
            } else {
                let dir = Node::Directory {
                    path: loc,
                    nodes: vec![],
                };
                Ok(dir)
            }
        });
        nodes.collect()
    }

    fn process(mut nodes: Vec<Node>) -> Vec<Node> {
        let grouped = nodes.iter_mut().group_by(|n| n.path().base());
        grouped
            .into_iter()
            .map(|(key, group)| {
                let (leafs, dirs): (Vec<&Node>, Vec<&Node>) = group
                    .into_iter()
                    .map(|n| n.drop_base_mut())
                    .partition(|n| n.path().ns.is_empty());

                let lfs: Vec<&Node> = leafs
                    .into_iter()
                    .filter(|n| match n {
                        Node::Directory { nodes, .. } => !nodes.is_empty(),
                        _ => true,
                    })
                    .collect();
                let mut n1: Vec<Node> = lfs.into_iter().map(move |x| x.clone()).collect();
                let nodes = if dirs.is_empty() {
                    n1
                } else {
                    let dir2: Vec<Node> = dirs.into_iter().map(move |x| x.clone()).collect();
                    let mut s = Node::process(dir2);
                    n1.append(&mut s);
                    n1
                };
                Node::Directory {
                    path: key,
                    nodes: nodes,
                }
            })
            .collect::<Vec<Node>>()
    }

    fn drop_base_mut(&mut self) -> &Node {
        match self {
            Node::KeyValue { path, .. } => {
                path.drop_base_mut();
                self
            }
            Node::Directory { path, .. } => {
                path.drop_base_mut();
                self
            }
        }
    }

    fn path(&self) -> &Location {
        match self {
            Node::KeyValue { path, .. } => path,
            Node::Directory { path, .. } => path,
        }
    }
}
\$\endgroup\$
4
\$\begingroup\$

You're doing multiple clone operations because of the types you've used. Due to the fact that consul's KV store is a tree, it does not really make that much sense to represent it as a set of nested Vecs. As I'm sure you found out, you're running into serious issues trying to figure out whether keys are inserted, what keys are inserted, and how to modify them without getting the borrow checker in a twist.

I'd like to share with you a simpler implementation, one that sidesteps all the problems you've had entirely. Since we're dealing with a KV tree, we can and should take advantage of a HashMap for this - its type is absolutely ideal for what we are about to do. We're going to define two enum types, one for the keys, one for the value:

#[derive(Eq, PartialEq, Hash, Debug)]
pub enum Key {
    Leaf(String),
    Branch(String)
}

#[derive(Eq, PartialEq, Debug)]
pub enum Node {
    Leaf {
        key: String,
        value: String
    },
    Branch {
        key: String,
        children: HashMap<Key, Node>
    }
}

We're going to then proceed through a little game of indirection to avoid some of the reference issues we might encounter...

impl Node {
    fn insert_key(children: &mut HashMap<Key, Node>, key: String, value: String) {
        match children.entry(Key::Leaf(key.clone())) {
            Entry::Occupied(mut state) => {
                state.insert(Node::Leaf {
                    key: key,
                    value: value
                });
            },
            Entry::Vacant(state) => {

                state.insert(Node::Leaf {
                    key: key,
                    value: value
                });
            }
        }
    }
    fn branch(children: &mut HashMap<Key, Node>, key: String, remainder: Vec<String>, value: String) {
        match children.entry(Key::Branch(key.clone())) {
            Entry::Occupied(mut state) => {
                // We already have a branch of that name, we just
                // forward the call and move on
                state.get_mut().add_value(remainder, value)
            },
            Entry::Vacant(state) => {
                // We need to create the node
                let mut node = Node::Branch {
                    key: key,
                    children: HashMap::new()
                };
                let status = node.add_value(remainder, value);
                state.insert(node);
                status
            }
        };
    }
    pub fn get(&self, test: &Key) -> Option<&Node> {
        match self {
            Node::Branch {
                key: _key,
                ref children
            } => children.get(test),
            _ => None
        }
    }
    pub fn add_value(&mut self, mut path: Vec<String>, value: String) {
        (match self {
            Node::Leaf {
                key: _key,
                value: _value
            } => None,
            Node::Branch {
                key: _key,
                ref mut children
            } => Some(children)
        }).map(|contents| {
            match path.len() {
                0 => panic!("Path cannot be empty"),
                1 => Node::insert_key(contents, path.pop().unwrap(), value),
                _ => Node::branch(contents, path.pop().unwrap(), path, value)
            }
        });
    }
}

And finally, we create a method to construct our tree:

pub fn into_tree(collection: Vec<RawKV>) -> Node {
    // Create the root namespace
    println!("Creating nodes");
    let mut root_node = Node::Branch {
        key: "/".to_string(),
        children: HashMap::new()
    };

    for node in collection {
        let mut path_elements:Vec<String> = node.key.split("/").map(|r| r.to_string()).collect();
        path_elements.reverse();
        root_node.add_value(path_elements, node.value);
    }

    root_node
}

This is more efficient than your method in multiple respects:

  • Less allocations (I am not copying nodes unless I absolutely have to)
  • Less re-processing (no use of group_by, no partition, nothing but recursive tree access)
  • Clearer code

Let me know what you think :-)

\$\endgroup\$
  • \$\begingroup\$ I'm really impressed with your approach and the quickness you reacted to my question. Well done and huge thanks. I'm currently trying to integrate your approach with some additions, like dropping the empty children directories. Once done, I will cargo bench it against my implementation. I'm super curious about the results. Once again, great approach and neat code. \$\endgroup\$ – RedDree Oct 7 '19 at 20:00
  • \$\begingroup\$ I've had a play with my own solution on benchmarks; deeply-nested trees are O(N) worst-case, as expected (2000 keys, 256 nesting levels - this is impossible on consul BTW). Also as expected, the flatter the tree, the faster the tree-building is; 2000 keys take 3ms if deeply nested, <1ms if sparsely nested. Not bad for a first draft \$\endgroup\$ – Sébastien Renauld Oct 7 '19 at 22:35
  • \$\begingroup\$ I have written a wrapper struct in order to serialise recursively and print the tree to json: the link to the code. It was a bit easier than to write a serialiser for the Node itself. Now working on reverse functionality (making Consul config from json), and the hocon version. \$\endgroup\$ – RedDree Oct 9 '19 at 9:25

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.