use std::fmt;
use super::arena::Arena;
use super::hash::FxHashMap;
use super::string::{StringRecord, StringRepr};
use crate::value::{StringId, SymbolId, Value};
const INDEX_TAG: u32 = 1 << 31;
const SYMBOL_TAG: u32 = 1 << 30;
pub const KEY_INDEX_LIMIT: u32 = SYMBOL_TAG;
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct PropertyKey(u32);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KeyKind {
Index(u32),
String(StringId),
Symbol(SymbolId),
}
impl PropertyKey {
#[inline]
pub const fn index(index: u32) -> Option<PropertyKey> {
if index < INDEX_TAG {
Some(PropertyKey(index | INDEX_TAG))
} else {
None
}
}
#[inline]
pub fn atom(id: StringId) -> PropertyKey {
assert!(id.0 < KEY_INDEX_LIMIT, "atom index out of key range");
PropertyKey(id.0)
}
#[inline]
pub fn symbol(id: SymbolId) -> PropertyKey {
assert!(id.0 < KEY_INDEX_LIMIT, "symbol index out of key range");
PropertyKey(id.0 | SYMBOL_TAG)
}
#[inline]
pub const fn kind(self) -> KeyKind {
if self.0 & INDEX_TAG != 0 {
KeyKind::Index(self.0 & !INDEX_TAG)
} else if self.0 & SYMBOL_TAG != 0 {
KeyKind::Symbol(SymbolId(self.0 & !SYMBOL_TAG))
} else {
KeyKind::String(StringId(self.0))
}
}
#[inline]
pub const fn raw(self) -> u32 {
self.0
}
#[inline]
pub const fn from_raw(bits: u32) -> PropertyKey {
PropertyKey(bits)
}
#[inline]
pub const fn to_value(self) -> Value {
match self.kind() {
KeyKind::Index(index) => Value::int32(index as i32),
KeyKind::String(id) => Value::string(id),
KeyKind::Symbol(id) => Value::symbol(id),
}
}
#[inline]
pub const fn heap_value(self) -> Option<Value> {
match self.kind() {
KeyKind::Index(_) => None,
KeyKind::String(id) => Some(Value::string(id)),
KeyKind::Symbol(id) => Some(Value::symbol(id)),
}
}
}
impl fmt::Debug for PropertyKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.kind() {
KeyKind::Index(index) => write!(f, "Key::Index({index})"),
KeyKind::String(id) => write!(f, "Key::Atom({})", id.0),
KeyKind::Symbol(id) => write!(f, "Key::Symbol({})", id.0),
}
}
}
const OVERFLOW: u32 = 1 << 31;
#[derive(Default)]
pub struct AtomTable {
buckets: FxHashMap<u32, u32>,
collisions: Vec<Vec<StringId>>,
free_collisions: Vec<u32>,
count: usize,
}
impl AtomTable {
pub fn find(
&self,
strings: &Arena<StringRecord>,
repr: &StringRepr,
hash: u32,
) -> Option<StringId> {
self.find_by(hash, |id| strings.get(id.0).flat() == Some(repr))
}
pub fn find_one_byte(
&self,
strings: &Arena<StringRecord>,
bytes: &[u8],
hash: u32,
) -> Option<StringId> {
self.find_by(hash, |id| {
matches!(strings.get(id.0).flat(), Some(StringRepr::OneByte(atom)) if **atom == *bytes)
})
}
fn find_by(&self, hash: u32, matches: impl Fn(&StringId) -> bool) -> Option<StringId> {
let bucket = *self.buckets.get(&hash)?;
if bucket & OVERFLOW == 0 {
let id = StringId(bucket);
matches(&id).then_some(id)
} else {
self.collisions[(bucket & !OVERFLOW) as usize]
.iter()
.copied()
.find(matches)
}
}
pub fn insert(&mut self, hash: u32, id: StringId) {
debug_assert!(id.0 < KEY_INDEX_LIMIT);
self.count += 1;
match self.buckets.get_mut(&hash) {
None => {
self.buckets.insert(hash, id.0);
}
Some(bucket) if *bucket & OVERFLOW == 0 => {
let ids = vec![StringId(*bucket), id];
let list = match self.free_collisions.pop() {
Some(list) => {
self.collisions[list as usize] = ids;
list
}
None => {
self.collisions.push(ids);
self.collisions.len() as u32 - 1
}
};
*bucket = list | OVERFLOW;
}
Some(bucket) => self.collisions[(*bucket & !OVERFLOW) as usize].push(id),
}
}
pub fn remove(&mut self, hash: u32, id: StringId) {
let Some(bucket) = self.buckets.get_mut(&hash) else {
debug_assert!(false, "an atom outside the table");
return;
};
if *bucket & OVERFLOW == 0 {
debug_assert_eq!(*bucket, id.0);
self.buckets.remove(&hash);
} else {
let list = *bucket & !OVERFLOW;
let ids = &mut self.collisions[list as usize];
let Some(position) = ids.iter().position(|&other| other == id) else {
debug_assert!(false, "an atom outside the table");
return;
};
ids.swap_remove(position);
if let [last] = ids[..] {
*bucket = last.0;
self.collisions[list as usize] = Vec::new();
self.free_collisions.push(list);
}
}
self.count -= 1;
}
pub fn len(&self) -> usize {
self.count
}
pub fn is_empty(&self) -> bool {
self.count == 0
}
pub fn ids(&self) -> impl Iterator<Item = StringId> + '_ {
self.buckets.values().flat_map(move |&bucket| {
let ids: &[StringId] = if bucket & OVERFLOW == 0 {
&[]
} else {
&self.collisions[(bucket & !OVERFLOW) as usize]
};
let single = (bucket & OVERFLOW == 0).then_some(StringId(bucket));
single.into_iter().chain(ids.iter().copied())
})
}
pub fn bytes(&self) -> usize {
self.buckets.capacity() * (size_of::<(u32, u32)>() + 1)
+ self
.collisions
.iter()
.map(|ids| 24 + ids.capacity() * 4)
.sum::<usize>()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn key_tags_round_trip() {
assert_eq!(PropertyKey::index(0).unwrap().kind(), KeyKind::Index(0));
let max = INDEX_TAG - 1;
assert_eq!(PropertyKey::index(max).unwrap().kind(), KeyKind::Index(max));
assert_eq!(PropertyKey::index(INDEX_TAG), None);
assert_eq!(
PropertyKey::atom(StringId(5)).kind(),
KeyKind::String(StringId(5))
);
assert_eq!(
PropertyKey::symbol(SymbolId(5)).kind(),
KeyKind::Symbol(SymbolId(5))
);
assert_ne!(
PropertyKey::atom(StringId(5)),
PropertyKey::symbol(SymbolId(5))
);
assert_eq!(PropertyKey::index(7).unwrap().to_value(), Value::int32(7));
assert_eq!(PropertyKey::index(7).unwrap().heap_value(), None);
}
}