vtcode_commons/
interner.rs

1#![expect(
2    clippy::indexing_slicing,
3    clippy::cast_possible_truncation,
4    reason = "Arena offsets are bounded by the compact interner representation and IDs are range-checked by construction."
5)]
6
7//! Arena-based string interner for memory-efficient string deduplication.
8//!
9//! Stores all strings in a single contiguous buffer to minimize allocations
10//! and improve cache locality. Uses a hash-based lookup for O(1) interning.
11//!
12//! # Example
13//!
14//! ```
15//! use vtcode_commons::interner::StringInterner;
16//!
17//! let mut interner = StringInterner::new();
18//! let id1 = interner.intern("src/lib.rs");
19//! let id2 = interner.intern("src/lib.rs");
20//! assert_eq!(id1, id2);
21//! assert_eq!(interner.get(id1), Some("src/lib.rs"));
22//! ```
23
24use std::hash::{Hash, Hasher};
25
26use hashbrown::HashMap;
27use rustc_hash::FxHasher;
28use serde::{Deserialize, Serialize};
29use smallvec::SmallVec;
30
31/// Type alias for HashMap with u64 keys that are already hashed.
32type U64NoHashMap<V> = HashMap<u64, V, rustc_hash::FxBuildHasher>;
33
34/// A compact identifier for an interned string.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize)]
36pub struct StringId(u32);
37
38impl StringId {
39    /// Create a new StringId from a raw u32 value.
40    #[inline]
41    const fn new(id: u32) -> Self {
42        Self(id)
43    }
44
45    /// Get the raw u32 value.
46    #[inline]
47    pub const fn as_u32(self) -> u32 {
48        self.0
49    }
50}
51
52/// Arena-based string interner for efficient string deduplication.
53#[derive(Debug, Clone, Default)]
54pub struct StringInterner {
55    arena: Vec<u8>,
56    lookup: U64NoHashMap<SmallVec<[StringId; 1]>>,
57    offsets: Vec<(u32, u32)>,
58}
59
60impl StringInterner {
61    /// Create a new empty interner.
62    #[must_use]
63    pub fn new() -> Self {
64        Self::default()
65    }
66
67    /// Create an interner with pre-allocated capacity.
68    #[must_use]
69    pub fn with_capacity(string_bytes: usize, num_strings: usize) -> Self {
70        Self {
71            arena: Vec::with_capacity(string_bytes),
72            lookup: U64NoHashMap::with_capacity_and_hasher(num_strings, rustc_hash::FxBuildHasher),
73            offsets: Vec::with_capacity(num_strings),
74        }
75    }
76
77    /// Intern a byte string, returning its StringId.
78    fn intern_bytes(&mut self, s: &[u8]) -> StringId {
79        let hash = Self::hash_bytes(s);
80
81        if let Some(ids) = self.lookup.get(&hash) {
82            for &id in ids {
83                if self.get_bytes(id) == Some(s) {
84                    return id;
85                }
86            }
87        }
88
89        let start = self.arena.len() as u32;
90        let len = s.len() as u32;
91        self.arena.extend_from_slice(s);
92        let id = StringId::new(self.offsets.len() as u32);
93        self.offsets.push((start, len));
94        self.lookup.entry(hash).or_default().push(id);
95        id
96    }
97
98    /// Intern a UTF-8 string. Convenience wrapper around `intern_bytes`.
99    #[inline]
100    pub fn intern(&mut self, s: &str) -> StringId {
101        self.intern_bytes(s.as_bytes())
102    }
103
104    /// Get the StringId for a byte string without interning it.
105    fn get_bytes_id(&self, s: &[u8]) -> Option<StringId> {
106        let hash = Self::hash_bytes(s);
107        if let Some(ids) = self.lookup.get(&hash) {
108            for &id in ids {
109                if self.get_bytes(id) == Some(s) {
110                    return Some(id);
111                }
112            }
113        }
114        None
115    }
116
117    /// Get the StringId for a UTF-8 string without interning it.
118    #[inline]
119    fn get_id(&self, s: &str) -> Option<StringId> {
120        self.get_bytes_id(s.as_bytes())
121    }
122
123    /// Get the raw bytes for a StringId.
124    fn get_bytes(&self, id: StringId) -> Option<&[u8]> {
125        let (start, len) = self.offsets[id.0 as usize];
126        self.arena.get(start as usize..(start as usize + len as usize))
127    }
128
129    /// Get the string for a StringId, if it's valid UTF-8.
130    pub fn get(&self, id: StringId) -> Option<&str> {
131        self.get_bytes(id).and_then(|b| std::str::from_utf8(b).ok())
132    }
133
134    /// Number of interned strings.
135    #[inline]
136    fn len(&self) -> usize {
137        self.offsets.len()
138    }
139
140    /// Check if the interner is empty.
141    #[inline]
142    pub fn is_empty(&self) -> bool {
143        self.offsets.is_empty()
144    }
145
146    /// Compute FxHash of a byte slice.
147    #[inline]
148    fn hash_bytes(s: &[u8]) -> u64 {
149        let mut hasher = FxHasher::default();
150        s.hash(&mut hasher);
151        hasher.finish()
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn test_basic_interning() {
161        let mut interner = StringInterner::new();
162        let id1 = interner.intern("src");
163        let id2 = interner.intern("lib");
164        let id3 = interner.intern("src");
165        assert_eq!(id1, id3);
166        assert_ne!(id1, id2);
167        assert_eq!(interner.get(id1), Some("src"));
168        assert_eq!(interner.len(), 2);
169    }
170
171    #[test]
172    fn test_bytes_interning() {
173        let mut interner = StringInterner::new();
174        let id1 = interner.intern_bytes(b"hello");
175        assert_eq!(interner.get(id1), Some("hello"));
176    }
177
178    #[test]
179    fn test_get_id() {
180        let mut interner = StringInterner::new();
181        let id_src = interner.intern("src");
182        assert_eq!(interner.get_id("src"), Some(id_src));
183        assert_eq!(interner.get_id("nonexistent"), None);
184    }
185}
vtcode_commons/interner.rs

vtcode_commons/
interner.rs
