vtcode_commons/
interner.rs

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

vtcode_commons/
interner.rs
