1
\$\begingroup\$

I've implemented a map using a Patricia tree, using u64 as keys. I would like general feedback on my code, but I have one thing that I think could be implemented in a lot cleaner way, but I can't figure out how to do it:

As a Patricia tree is a full binary tree (meaning all internal nodes have exactly 2 children), a node can either be a leaf containing a key and a value or an internal node containing two non-null pointers to its left and right child (and some values to figure out which sub-tree to search in). I've of course used an enum to implement this:

#[derive(Debug)]
enum Node<V> {
    Leaf {
        key: u64,
        value: V,
    },
    Internal {
        key_prefix: u64,
        branch_bit: u8,
        left: Box<Node<V>>,
        right: Box<Node<V>>,
    },

    // Only used temporarily during insertion
    _TemporaryUnused,
}

To implement the insert function I had to add the _TemporaryUnused variant, which makes pattern matching on a node a lot more annoying as a node will never be a _TemporaryUnused outside of the insert function. Is there a way to remove _TemporaryUnused, possibly using unsafe code?


Code

Cargo.toml:

[package]
name = "patricia_tree"
version = "0.1.0"
edition = "2018"

[dependencies]
duplicate = "*"

[dev-dependencies]
proptest = "*"
$ cargo tree --depth 1
patricia_tree v0.1.0 (/home/tyilo/repos/patricia_tree)
└── duplicate v0.3.0 (proc-macro)
[dev-dependencies]
└── proptest v1.0.0

lib.rs:

use duplicate::duplicate;
use std::hint::unreachable_unchecked;
use std::mem;

#[derive(Debug)]
enum Node<V> {
    Leaf {
        key: u64,
        value: V,
    },
    Internal {
        key_prefix: u64,
        branch_bit: u8,
        left: Box<Node<V>>,
        right: Box<Node<V>>,
    },

    // Only used temporarily during insertion
    _TemporaryUnused,
}

#[derive(Debug)]
pub struct PatriciaTreeMap<V> {
    size: usize,
    root: Option<Box<Node<V>>>,
}

impl<V> PatriciaTreeMap<V> {
    pub fn new() -> Self {
        Self {
            size: 0,
            root: None,
        }
    }

    pub fn len(&self) -> usize {
        self.size
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn get_prefix(key: u64, branch_bit: u8) -> u64 {
        let mask = (1 << branch_bit) - 1;
        key & mask
    }

    fn is_left(key: u64, branch_bit: u8) -> bool {
        key & (1 << branch_bit) == 0
    }

    #[duplicate(
      method                     reference(type) as_ref(v);
      [find_insertion_point]     [& type]        [v.as_ref()];
      [find_insertion_point_mut] [&mut type]     [v.as_mut()];
    )]
    #[allow(clippy::needless_arbitrary_self_type)]
    fn method(self: reference([Self]), key: u64) -> Option<reference([Node<V>])> {
        fn aux<V>(node: reference([Node<V>]), key: u64) -> reference([Node<V>]) {
            match node {
                Node::Leaf { .. } => node,
                Node::Internal {
                    key_prefix,
                    branch_bit,
                    ..
                } if *key_prefix != PatriciaTreeMap::<V>::get_prefix(key, *branch_bit) => node,
                Node::Internal {
                    branch_bit,
                    right,
                    left,
                    ..
                } => {
                    if PatriciaTreeMap::<V>::is_left(key, *branch_bit) {
                        aux(left, key)
                    } else {
                        aux(right, key)
                    }
                }
                Node::_TemporaryUnused => unsafe { unreachable_unchecked() },
            }
        }

        as_ref([self.root]).map(|r| aux(r, key))
    }

    pub fn get(&self, key: u64) -> Option<&V> {
        match self.find_insertion_point(key) {
            Some(Node::Leaf { key: k, value: v }) if k == &key => Some(v),
            _ => None,
        }
    }

    pub fn contains(&self, key: u64) -> bool {
        self.get(key).is_some()
    }

    pub fn insert(&mut self, key: u64, value: V) -> Option<V> {
        fn aux<V>(tree: &mut PatriciaTreeMap<V>, key: u64, value: V) -> Option<V> {
            fn do_insert<V>(diff: u64, key: u64, value: V, node: &mut Node<V>) -> Option<V> {
                let branch_bit = diff.trailing_zeros() as u8;
                let key_prefix = PatriciaTreeMap::<V>::get_prefix(key, branch_bit);

                let leaf = Node::Leaf { key, value };
                let old_node = mem::replace(node, Node::_TemporaryUnused);

                let (left, right) = if PatriciaTreeMap::<V>::is_left(key, branch_bit) {
                    (leaf, old_node)
                } else {
                    (old_node, leaf)
                };

                *node = Node::Internal {
                    branch_bit,
                    key_prefix,
                    left: Box::new(left),
                    right: Box::new(right),
                };

                None
            }

            let node = tree.find_insertion_point_mut(key);
            match node {
                None => {
                    tree.root = Some(Box::new(Node::Leaf { key, value }));
                    None
                }
                Some(node) => match node {
                    Node::Leaf { key: k, value: v } => {
                        if k != &key {
                            let diff = *k ^ key;
                            do_insert(diff, key, value, node)
                        } else {
                            Some(mem::replace(v, value))
                        }
                    }
                    Node::Internal { key_prefix, .. } => {
                        let diff = *key_prefix ^ key;
                        do_insert(diff, key, value, node)
                    }
                    Node::_TemporaryUnused => unsafe { unreachable_unchecked() },
                },
            }
        }

        let res = aux(self, key, value);
        self.size += res.is_none() as usize;
        res
    }
}

