413 lines
12 KiB
Rust
413 lines
12 KiB
Rust
//! Arena-based string interner for memory-efficient string deduplication.
|
||
//!
|
||
//! This module provides a string interner that stores all strings in a single
|
||
//! contiguous buffer, minimizing allocations and improving cache locality.
|
||
//! It uses hash-based lookup for O(1) interning operations.
|
||
//!
|
||
//! # Design
|
||
//!
|
||
//! The interner uses a two-level lookup approach:
|
||
//! 1. **Primary lookup**: HashMap from 64-bit hash -> list of StringIds with that hash
|
||
//! 2. **Collision resolution**: When hashes collide, actual string content is compared
|
||
//!
|
||
//! This gives O(1) average case for both `intern()` and `get_id()` operations.
|
||
//!
|
||
//! # Example
|
||
//!
|
||
//! ```
|
||
//! use xai_codebase_graph::interner::StringInterner;
|
||
//!
|
||
//! let mut interner = StringInterner::new();
|
||
//!
|
||
//! let id1 = interner.intern("hello");
|
||
//! let id2 = interner.intern("world");
|
||
//! let id3 = interner.intern("hello"); // Returns same id as id1
|
||
//!
|
||
//! assert_eq!(id1, id3);
|
||
//! assert_ne!(id1, id2);
|
||
//! assert_eq!(interner.get(id1), Some("hello"));
|
||
//! ```
|
||
|
||
use std::hash::{Hash, Hasher};
|
||
|
||
use hashbrown::HashMap;
|
||
use nohash_hasher::BuildNoHashHasher;
|
||
use rustc_hash::FxHasher;
|
||
use serde::{Deserialize, Serialize};
|
||
use smallvec::SmallVec;
|
||
|
||
/// Type alias for HashMap with u64 keys that are already hashed.
|
||
/// Uses NoHashHasher since keys don't need re-hashing.
|
||
type U64NoHashMap<V> = HashMap<u64, V, BuildNoHashHasher<u64>>;
|
||
|
||
/// A compact identifier for an interned string.
|
||
/// Using u32 allows up to 4 billion unique strings.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||
pub struct StringId(u32);
|
||
|
||
impl StringId {
|
||
/// Create a new StringId from a raw u32 value.
|
||
#[inline]
|
||
pub const fn new(id: u32) -> Self {
|
||
Self(id)
|
||
}
|
||
|
||
/// Get the raw u32 value.
|
||
#[inline]
|
||
pub const fn as_u32(self) -> u32 {
|
||
self.0
|
||
}
|
||
}
|
||
|
||
/// Arena-based string interner for efficient string deduplication.
|
||
///
|
||
/// Stores all strings in a single contiguous buffer to minimize allocations
|
||
/// and improve cache locality. Uses a hash-based lookup for O(1) interning.
|
||
///
|
||
/// The interner stores arbitrary byte sequences, supporting paths and strings
|
||
/// that may not be valid UTF-8.
|
||
#[derive(Debug, Clone)]
|
||
pub struct StringInterner {
|
||
/// Contiguous storage for all interned byte strings
|
||
arena: Vec<u8>,
|
||
/// Maps hash -> StringId(s). Most buckets have exactly one entry.
|
||
/// Using SmallVec<[StringId; 1]> optimizes for the common case of no collisions.
|
||
/// Uses NoHashHasher since keys are already hashed.
|
||
lookup: U64NoHashMap<SmallVec<[StringId; 1]>>,
|
||
/// Maps StringId to (start, len) in arena
|
||
offsets: Vec<(u32, u16)>,
|
||
}
|
||
|
||
impl Default for StringInterner {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
impl StringInterner {
|
||
/// Create a new empty interner.
|
||
pub fn new() -> Self {
|
||
Self {
|
||
arena: Vec::new(),
|
||
lookup: U64NoHashMap::default(),
|
||
offsets: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// Create an interner with pre-allocated capacity.
|
||
///
|
||
/// # Arguments
|
||
/// * `string_bytes` - Estimated total bytes for all strings
|
||
/// * `num_strings` - Estimated number of unique strings
|
||
pub fn with_capacity(string_bytes: usize, num_strings: usize) -> Self {
|
||
Self {
|
||
arena: Vec::with_capacity(string_bytes),
|
||
lookup: U64NoHashMap::with_capacity_and_hasher(
|
||
num_strings,
|
||
BuildNoHashHasher::default(),
|
||
),
|
||
offsets: Vec::with_capacity(num_strings),
|
||
}
|
||
}
|
||
|
||
/// Intern a byte string, returning its StringId.
|
||
/// If the string is already interned, returns the existing id.
|
||
///
|
||
/// # Complexity
|
||
/// O(1) average case, O(k) worst case where k is the number of
|
||
/// hash collisions (typically 0 or 1).
|
||
pub fn intern_bytes(&mut self, s: &[u8]) -> StringId {
|
||
let hash = Self::hash_bytes(s);
|
||
|
||
// Check if already interned
|
||
if let Some(ids) = self.lookup.get(&hash) {
|
||
for &id in ids {
|
||
if self.get_bytes(id) == Some(s) {
|
||
return id;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Not found, add new
|
||
let start = self.arena.len() as u32;
|
||
let len = s.len() as u16;
|
||
|
||
self.arena.extend_from_slice(s);
|
||
|
||
let id = StringId::new(self.offsets.len() as u32);
|
||
self.offsets.push((start, len));
|
||
|
||
// Add to lookup
|
||
self.lookup.entry(hash).or_default().push(id);
|
||
|
||
id
|
||
}
|
||
|
||
/// Intern a UTF-8 string. Convenience wrapper around `intern_bytes`.
|
||
#[inline]
|
||
pub fn intern(&mut self, s: &str) -> StringId {
|
||
self.intern_bytes(s.as_bytes())
|
||
}
|
||
|
||
/// Get the StringId for a byte string without interning it.
|
||
/// Returns None if the string is not in the interner.
|
||
///
|
||
/// # Complexity
|
||
/// O(1) average case.
|
||
pub fn get_bytes_id(&self, s: &[u8]) -> Option<StringId> {
|
||
let hash = Self::hash_bytes(s);
|
||
|
||
if let Some(ids) = self.lookup.get(&hash) {
|
||
for &id in ids {
|
||
if self.get_bytes(id) == Some(s) {
|
||
return Some(id);
|
||
}
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// Get the StringId for a UTF-8 string without interning it.
|
||
#[inline]
|
||
pub fn get_id(&self, s: &str) -> Option<StringId> {
|
||
self.get_bytes_id(s.as_bytes())
|
||
}
|
||
|
||
/// Get the raw bytes for a StringId.
|
||
///
|
||
/// # Complexity
|
||
/// O(1)
|
||
pub fn get_bytes(&self, id: StringId) -> Option<&[u8]> {
|
||
let (start, len) = *self.offsets.get(id.0 as usize)?;
|
||
self.arena
|
||
.get(start as usize..(start as usize + len as usize))
|
||
}
|
||
|
||
/// Get the string for a StringId, if it's valid UTF-8.
|
||
///
|
||
/// # Complexity
|
||
/// O(1)
|
||
pub fn get(&self, id: StringId) -> Option<&str> {
|
||
self.get_bytes(id).and_then(|b| std::str::from_utf8(b).ok())
|
||
}
|
||
|
||
/// Get the string for a StringId, with lossy UTF-8 conversion.
|
||
/// Invalid UTF-8 sequences are replaced with the replacement character.
|
||
pub fn get_lossy(&self, id: StringId) -> Option<std::borrow::Cow<'_, str>> {
|
||
self.get_bytes(id).map(String::from_utf8_lossy)
|
||
}
|
||
|
||
/// Number of interned strings.
|
||
#[inline]
|
||
pub fn len(&self) -> usize {
|
||
self.offsets.len()
|
||
}
|
||
|
||
/// Check if the interner is empty.
|
||
#[inline]
|
||
pub fn is_empty(&self) -> bool {
|
||
self.offsets.is_empty()
|
||
}
|
||
|
||
/// Total bytes used by the arena.
|
||
#[inline]
|
||
pub fn arena_bytes(&self) -> usize {
|
||
self.arena.len()
|
||
}
|
||
|
||
/// Compute FxHash of a byte slice.
|
||
#[inline]
|
||
fn hash_bytes(s: &[u8]) -> u64 {
|
||
let mut hasher = FxHasher::default();
|
||
s.hash(&mut hasher);
|
||
hasher.finish()
|
||
}
|
||
|
||
/// Iterate over all strings with their IDs (only valid UTF-8).
|
||
pub fn iter(&self) -> impl Iterator<Item = (StringId, &str)> {
|
||
self.offsets
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(idx, &(start, len))| {
|
||
let bytes = self
|
||
.arena
|
||
.get(start as usize..(start as usize + len as usize))?;
|
||
let s = std::str::from_utf8(bytes).ok()?;
|
||
Some((StringId::new(idx as u32), s))
|
||
})
|
||
}
|
||
|
||
/// Iterate over all byte strings with their IDs.
|
||
pub fn iter_bytes(&self) -> impl Iterator<Item = (StringId, &[u8])> {
|
||
self.offsets
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(idx, &(start, len))| {
|
||
let bytes = self
|
||
.arena
|
||
.get(start as usize..(start as usize + len as usize))?;
|
||
Some((StringId::new(idx as u32), bytes))
|
||
})
|
||
}
|
||
|
||
/// Clear the interner, removing all strings but keeping allocated capacity.
|
||
pub fn clear(&mut self) {
|
||
self.arena.clear();
|
||
self.lookup.clear();
|
||
self.offsets.clear();
|
||
}
|
||
|
||
/// Get the internal arena for serialization purposes.
|
||
pub fn arena(&self) -> &[u8] {
|
||
&self.arena
|
||
}
|
||
|
||
/// Get the internal offsets for serialization purposes.
|
||
pub fn offsets(&self) -> &[(u32, u16)] {
|
||
&self.offsets
|
||
}
|
||
|
||
/// Release over-allocated capacity in the arena and offsets buffers.
|
||
///
|
||
/// After a bulk build the arena and offsets Vecs may hold up to 2× their
|
||
/// actual content due to doubling growth. Calling this reclaims that
|
||
/// wasted heap. The lookup table is intentionally left unshrunk because
|
||
/// it benefits from load-factor headroom.
|
||
///
|
||
/// This is an internal maintenance hook called by `ScopeGraphIndex::compact()`.
|
||
pub(crate) fn shrink_to_fit(&mut self) {
|
||
self.arena.shrink_to_fit();
|
||
self.offsets.shrink_to_fit();
|
||
}
|
||
|
||
/// Reconstruct an interner from serialized data.
|
||
///
|
||
/// This rebuilds the lookup table from the arena and offsets.
|
||
pub fn from_parts(arena: Vec<u8>, offsets: Vec<(u32, u16)>) -> Self {
|
||
let mut lookup: U64NoHashMap<SmallVec<[StringId; 1]>> =
|
||
U64NoHashMap::with_capacity_and_hasher(offsets.len(), BuildNoHashHasher::default());
|
||
|
||
for (idx, &(start, len)) in offsets.iter().enumerate() {
|
||
if let Some(bytes) = arena.get(start as usize..(start as usize + len as usize)) {
|
||
let hash = Self::hash_bytes(bytes);
|
||
let id = StringId::new(idx as u32);
|
||
lookup.entry(hash).or_default().push(id);
|
||
}
|
||
}
|
||
|
||
Self {
|
||
arena,
|
||
lookup,
|
||
offsets,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_basic_interning() {
|
||
let mut interner = StringInterner::new();
|
||
|
||
let id1 = interner.intern("src");
|
||
let id2 = interner.intern("lib");
|
||
let id3 = interner.intern("src"); // duplicate
|
||
|
||
assert_eq!(id1, id3);
|
||
assert_ne!(id1, id2);
|
||
assert_eq!(interner.get(id1), Some("src"));
|
||
assert_eq!(interner.get(id2), Some("lib"));
|
||
assert_eq!(interner.len(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_id() {
|
||
let mut interner = StringInterner::new();
|
||
|
||
let id_src = interner.intern("src");
|
||
let id_lib = interner.intern("lib");
|
||
|
||
assert_eq!(interner.get_id("src"), Some(id_src));
|
||
assert_eq!(interner.get_id("lib"), Some(id_lib));
|
||
assert_eq!(interner.get_id("nonexistent"), None);
|
||
|
||
// get_id should not modify the interner
|
||
assert_eq!(interner.len(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_bytes_interning() {
|
||
let mut interner = StringInterner::new();
|
||
|
||
// Valid UTF-8
|
||
let id1 = interner.intern_bytes(b"hello");
|
||
assert_eq!(interner.get(id1), Some("hello"));
|
||
|
||
// Invalid UTF-8
|
||
let invalid_utf8: &[u8] = &[0x80, 0x81, 0x82];
|
||
let id2 = interner.intern_bytes(invalid_utf8);
|
||
assert_eq!(interner.get(id2), None); // Not valid UTF-8
|
||
assert_eq!(interner.get_bytes(id2), Some(invalid_utf8));
|
||
|
||
// Duplicate bytes return same ID
|
||
let id3 = interner.intern_bytes(invalid_utf8);
|
||
assert_eq!(id2, id3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_many_strings() {
|
||
let mut interner = StringInterner::new();
|
||
|
||
let count = 10_000;
|
||
let mut ids = Vec::with_capacity(count);
|
||
|
||
for i in 0..count {
|
||
let s = format!("string_{}", i);
|
||
ids.push(interner.intern(&s));
|
||
}
|
||
|
||
assert_eq!(interner.len(), count);
|
||
|
||
// Verify all strings can be looked up
|
||
for (i, &id) in ids.iter().enumerate() {
|
||
let s = format!("string_{}", i);
|
||
assert_eq!(interner.get_id(&s), Some(id));
|
||
assert_eq!(interner.get(id), Some(s.as_str()));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_from_parts() {
|
||
let mut interner = StringInterner::new();
|
||
interner.intern("hello");
|
||
interner.intern("world");
|
||
interner.intern("foo");
|
||
|
||
let arena = interner.arena().to_vec();
|
||
let offsets = interner.offsets().to_vec();
|
||
|
||
let restored = StringInterner::from_parts(arena, offsets);
|
||
|
||
assert_eq!(restored.len(), 3);
|
||
assert_eq!(restored.get_id("hello"), Some(StringId::new(0)));
|
||
assert_eq!(restored.get_id("world"), Some(StringId::new(1)));
|
||
assert_eq!(restored.get_id("foo"), Some(StringId::new(2)));
|
||
}
|
||
|
||
#[test]
|
||
fn test_clear() {
|
||
let mut interner = StringInterner::new();
|
||
interner.intern("hello");
|
||
interner.intern("world");
|
||
|
||
assert_eq!(interner.len(), 2);
|
||
|
||
interner.clear();
|
||
|
||
assert_eq!(interner.len(), 0);
|
||
assert!(interner.is_empty());
|
||
assert_eq!(interner.get_id("hello"), None);
|
||
}
|
||
}
|