use crate::node;
#[cfg(any(test, feature = "serde"))]
use crate::node::Node;
use crate::tree::{self, PatriciaTree};
use crate::{BorrowedBytes, Bytes};
use alloc::borrow::ToOwned;
use alloc::string::String;
use alloc::vec::Vec;
use core::fmt;
use core::iter::FromIterator;
use core::marker::PhantomData;
pub type PatriciaMap<V> = GenericPatriciaMap<Vec<u8>, V>;
pub type StringPatriciaMap<V> = GenericPatriciaMap<String, V>;
pub struct GenericPatriciaMap<K, V> {
tree: PatriciaTree<V>,
_key: PhantomData<K>,
}
impl<K, V> GenericPatriciaMap<K, V> {
pub fn new() -> Self {
GenericPatriciaMap {
tree: PatriciaTree::new(),
_key: PhantomData,
}
}
pub fn clear(&mut self) {
self.tree.clear();
}
pub fn len(&self) -> usize {
self.tree.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[cfg(feature = "serde")]
pub(crate) fn from_node(node: Node<V>) -> Self {
Self {
tree: node.into(),
_key: PhantomData,
}
}
#[cfg(any(test, feature = "serde"))]
pub(crate) fn as_node(&self) -> &Node<V> {
self.tree.root()
}
#[cfg(test)]
pub(crate) fn into_node(self) -> Node<V> {
self.tree.into_root()
}
}
impl<K: Bytes, V> GenericPatriciaMap<K, V> {
pub fn contains_key<Q: AsRef<K::Borrowed>>(&self, key: Q) -> bool {
self.tree.get(key.as_ref()).is_some()
}
pub fn get<Q: AsRef<K::Borrowed>>(&self, key: Q) -> Option<&V> {
self.tree.get(key.as_ref())
}
pub fn get_mut<Q: AsRef<K::Borrowed>>(&mut self, key: Q) -> Option<&mut V> {
self.tree.get_mut(key.as_ref())
}
pub fn get_longest_common_prefix<'a, Q>(&self, key: &'a Q) -> Option<(&'a K::Borrowed, &V)>
where
Q: ?Sized + AsRef<K::Borrowed>,
{
let (key, value) = self.tree.get_longest_common_prefix(key.as_ref())?;
Some((K::Borrowed::from_bytes(key), value))
}
pub fn get_longest_common_prefix_mut<'a, Q>(
&mut self,
key: &'a Q,
) -> Option<(&'a K::Borrowed, &mut V)>
where
Q: ?Sized + AsRef<K::Borrowed>,
{
let (key, value) = self.tree.get_longest_common_prefix_mut(key.as_ref())?;
Some((K::Borrowed::from_bytes(key), value))
}
pub fn longest_common_prefix_len<Q>(&self, key: &Q) -> usize
where
Q: ?Sized + AsRef<K::Borrowed>,
{
self.tree.longest_common_prefix_len(key.as_ref())
}
pub fn insert<Q: AsRef<K::Borrowed>>(&mut self, key: Q, value: V) -> Option<V> {
self.tree.insert(key.as_ref(), value)
}
pub fn remove<Q: AsRef<K::Borrowed>>(&mut self, key: Q) -> Option<V> {
self.tree.remove(key.as_ref())
}
pub fn common_prefixes<'a, 'b, Q>(
&'a self,
key: &'b Q,
) -> CommonPrefixesIter<'a, 'b, K::Borrowed, V>
where
Q: ?Sized + AsRef<K::Borrowed>,
{
CommonPrefixesIter {
key_bytes: key.as_ref().as_bytes(),
iterator: self.tree.common_prefixes(key.as_ref()),
}
}
pub fn common_prefix_values<'a, 'b, Q>(
&'a self,
key: &'b Q,
) -> impl Iterator<Item = &'a V> + use<'a, 'b, Q, K, V>
where
Q: ?Sized + AsRef<K::Borrowed>,
<K as Bytes>::Borrowed: 'b,
{
self.tree
.common_prefixes(key.as_ref())
.filter_map(|(_, n)| n.value())
}
pub fn common_prefix_values_owned<'a>(
&'a self,
key: K,
) -> impl Iterator<Item = &'a V> + use<'a, K, V>
where
K: AsRef<K::Borrowed>,
{
self.tree
.common_prefixes_owned(key)
.filter_map(|(_, n)| n.value())
}
pub fn split_by_prefix<Q: AsRef<K::Borrowed>>(&mut self, prefix: Q) -> Self {
let subtree = self.tree.split_by_prefix(prefix.as_ref());
GenericPatriciaMap {
tree: subtree,
_key: PhantomData,
}
}
pub fn iter(&self) -> Iter<'_, K, V> {
Iter::new(self.tree.nodes(), Vec::new())
}
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
IterMut::new(self.tree.nodes_mut(), Vec::new())
}
pub fn keys(&self) -> Keys<'_, K, V> {
Keys(self.iter())
}
pub fn values(&self) -> Values<'_, V> {
Values {
nodes: self.tree.nodes(),
}
}
pub fn values_mut(&mut self) -> ValuesMut<'_, V> {
ValuesMut {
nodes: self.tree.nodes_mut(),
}
}
}
impl<K: Bytes, V> GenericPatriciaMap<K, V> {
pub fn iter_prefix<'a, 'b>(
&'a self,
prefix: &'b K::Borrowed,
) -> impl Iterator<Item = (K, &'a V)> + use<'a, 'b, K, V> {
self.tree
.iter_prefix(prefix)
.into_iter()
.flat_map(move |(prefix_len, nodes)| {
Iter::<K, V>::new(nodes, Vec::from(&prefix.as_bytes()[..prefix_len]))
})
}
pub fn iter_prefix_mut<'a, 'b>(
&'a mut self,
prefix: &'b K::Borrowed,
) -> impl Iterator<Item = (K, &'a mut V)> + use<'a, 'b, K, V> {
self.tree
.iter_prefix_mut(prefix)
.into_iter()
.flat_map(move |(prefix_len, nodes)| {
IterMut::<K, V>::new(nodes, Vec::from(&prefix.as_bytes()[..prefix_len]))
})
}
}
impl<K: Bytes + fmt::Debug, V: fmt::Debug> fmt::Debug for GenericPatriciaMap<K, V> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_map().entries(self.iter()).finish()
}
}
impl<K, V: Clone> Clone for GenericPatriciaMap<K, V> {
fn clone(&self) -> Self {
Self {
tree: self.tree.clone(),
_key: PhantomData,
}
}
}
impl<K, V> Default for GenericPatriciaMap<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K: Bytes, V> IntoIterator for GenericPatriciaMap<K, V> {
type Item = (K, V);
type IntoIter = IntoIter<K, V>;
fn into_iter(self) -> Self::IntoIter {
IntoIter {
nodes: self.tree.into_nodes(),
key_bytes: Vec::new(),
_key: PhantomData,
}
}
}
impl<K, Q, V> FromIterator<(Q, V)> for GenericPatriciaMap<K, V>
where
K: Bytes,
Q: AsRef<K::Borrowed>,
{
fn from_iter<I>(iter: I) -> Self
where
I: IntoIterator<Item = (Q, V)>,
{
let mut map = GenericPatriciaMap::new();
for (k, v) in iter {
map.insert(k, v);
}
map
}
}
impl<K, Q, V> Extend<(Q, V)> for GenericPatriciaMap<K, V>
where
K: Bytes,
Q: AsRef<K::Borrowed>,
{
fn extend<I>(&mut self, iter: I)
where
I: IntoIterator<Item = (Q, V)>,
{
for (k, v) in iter {
self.insert(k, v);
}
}
}
#[derive(Debug)]
pub struct Iter<'a, K, V: 'a> {
nodes: tree::Nodes<'a, V>,
key_bytes: Vec<u8>,
key_offset: usize,
_key: PhantomData<K>,
}
impl<'a, K, V: 'a> Iter<'a, K, V> {
fn new(nodes: tree::Nodes<'a, V>, key: Vec<u8>) -> Self {
let key_offset = key.len();
Self {
nodes,
key_bytes: key,
key_offset,
_key: PhantomData,
}
}
}
impl<'a, K: Bytes, V: 'a> Iterator for Iter<'a, K, V> {
type Item = (K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
for (key_len, node) in &mut self.nodes {
self.key_bytes.truncate(self.key_offset + key_len);
self.key_bytes.extend(node.label());
if let Some(value) = node.value() {
return Some((K::Borrowed::from_bytes(&self.key_bytes).to_owned(), value));
}
}
None
}
}
#[derive(Debug)]
pub struct IntoIter<K, V> {
nodes: tree::IntoNodes<V>,
key_bytes: Vec<u8>,
_key: PhantomData<K>,
}
impl<K: Bytes, V> Iterator for IntoIter<K, V> {
type Item = (K, V);
fn next(&mut self) -> Option<Self::Item> {
for (key_len, mut node) in &mut self.nodes {
self.key_bytes.truncate(key_len);
self.key_bytes.extend(node.label());
if let Some(value) = node.take_value() {
return Some((K::Borrowed::from_bytes(&self.key_bytes).to_owned(), value));
}
}
None
}
}
#[derive(Debug)]
pub struct IterMut<'a, K, V: 'a> {
nodes: tree::NodesMut<'a, V>,
key_bytes: Vec<u8>,
key_offset: usize,
_key: PhantomData<K>,
}
impl<'a, K, V: 'a> IterMut<'a, K, V> {
fn new(nodes: tree::NodesMut<'a, V>, key: Vec<u8>) -> Self {
let key_offset = key.len();
Self {
nodes,
key_bytes: key,
key_offset,
_key: PhantomData,
}
}
}
impl<'a, K: Bytes, V: 'a> Iterator for IterMut<'a, K, V> {
type Item = (K, &'a mut V);
fn next(&mut self) -> Option<Self::Item> {
for (key_len, node) in &mut self.nodes {
self.key_bytes.truncate(self.key_offset + key_len);
self.key_bytes.extend(node.label());
if let Some(value) = node.into_value_mut() {
return Some((K::Borrowed::from_bytes(&self.key_bytes).to_owned(), value));
}
}
None
}
}
#[derive(Debug)]
pub struct Keys<'a, K, V: 'a>(Iter<'a, K, V>);
impl<'a, K: Bytes, V: 'a> Iterator for Keys<'a, K, V> {
type Item = K;
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|(k, _)| k)
}
}
#[derive(Debug)]
pub struct Values<'a, V: 'a> {
nodes: tree::Nodes<'a, V>,
}
impl<'a, V: 'a> Iterator for Values<'a, V> {
type Item = &'a V;
fn next(&mut self) -> Option<Self::Item> {
for (_, node) in &mut self.nodes {
if let Some(value) = node.value() {
return Some(value);
}
}
None
}
}
#[derive(Debug)]
pub struct ValuesMut<'a, V: 'a> {
nodes: tree::NodesMut<'a, V>,
}
impl<'a, V: 'a> Iterator for ValuesMut<'a, V> {
type Item = &'a mut V;
fn next(&mut self) -> Option<Self::Item> {
for (_, node) in &mut self.nodes {
if let Some(value) = node.into_value_mut() {
return Some(value);
}
}
None
}
}
#[derive(Debug)]
pub struct CommonPrefixesIter<'a, 'b, K: ?Sized, V> {
key_bytes: &'b [u8],
iterator: node::CommonPrefixesIter<'a, 'b, K, V>,
}
impl<'a, 'b, K, V> Iterator for CommonPrefixesIter<'a, 'b, K, V>
where
K: 'b + ?Sized + BorrowedBytes,
{
type Item = (&'b K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
for (prefix_len, n) in self.iterator.by_ref() {
if let Some(v) = n.value() {
return Some((K::from_bytes(&self.key_bytes[..prefix_len]), v));
}
}
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use rand::seq::SliceRandom;
#[test]
fn it_works() {
let input = [
("7", 7),
("43", 43),
("92", 92),
("37", 37),
("31", 31),
("21", 21),
("0", 0),
("35", 35),
("47", 47),
("82", 82),
("61", 61),
("9", 9),
];
let mut map = PatriciaMap::new();
for &(ref k, v) in input.iter() {
assert_eq!(map.insert(k, v), None);
assert_eq!(map.get(k), Some(&v));
}
}
#[test]
fn debug_works() {
let map: PatriciaMap<_> = vec![("foo", 1), ("bar", 2), ("baz", 3)]
.into_iter()
.collect();
assert_eq!(
format!("{map:?}"),
"{[98, 97, 114]: 2, [98, 97, 122]: 3, [102, 111, 111]: 1}"
);
}
#[test]
fn clear_works() {
let mut map = PatriciaMap::new();
assert!(map.is_empty());
map.insert("foo", 1);
assert!(!map.is_empty());
map.clear();
assert!(map.is_empty());
}
#[test]
fn into_iter_works() {
let map: PatriciaMap<_> = vec![("foo", 1), ("bar", 2), ("baz", 3)]
.into_iter()
.collect();
assert_eq!(
map.into_iter().collect::<Vec<_>>(),
[(Vec::from("bar"), 2), ("baz".into(), 3), ("foo".into(), 1)]
);
}
#[test]
fn iter_mut_works() {
let mut map: PatriciaMap<_> = vec![("foo", 1), ("bar", 2), ("baz", 3)]
.into_iter()
.collect();
for (_key, x) in map.iter_mut() {
(*x) *= 2;
}
assert_eq!(
map.into_iter().collect::<Vec<_>>(),
[(Vec::from("bar"), 4), ("baz".into(), 6), ("foo".into(), 2)]
);
}
#[test]
#[cfg_attr(miri, ignore)]
fn large_map_works() {
let mut input = (0..10000).map(|i| (i.to_string(), i)).collect::<Vec<_>>();
input.shuffle(&mut rand::rng());
let mut map = input.iter().cloned().collect::<PatriciaMap<_>>();
assert_eq!(map.len(), input.len());
for &(ref k, v) in input.iter() {
assert_eq!(map.get(k), Some(&v));
}
for &(ref k, v) in input.iter().take(input.len() / 2) {
assert_eq!(map.remove(k), Some(v));
assert_eq!(map.remove(k), None);
}
for (k, _) in input.iter().take(input.len() / 2) {
assert_eq!(map.get(k), None);
}
for &(ref k, v) in input.iter().skip(input.len() / 2) {
assert_eq!(map.get(k), Some(&v));
}
for &(ref k, v) in input.iter().take(input.len() / 2) {
assert_eq!(map.insert(k, v), None);
}
for &(ref k, v) in input.iter().skip(input.len() / 2) {
assert_eq!(map.insert(k, v), Some(v));
}
for &(ref k, v) in input.iter() {
assert_eq!(map.get(k), Some(&v));
}
}
#[test]
fn test_common_word_prefixes() {
let mut t = PatriciaMap::new();
t.insert(".com.foo.", vec!["b"]);
t.insert(".", vec!["a"]);
t.insert(".com.foo.bar.", vec!["c"]);
t.insert("..", vec!["e"]);
t.insert("x", vec!["d"]);
let results = t
.common_prefixes(b".com.foo.bar.baz.")
.flat_map(|(_, v)| v)
.cloned()
.collect::<Vec<_>>();
assert!(results.iter().eq(vec![&"a", &"b", &"c"].into_iter()));
}
#[test]
fn test_letter_prefixes() {
let mut t = PatriciaMap::new();
t.insert("x", vec!["x"]);
t.insert("a", vec!["a"]);
t.insert("ab", vec!["b"]);
t.insert("abc", vec!["c"]);
t.insert("abcd", vec!["d"]);
t.insert("abcdf", vec!["f"]);
let results = t
.common_prefixes(b"abcde")
.flat_map(|(_, v)| v)
.cloned()
.collect::<Vec<_>>();
assert!(results.iter().eq(vec![&"a", &"b", &"c", &"d"].into_iter()));
}
#[test]
fn test_common_prefixes() {
let mut t = PatriciaMap::new();
t.insert("b", vec!["b"]);
t.insert("a", vec!["a"]);
t.insert("c", vec!["c"]);
t.insert("..", vec!["e"]);
t.insert("x", vec!["d"]);
let results = t
.common_prefixes(b"abc")
.flat_map(|(k, v)| {
unsafe {
println!("{:?}", std::str::from_utf8_unchecked(k));
}
v
})
.cloned()
.collect::<Vec<_>>();
dbg!(&results);
assert!(results.iter().eq(vec![&"a"].into_iter()));
let mut t = PatriciaMap::new();
t.insert("ab", vec!["b"]);
t.insert("a", vec!["a"]);
t.insert("abc", vec!["c"]);
t.insert("..", vec!["e"]);
t.insert("x", vec!["d"]);
let results = t
.common_prefixes(b"abcd")
.flat_map(|(_, v)| v)
.cloned()
.collect::<Vec<_>>();
assert!(results.iter().eq(vec![&"a", &"b", &"c"].into_iter()));
let mut list = PatriciaMap::new();
list.insert(b".com.foocatnetworks.".as_ref(), vec![0_u16]);
list.insert(b".com.foocatnetworks.foo.".as_ref(), vec![1]);
list.insert(b".com.foocatnetworks.foo.baz.".as_ref(), vec![2]);
list.insert(b".com.google.".as_ref(), vec![0]);
list.insert(b".com.cisco.".as_ref(), vec![0]);
list.insert(b".org.wikipedia.".as_ref(), vec![0]);
let results = list
.common_prefixes(b".com.foocatnetworks.foo.baz.")
.flat_map(|(_, v)| v)
.cloned()
.collect::<Vec<_>>();
assert!(vec![0_u16, 1, 2].into_iter().eq(results.into_iter()));
}
#[test]
fn string_patricia_map_works() {
let mut t = PatriciaMap::new();
t.insert("🌏🗻", ()); t.insert("🌏🍔", ());
let first_label = t.as_node().child().unwrap().label();
assert!(std::str::from_utf8(first_label).is_err());
assert_eq!(first_label, [240, 159, 140, 143, 240, 159]);
let mut t = StringPatriciaMap::new();
t.insert("🌏🗻", ());
t.insert("🌏🍔", ());
let first_label = t.as_node().child().unwrap().label();
assert_eq!(std::str::from_utf8(first_label).ok(), Some("🌏"));
}
#[test]
fn issue21() {
let mut map = PatriciaMap::new();
map.insert("1", 0);
map.insert("2", 0);
map.remove("2");
map.insert("2", 0);
assert_eq!(map.len(), map.iter().count());
assert_eq!(map.len(), map.iter_mut().count());
}
#[test]
fn issue35() {
let mut map = StringPatriciaMap::<u8>::new();
map.insert("インターポール", 1);
map.insert("インターポル", 2);
map.insert("インターリーブ", 3);
map.insert("インターン", 4);
assert_eq!(map.get("インターポール"), Some(&1));
assert_eq!(map.get("インターポル"), Some(&2));
}
#[test]
fn issue42_iter_prefix() {
let mut map = StringPatriciaMap::new();
map.insert("a0/b0", 0);
map.insert("a1/b1", 0);
let items: Vec<_> = {
let prefix = "a0".to_owned();
map.iter_prefix(&prefix).collect()
};
assert_eq!(items, vec![("a0/b0".to_owned(), &0)])
}
#[test]
fn issue42_iter_prefix_mut() {
let mut map = StringPatriciaMap::new();
map.insert("a0/b0", 0);
map.insert("a1/b1", 0);
let items: Vec<_> = {
let prefix = "a0".to_owned();
map.iter_prefix_mut(&prefix).collect()
};
assert_eq!(items, vec![("a0/b0".to_owned(), &mut 0)])
}
#[test]
fn issue42_common_prefix_values() {
let mut map = StringPatriciaMap::new();
map.insert("a0/b0", 0);
map.insert("a1/b1", 0);
let items: Vec<_> = {
let prefix = "a0/b0/c0".to_owned();
map.common_prefix_values(&prefix).collect()
};
assert_eq!(items, vec![&0])
}
#[test]
fn test_owned_impl_iter() {
struct TestTrie<T> {
map: GenericPatriciaMap<Vec<u8>, T>,
}
impl<T> TestTrie<T> {
#[expect(dead_code)]
fn common_prefix_test<'a>(
&'a self,
domain: &[u8],
) -> impl Iterator<Item = &'a T> + use<'a, T> {
let domain = domain.to_vec();
self.map.common_prefix_values_owned(domain)
}
}
}
}