impl<V> Default for PatriciaTreeMap<V> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod test {
    use super::PatriciaTreeMap;
    use proptest::bits;
    use proptest::collection::hash_set;
    use proptest::collection::vec;
    use proptest::collection::SizeRange;
    use proptest::prelude::*;
    use std::collections::HashSet;
    use std::hash::Hash;

    #[test]
    fn test_empty_map() {
        let map = PatriciaTreeMap::<String>::new();
        assert_eq!(map.len(), 0);
        assert_eq!(map.get(0), None);
    }

    #[test]
    fn test_insert_return_value() {
        let mut map = PatriciaTreeMap::<String>::new();
        assert_eq!(map.get(123), None);
        assert_eq!(map.insert(123, "A".into()), None);
        assert_eq!(map.get(123), Some(&"A".into()));
        assert_eq!(map.insert(123, "B".into()), Some("A".into()));
        assert_eq!(map.get(123), Some(&"B".into()));
    }

    fn unique_vec<T>(element: T, size: impl Into<SizeRange>) -> impl Strategy<Value = Vec<T::Value>>
    where
        T: Strategy,
        T::Value: Hash + Eq,
    {
        let x = hash_set(element, size);
        x.prop_map(|v| v.into_iter().collect())
    }

    fn test_insertion_impl(keys: Vec<u64>) {
        let tree = {
            let mut tree = PatriciaTreeMap::<String>::new();
            for v in keys.iter() {
                tree.insert(*v, format!("{}", *v));
            }
            tree
        };

        let unique_keys = keys.into_iter().collect::<HashSet<u64>>();

        assert_eq!(tree.len(), unique_keys.len());

        for v in unique_keys.iter() {
            assert_eq!(tree.get(*v), Some(&format!("{}", *v)));
        }
    }

    proptest! {
        #[test]
        fn test_insert_with_duplicates(keys in vec(bits::u64::between(0, 10), 0..100)) {
            test_insertion_impl(keys)
        }

        #[test]
        fn test_insert_unique(keys in unique_vec(bits::u64::between(0, 10), 0..100)) {
            test_insertion_impl(keys)
        }
    }
}
\$\endgroup\$
2
\$\begingroup\$

The Node::_TemporaryUnused is very sad.

You need to use it here:

let old_node = mem::replace(node, Node::_TemporaryUnused);

You can look at how std::mem::replace is implemented

unsafe {
    let result = ptr::read(dest);
    ptr::write(dest, src);
    result
}

You can simply use the ptr::read and ptr::write functions directly to do what you want.

However, you do need to be careful. If a panic occurs between your call to ptr::read and ptr::write, you'll end up leaving an invalid state and invoke undefined behavior.

There are a couple of crates which encapsulate this operation: take_mut and replace_with. You'll see they make render the operation safe by handling panics that might occur.

\$\endgroup\$
1
  • \$\begingroup\$ I have now changed the code to use replace_with::replace_with_or_abort and removed Node::_TemporaryUnused, which works perfectly. \$\endgroup\$
    – Tyilo
    Sep 24 '21 at 6:53

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.