use crate::types::{read_varint, write_varint, ValueError};
use std::collections::HashMap;
pub const MAGIC: &[u8; 4] = b"RKDX";
pub const VERSION: u8 = 0x01;
#[derive(Debug, PartialEq)]
pub enum KeyDictError {
TruncatedData,
BadMagic,
UnsupportedVersion(u8),
InvalidName,
Varint(ValueError),
}
impl From<ValueError> for KeyDictError {
fn from(e: ValueError) -> Self {
KeyDictError::Varint(e)
}
}
impl std::fmt::Display for KeyDictError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::TruncatedData => write!(f, "key dictionary: truncated data"),
Self::BadMagic => write!(f, "key dictionary: bad magic bytes (expected RKDX)"),
Self::UnsupportedVersion(v) => write!(f, "key dictionary: unsupported version {v}"),
Self::InvalidName => write!(f, "key dictionary: name is not valid UTF-8"),
Self::Varint(e) => write!(f, "key dictionary: varint error: {e}"),
}
}
}
impl std::error::Error for KeyDictError {}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct KeyDictionary {
names: Vec<String>,
ids: HashMap<String, u32>,
}
impl KeyDictionary {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.names.len()
}
pub fn is_empty(&self) -> bool {
self.names.is_empty()
}
pub fn intern(&mut self, name: &str) -> u32 {
if let Some(&id) = self.ids.get(name) {
return id;
}
let id = self.names.len() as u32;
self.names.push(name.to_string());
self.ids.insert(name.to_string(), id);
id
}
pub fn id_of(&self, name: &str) -> Option<u32> {
self.ids.get(name).copied()
}
pub fn name_of(&self, id: u32) -> Option<&str> {
self.names.get(id as usize).map(String::as_str)
}
pub fn encode(&self, out: &mut Vec<u8>) {
out.extend_from_slice(MAGIC);
out.push(VERSION);
write_varint(out, self.names.len() as u64);
for name in &self.names {
let bytes = name.as_bytes();
write_varint(out, bytes.len() as u64);
out.extend_from_slice(bytes);
}
}
pub fn decode(data: &[u8]) -> Result<Self, KeyDictError> {
if data.len() < 5 {
return Err(KeyDictError::TruncatedData);
}
if &data[0..4] != MAGIC.as_slice() {
return Err(KeyDictError::BadMagic);
}
if data[4] != VERSION {
return Err(KeyDictError::UnsupportedVersion(data[4]));
}
let mut cursor = 5;
let (count, n) = read_varint(&data[cursor..])?;
cursor += n;
let mut dict = KeyDictionary::new();
for _ in 0..count {
let (len, n) = read_varint(&data[cursor..])?;
cursor += n;
let end = cursor
.checked_add(len as usize)
.ok_or(KeyDictError::TruncatedData)?;
if end > data.len() {
return Err(KeyDictError::TruncatedData);
}
let name =
std::str::from_utf8(&data[cursor..end]).map_err(|_| KeyDictError::InvalidName)?;
dict.intern(name);
cursor = end;
}
Ok(dict)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_dictionary() {
let dict = KeyDictionary::new();
assert!(dict.is_empty());
assert_eq!(dict.len(), 0);
assert_eq!(dict.id_of("anything"), None);
assert_eq!(dict.name_of(0), None);
}
#[test]
fn intern_assigns_dense_ids_in_order() {
let mut dict = KeyDictionary::new();
assert_eq!(dict.intern("name"), 0);
assert_eq!(dict.intern("email"), 1);
assert_eq!(dict.intern("active"), 2);
assert_eq!(dict.len(), 3);
}
#[test]
fn intern_is_idempotent_and_append_only() {
let mut dict = KeyDictionary::new();
let first = dict.intern("status");
dict.intern("createdAt");
let again = dict.intern("status");
assert_eq!(first, again);
assert_eq!(dict.len(), 2);
}
#[test]
fn id_of_does_not_mutate() {
let mut dict = KeyDictionary::new();
dict.intern("a");
assert_eq!(dict.id_of("a"), Some(0));
assert_eq!(dict.id_of("b"), None);
assert_eq!(dict.len(), 1);
}
#[test]
fn name_of_round_trips_ids() {
let mut dict = KeyDictionary::new();
dict.intern("first");
dict.intern("second");
assert_eq!(dict.name_of(0), Some("first"));
assert_eq!(dict.name_of(1), Some("second"));
assert_eq!(dict.name_of(2), None);
}
#[test]
fn encode_decode_round_trip() {
let mut dict = KeyDictionary::new();
for k in ["id", "name", "email", "país", "🔑"] {
dict.intern(k);
}
let mut buf = Vec::new();
dict.encode(&mut buf);
let restored = KeyDictionary::decode(&buf).expect("decode");
assert_eq!(restored, dict);
assert_eq!(restored.id_of("país"), Some(3));
assert_eq!(restored.name_of(4), Some("🔑"));
}
#[test]
fn empty_dictionary_round_trips() {
let dict = KeyDictionary::new();
let mut buf = Vec::new();
dict.encode(&mut buf);
let restored = KeyDictionary::decode(&buf).expect("decode");
assert!(restored.is_empty());
}
#[test]
fn decode_rejects_bad_magic_and_version() {
let mut buf = Vec::new();
KeyDictionary::new().encode(&mut buf);
let mut bad = buf.clone();
bad[0] = b'X';
assert_eq!(KeyDictionary::decode(&bad), Err(KeyDictError::BadMagic));
let mut bad = buf.clone();
bad[4] = 0x99;
assert_eq!(
KeyDictionary::decode(&bad),
Err(KeyDictError::UnsupportedVersion(0x99))
);
}
#[test]
fn decode_rejects_truncated() {
assert_eq!(KeyDictionary::decode(&[]), Err(KeyDictError::TruncatedData));
assert_eq!(
KeyDictionary::decode(b"RKDX"),
Err(KeyDictError::TruncatedData)
);
}
}