dynamo_tokens/
radix.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use dashmap::DashMap;
5use std::hash::Hash;
6
7use crate::{PositionalHash, PositionalSequenceHash};
8
9/// Positionally sparse radix tree for efficient indexing of [PositionalSequenceHashes][`crate::PositionalSequenceHash`].
10#[derive(Clone)]
11pub struct PositionalRadixTree<V, K = PositionalSequenceHash>
12where
13    K: PositionalHash + Hash + Eq + Clone,
14{
15    map: DashMap<u64, DashMap<K, V>>,
16}
17
18impl<V, K> PositionalRadixTree<V, K>
19where
20    K: PositionalHash + Hash + Eq + Clone,
21{
22    /// Creates a new empty [`PositionalRadixTree`].
23    pub fn new() -> Self {
24        Self {
25            map: DashMap::new(),
26        }
27    }
28
29    /// Provides the entry for the key at the given position.
30    pub fn prefix(&self, key: &K) -> dashmap::mapref::one::RefMut<'_, u64, DashMap<K, V>> {
31        let position = key.position();
32        self.map.entry(position).or_default()
33    }
34
35    /// Provides the sub-map for all entries at the given position.
36    pub fn position(
37        &self,
38        position: u64,
39    ) -> Option<dashmap::mapref::one::RefMut<'_, u64, DashMap<K, V>>> {
40        self.map.get_mut(&position)
41    }
42
43    /// Returns the number of entries in the [`PositionalRadixTree`].
44    pub fn len(&self) -> usize {
45        if self.map.is_empty() {
46            return 0;
47        }
48        self.map.iter().map(|level| level.len()).sum()
49    }
50
51    /// Returns true if the [`PositionalRadixTree`] is empty.
52    pub fn is_empty(&self) -> bool {
53        self.len() == 0
54    }
55}
56
57impl<V, K> Default for PositionalRadixTree<V, K>
58where
59    K: PositionalHash + Hash + Eq + Clone,
60{
61    fn default() -> Self {
62        Self {
63            map: DashMap::new(),
64        }
65    }
66}
dynamo_tokens/radix.rs

dynamo_tokens/
radix.rs
