use std::collections::HashMap;
pub struct InternPool {
map: HashMap<String, u32>,
strings: Vec<String>,
}
impl InternPool {
#[inline]
pub fn new() -> Self {
Self {
map: HashMap::new(),
strings: Vec::new(),
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self {
map: HashMap::with_capacity(capacity),
strings: Vec::with_capacity(capacity),
}
}
pub fn intern(&mut self, s: &str) -> u32 {
if let Some(&idx) = self.map.get(s) {
return idx;
}
let idx = self.strings.len() as u32;
self.strings.push(s.to_string());
self.map.insert(s.to_string(), idx);
idx
}
#[inline]
pub fn get(&self, idx: u32) -> Option<&str> {
self.strings.get(idx as usize).map(|s| s.as_str())
}
#[inline]
pub fn len(&self) -> usize {
self.strings.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.strings.is_empty()
}
pub fn serialize(&self) -> Vec<u8> {
if self.strings.is_empty() {
return vec![0, 0, 0, 0, 0, 0, 0, 0];
}
let count = self.strings.len() as u32;
let mut offsets = Vec::with_capacity(self.strings.len());
let mut data = Vec::new();
for s in &self.strings {
offsets.push(data.len() as u32);
data.extend_from_slice(s.as_bytes());
}
let pool_size = 8 + (offsets.len() * 4) + data.len();
let mut bytes = Vec::with_capacity(pool_size);
bytes.extend_from_slice(&(pool_size as u32).to_le_bytes());
bytes.extend_from_slice(&count.to_le_bytes());
for offset in offsets {
bytes.extend_from_slice(&offset.to_le_bytes());
}
bytes.extend_from_slice(&data);
bytes
}
pub fn deserialize(bytes: &[u8]) -> Result<(Self, usize), InternError> {
if bytes.len() < 8 {
return Err(InternError::BufferTooSmall);
}
let pool_size = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize;
if pool_size == 0 {
return Ok((Self::new(), 8));
}
if bytes.len() < pool_size {
return Err(InternError::BufferTooSmall);
}
let count = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]) as usize;
if count == 0 {
return Ok((Self::new(), pool_size));
}
let offsets_start = 8;
let offsets_end = offsets_start + (count * 4);
if bytes.len() < offsets_end {
return Err(InternError::BufferTooSmall);
}
let mut offsets = Vec::with_capacity(count);
for i in 0..count {
let offset_pos = offsets_start + (i * 4);
let offset = u32::from_le_bytes([
bytes[offset_pos],
bytes[offset_pos + 1],
bytes[offset_pos + 2],
bytes[offset_pos + 3],
]) as usize;
offsets.push(offset);
}
let data_start = offsets_end;
let data = &bytes[data_start..pool_size];
let mut strings = Vec::with_capacity(count);
let mut map = HashMap::with_capacity(count);
for (idx, &offset) in offsets.iter().enumerate() {
let end = if idx + 1 < offsets.len() {
offsets[idx + 1]
} else {
data.len()
};
if offset > data.len() || end > data.len() || offset > end {
return Err(InternError::InvalidOffset);
}
let s = std::str::from_utf8(&data[offset..end])
.map_err(|_| InternError::InvalidUtf8)?
.to_string();
map.insert(s.clone(), idx as u32);
strings.push(s);
}
Ok((Self { map, strings }, pool_size))
}
}
impl Default for InternPool {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InternError {
BufferTooSmall,
InvalidOffset,
InvalidUtf8,
}
impl std::fmt::Display for InternError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::BufferTooSmall => write!(f, "Buffer too small to contain intern pool"),
Self::InvalidOffset => write!(f, "Invalid string offset in intern pool"),
Self::InvalidUtf8 => write!(f, "Invalid UTF-8 in intern pool string data"),
}
}
}
impl std::error::Error for InternError {}
pub struct InterningSerializer {
pool: InternPool,
}
impl InterningSerializer {
#[inline]
pub fn new() -> Self {
Self {
pool: InternPool::new(),
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self {
pool: InternPool::with_capacity(capacity),
}
}
#[inline]
pub fn pool(&self) -> &InternPool {
&self.pool
}
#[inline]
pub fn pool_mut(&mut self) -> &mut InternPool {
&mut self.pool
}
#[inline]
pub fn intern(&mut self, s: &str) -> u32 {
self.pool.intern(s)
}
#[inline]
pub fn serialize_pool(&self) -> Vec<u8> {
self.pool.serialize()
}
}
impl Default for InterningSerializer {
fn default() -> Self {
Self::new()
}
}
pub struct InterningDeserializer {
pool: InternPool,
}
impl InterningDeserializer {
pub fn new(bytes: &[u8]) -> Result<(Self, usize), InternError> {
let (pool, consumed) = InternPool::deserialize(bytes)?;
Ok((Self { pool }, consumed))
}
#[inline]
pub fn get(&self, idx: u32) -> Option<&str> {
self.pool.get(idx)
}
#[inline]
pub fn pool(&self) -> &InternPool {
&self.pool
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_intern_pool_basic() {
let mut pool = InternPool::new();
let idx1 = pool.intern("hello");
let idx2 = pool.intern("world");
let idx3 = pool.intern("hello");
assert_eq!(idx1, 0);
assert_eq!(idx2, 1);
assert_eq!(idx3, 0);
assert_eq!(pool.get(0), Some("hello"));
assert_eq!(pool.get(1), Some("world"));
assert_eq!(pool.get(2), None);
}
#[test]
fn test_intern_pool_roundtrip() {
let mut pool = InternPool::new();
pool.intern("foo");
pool.intern("bar");
pool.intern("baz");
pool.intern("foo");
let bytes = pool.serialize();
let (deserialized, consumed) = InternPool::deserialize(&bytes).unwrap();
assert_eq!(consumed, bytes.len());
assert_eq!(deserialized.len(), 3); assert_eq!(deserialized.get(0), Some("foo"));
assert_eq!(deserialized.get(1), Some("bar"));
assert_eq!(deserialized.get(2), Some("baz"));
}
#[test]
fn test_empty_pool() {
let pool = InternPool::new();
let bytes = pool.serialize();
let (deserialized, consumed) = InternPool::deserialize(&bytes).unwrap();
assert_eq!(consumed, 8);
assert!(deserialized.is_empty());
}
#[test]
fn test_interning_serializer() {
let mut serializer = InterningSerializer::new();
let idx1 = serializer.intern("test");
let idx2 = serializer.intern("data");
let idx3 = serializer.intern("test");
assert_eq!(idx1, 0);
assert_eq!(idx2, 1);
assert_eq!(idx3, 0);
let bytes = serializer.serialize_pool();
let (deserializer, _) = InterningDeserializer::new(&bytes).unwrap();
assert_eq!(deserializer.get(0), Some("test"));
assert_eq!(deserializer.get(1), Some("data"));
}
#[test]
fn test_unicode_strings() {
let mut pool = InternPool::new();
pool.intern("Hello 世界");
pool.intern("🦀 Rust");
pool.intern("Привет");
let bytes = pool.serialize();
let (deserialized, _) = InternPool::deserialize(&bytes).unwrap();
assert_eq!(deserialized.get(0), Some("Hello 世界"));
assert_eq!(deserialized.get(1), Some("🦀 Rust"));
assert_eq!(deserialized.get(2), Some("Привет"));
}
#[test]
fn test_large_pool() {
let mut pool = InternPool::new();
for i in 0..1000 {
pool.intern(&format!("string_{}", i));
}
for i in 0..500 {
pool.intern(&format!("string_{}", i));
}
assert_eq!(pool.len(), 1000);
let bytes = pool.serialize();
let (deserialized, _) = InternPool::deserialize(&bytes).unwrap();
assert_eq!(deserialized.len(), 1000);
for i in 0..1000 {
assert_eq!(deserialized.get(i), Some(format!("string_{}", i).as_str()));
}
}
}