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 rustc_hash::FxHashMap;
6use std::hash::Hash;
7
8use crate::{PositionalHash, PositionalSequenceHash};
9
10/// Positionally sparse radix tree for efficient indexing of [PositionalSequenceHashes][`crate::PositionalSequenceHash`].
11#[derive(Clone)]
12pub struct PositionalRadixTree<V, K = PositionalSequenceHash>
13where
14    K: PositionalHash + Hash + Eq + Clone,
15{
16    // Access to a position's inner map is serialized by the outer DashMap's
17    // RefMut guard, so sharding the inner map adds memory without concurrency.
18    map: DashMap<u64, FxHashMap<K, V>>,
19}
20
21impl<V, K> PositionalRadixTree<V, K>
22where
23    K: PositionalHash + Hash + Eq + Clone,
24{
25    /// Creates a new empty [`PositionalRadixTree`].
26    pub fn new() -> Self {
27        Self {
28            map: DashMap::new(),
29        }
30    }
31
32    /// Provides the entry for the key at the given position.
33    pub fn prefix(&self, key: &K) -> dashmap::mapref::one::RefMut<'_, u64, FxHashMap<K, V>> {
34        let position = key.position();
35        self.map.entry(position).or_default()
36    }
37
38    /// Provides the sub-map for all entries at the given position.
39    pub fn position(
40        &self,
41        position: u64,
42    ) -> Option<dashmap::mapref::one::RefMut<'_, u64, FxHashMap<K, V>>> {
43        self.map.get_mut(&position)
44    }
45
46    /// Returns the number of entries in the [`PositionalRadixTree`].
47    pub fn len(&self) -> usize {
48        if self.map.is_empty() {
49            return 0;
50        }
51        self.map.iter().map(|level| level.len()).sum()
52    }
53
54    /// Returns true if the [`PositionalRadixTree`] is empty.
55    pub fn is_empty(&self) -> bool {
56        self.len() == 0
57    }
58}
59
60impl<V, K> Default for PositionalRadixTree<V, K>
61where
62    K: PositionalHash + Hash + Eq + Clone,
63{
64    fn default() -> Self {
65        Self {
66            map: DashMap::new(),
67        }
68    }
69}
dynamo_tokens/radix.rs

dynamo_tokens/
radix.rs
