crates/mimir-ingest/src/atoms.rs
//! The open set's keys, interned.
//!
//! The open set held every key as its own `String`: 14 M rows, 2.6 GB of key
//! bytes on dwarf (2026-10-16), most of them the same monikers over and over
//! — an edge key is `from␟to␟kind`, a span key `edge␟from␞to␞kind␟file␟l␟l`,
//! and every moniker is `rust-analyzer cargo <crate> <version> <path>`. The
//! distinct words of all of them are 46 MB.
//!
//! So a key is kept as at most [`MAX_PARTS`] PART ids (a part is the text
//! between U+001F / U+001E; the last part takes the rest of the key), and a
//! part as the ids of its WORDS (split on ' '). Both tables are append-only
//! while the process lives — a part or word no key uses any more stays until
//! the next load builds the set afresh — and readable from any thread (the
//! base rewrite reads keys while the consumer adds more).
//!
//! Lossless: [`Atoms::resolve`] of [`Atoms::intern`] is the key, byte for
//! byte, for every string.
use std::sync::RwLock;
use hashbrown::HashTable;
use hashbrown::DefaultHashBuilder;
use std::hash::BuildHasher;
/// Parts a key is kept as at most; the last one is the rest of the key,
/// separators and all.
pub(crate) const MAX_PARTS: usize = 7;
/// The part separators of a key (`KEY_SEP` and the edge-in-span separator).
const SEP_US: u8 = 0x1f;
const SEP_RS: u8 = 0x1e;
/// A part id's top bit: the part is followed by U+001E (not U+001F).
const RS_BIT: u32 = 1 << 31;
/// A part id's second bit: the part is followed by a separator at all.
const MORE_BIT: u32 = 1 << 30;
const ID_MASK: u32 = MORE_BIT - 1;
/// A key as part ids, on the stack.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct KeyIds {
pub len: u8,
pub ids: [u32; MAX_PARTS],
}
impl KeyIds {
pub fn as_slice(&self) -> &[u32] {
&self.ids[..self.len as usize]
}
}
/// The interned words and parts. See the module docs.
#[derive(Default)]
pub(crate) struct Atoms {
inner: RwLock<Inner>,
}
#[derive(Default)]
struct Inner {
hasher: DefaultHashBuilder,
/// Word `i` is `words[word_end[i-1]..word_end[i]]`.
words: Vec<u8>,
word_end: Vec<u64>,
word_index: HashTable<u32>,
/// Part `i` is the words `part_words[part_end[i-1]..part_end[i]]`.
part_words: Vec<u32>,
part_end: Vec<u64>,
part_index: HashTable<u32>,
}
/// Split `key` into its parts: (part bytes, the separator after it).
fn parts(key: &str) -> impl Iterator<Item = (&[u8], Option<u8>)> {
let b = key.as_bytes();
let mut start = 0usize;
let mut n = 0usize;
let mut done = false;
std::iter::from_fn(move || {
if done {
return None;
}
n += 1;
if n < MAX_PARTS
&& let Some(off) = b[start..].iter().position(|c| *c == SEP_US || *c == SEP_RS)
{
let part = &b[start..start + off];
let sep = b[start + off];
start += off + 1;
return Some((part, Some(sep)));
}
done = true;
Some((&b[start..], None))
})
}
impl Inner {
fn word(&self, i: u32) -> &[u8] {
let i = i as usize;
let from = if i == 0 { 0 } else { self.word_end[i - 1] as usize };
&self.words[from..self.word_end[i] as usize]
}
fn part(&self, i: u32) -> &[u32] {
let i = i as usize;
let from = if i == 0 { 0 } else { self.part_end[i - 1] as usize };
&self.part_words[from..self.part_end[i] as usize]
}
fn find_word(&self, w: &[u8]) -> Option<u32> {
let h = self.hasher.hash_one(w);
self.word_index.find(h, |&i| self.word(i) == w).copied()
}
fn word_id(&mut self, w: &[u8]) -> u32 {
let h = self.hasher.hash_one(w);
if let Some(&i) = self.word_index.find(h, |&i| self.word(i) == w) {
return i;
}
let id = self.word_end.len() as u32;
assert!(id <= ID_MASK, "open-set atoms: more than 2^30 words");
self.words.extend_from_slice(w);
self.word_end.push(self.words.len() as u64);
let Inner { hasher, words, word_end, word_index, .. } = self;
word_index.insert_unique(h, id, |&i| {
let i = i as usize;
let from = if i == 0 { 0 } else { word_end[i - 1] as usize };
hasher.hash_one(&words[from..word_end[i] as usize])
});
id
}
fn find_part(&self, ws: &[u32]) -> Option<u32> {
let h = self.hasher.hash_one(ws);
self.part_index.find(h, |&i| self.part(i) == ws).copied()
}
fn part_id(&mut self, ws: &[u32]) -> u32 {
let h = self.hasher.hash_one(ws);
if let Some(&i) = self.part_index.find(h, |&i| self.part(i) == ws) {
return i;
}
let id = self.part_end.len() as u32;
assert!(id <= ID_MASK, "open-set atoms: more than 2^30 key parts");
self.part_words.extend_from_slice(ws);
self.part_end.push(self.part_words.len() as u64);
let Inner { hasher, part_words, part_end, part_index, .. } = self;
part_index.insert_unique(h, id, |&i| {
let i = i as usize;
let from = if i == 0 { 0 } else { part_end[i - 1] as usize };
hasher.hash_one(&part_words[from..part_end[i] as usize])
});
id
}
}
fn tag(id: u32, sep: Option<u8>) -> u32 {
match sep {
None => id,
Some(SEP_RS) => id | MORE_BIT | RS_BIT,
Some(_) => id | MORE_BIT,
}
}
impl Atoms {
/// The ids of `key`, adding the words and parts it is the first to use.
pub fn intern(&self, key: &str) -> KeyIds {
let mut inner = self.inner.write().unwrap_or_else(|e| e.into_inner());
let mut out = KeyIds { len: 0, ids: [0; MAX_PARTS] };
let mut ws = Vec::with_capacity(8);
for (part, sep) in parts(key) {
ws.clear();
for w in part.split(|c| *c == b' ') {
ws.push(inner.word_id(w));
}
let id = inner.part_id(&ws);
out.ids[out.len as usize] = tag(id, sep);
out.len += 1;
}
out
}
/// The ids of `key` if every part of it is known — `None` means no key
/// of the set can be `key`. Adds nothing.
pub fn find(&self, key: &str) -> Option<KeyIds> {
let inner = self.inner.read().unwrap_or_else(|e| e.into_inner());
let mut out = KeyIds { len: 0, ids: [0; MAX_PARTS] };
let mut ws = Vec::with_capacity(8);
for (part, sep) in parts(key) {
ws.clear();
for w in part.split(|c| *c == b' ') {
ws.push(inner.find_word(w)?);
}
let id = inner.find_part(&ws)?;
out.ids[out.len as usize] = tag(id, sep);
out.len += 1;
}
Some(out)
}
/// The key `ids` stand for.
pub fn resolve(&self, ids: &[u32]) -> String {
let inner = self.inner.read().unwrap_or_else(|e| e.into_inner());
let mut out = Vec::with_capacity(64);
for &id in ids {
for (n, &w) in inner.part(id & ID_MASK).iter().enumerate() {
if n > 0 {
out.push(b' ');
}
out.extend_from_slice(inner.word(w));
}
if id & MORE_BIT != 0 {
out.push(if id & RS_BIT != 0 { SEP_RS } else { SEP_US });
}
}
// cut at ASCII bytes only: what went in as UTF-8 comes out as UTF-8
String::from_utf8(out).expect("open-set atoms: a key is UTF-8")
}
/// Bytes the tables hold (their capacity).
pub fn heap_bytes(&self) -> usize {
let i = self.inner.read().unwrap_or_else(|e| e.into_inner());
i.words.capacity()
+ i.word_end.capacity() * 8
+ i.word_index.capacity() * 5
+ i.part_words.capacity() * 4
+ i.part_end.capacity() * 8
+ i.part_index.capacity() * 5
}
/// Give back what the tables grew past their use (after a load).
pub fn shrink(&self) {
let mut i = self.inner.write().unwrap_or_else(|e| e.into_inner());
let Inner { hasher, words, word_end, word_index, part_words, part_end, part_index } = &mut *i;
words.shrink_to_fit();
word_end.shrink_to_fit();
part_words.shrink_to_fit();
part_end.shrink_to_fit();
word_index.shrink_to_fit(|&i| {
let i = i as usize;
let from = if i == 0 { 0 } else { word_end[i - 1] as usize };
hasher.hash_one(&words[from..word_end[i] as usize])
});
part_index.shrink_to_fit(|&i| {
let i = i as usize;
let from = if i == 0 { 0 } else { part_end[i - 1] as usize };
hasher.hash_one(&part_words[from..part_end[i] as usize])
});
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Every shape of key comes back byte for byte: no separator, empty
/// parts and words, both separators, more parts than are kept apart,
/// non-ASCII text.
#[test]
fn a_key_comes_back_as_it_went_in() {
let a = Atoms::default();
let keys = [
"",
" ",
"rust-analyzer cargo nordisk-pairing 0.5.13 bootlog/Tap#partial.",
"a\u{1f}b\u{1f}CALLS",
"edge\u{1f}f a\u{1e}t b\u{1e}CALLS\u{1f}src/main.rs\u{1f}29\u{1f}29",
"\u{1f}\u{1f}\u{1e}",
"1\u{1f}2\u{1f}3\u{1f}4\u{1f}5\u{1f}6\u{1f}7\u{1e}8\u{1f}9 10",
" double spaces ",
"åäö — ünïcode\u{1f}日本",
];
let ids: Vec<KeyIds> = keys.iter().map(|k| a.intern(k)).collect();
for (k, id) in keys.iter().zip(&ids) {
assert_eq!(a.resolve(id.as_slice()), *k);
assert_eq!(a.find(k), Some(*id), "{k:?}");
}
assert_eq!(a.find("never seen"), None);
// a moniker shared by two keys is one part
let e = a.intern("x m\u{1f}y m\u{1f}CALLS");
let n = a.intern("x m");
assert_eq!(e.ids[0] & ID_MASK, n.ids[0] & ID_MASK);
}
}