use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::OnceLock;
use crate::varint::{read_uvarint, write_uvarint};
pub const DEFAULT_RESTART_INTERVAL: u32 = 16;
pub const ABSENT: u32 = 0;
#[derive(Debug, thiserror::Error)]
pub enum DictError {
#[error("malformed dictionary section: {0}")]
Malformed(&'static str),
}
#[derive(Default)]
pub struct DictSectionBuilder {
terms: Vec<String>,
restart_interval: u32,
}
pub fn env_restart_interval() -> u32 {
static R: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
*R.get_or_init(|| {
std::env::var("RETE_DICT_RESTART_INTERVAL")
.ok()
.and_then(|v| v.parse::<u32>().ok())
.filter(|&v| v >= 1)
.unwrap_or(DEFAULT_RESTART_INTERVAL)
})
}
impl DictSectionBuilder {
pub fn new() -> Self {
Self {
terms: Vec::new(),
restart_interval: env_restart_interval(),
}
}
pub fn with_restart_interval(mut self, r: u32) -> Self {
assert!(r >= 1, "restart interval must be >= 1");
self.restart_interval = r;
self
}
pub fn push(&mut self, term: impl Into<String>) {
self.terms.push(term.into());
}
pub fn build(mut self) -> Vec<u8> {
self.terms.sort_unstable();
self.terms.dedup();
let r = self.restart_interval as usize;
let n = self.terms.len();
let num_restarts = n.div_ceil(r);
let mut body = Vec::new();
let mut restart_offsets = Vec::with_capacity(num_restarts);
let mut prev = "";
for (i, term) in self.terms.iter().enumerate() {
if i % r == 0 {
restart_offsets.push(body.len() as u64);
write_uvarint(&mut body, 0);
write_uvarint(&mut body, term.len() as u64);
body.extend_from_slice(term.as_bytes());
} else {
let shared = common_prefix_len(prev, term);
let suffix = &term.as_bytes()[shared..];
write_uvarint(&mut body, shared as u64);
write_uvarint(&mut body, suffix.len() as u64);
body.extend_from_slice(suffix);
}
prev = term;
}
let mut out = Vec::new();
write_uvarint(&mut out, n as u64);
write_uvarint(&mut out, self.restart_interval as u64);
write_uvarint(&mut out, num_restarts as u64);
for off in &restart_offsets {
write_uvarint(&mut out, *off);
}
out.extend_from_slice(&body);
out
}
}
#[derive(Debug, Clone)]
pub struct SectionMeta {
pub term_count: u32,
pub restart_interval: u32,
pub restart_offsets: Vec<u64>,
}
pub fn parse_meta(bytes: &[u8]) -> Result<SectionMeta, DictError> {
let mut pos = 0;
let take = |pos: &mut usize| -> Result<u64, DictError> {
let (v, n) =
read_uvarint(&bytes[*pos..]).ok_or(DictError::Malformed("truncated header"))?;
*pos += n;
Ok(v)
};
let term_count = take(&mut pos)? as u32;
let restart_interval = take(&mut pos)? as u32;
let num_restarts = take(&mut pos)? as usize;
if restart_interval == 0 {
return Err(DictError::Malformed("zero restart interval"));
}
let mut rel = Vec::with_capacity(num_restarts.min(bytes.len()));
for _ in 0..num_restarts {
rel.push(take(&mut pos)?);
}
let body_start = pos as u64;
Ok(SectionMeta {
term_count,
restart_interval,
restart_offsets: rel.into_iter().map(|o| body_start + o).collect(),
})
}
#[inline]
fn entry_into(bytes: &[u8], pos: usize, buf: &mut Vec<u8>) -> Option<usize> {
let (shared, n1) = read_uvarint(bytes.get(pos..)?)?;
let p = pos + n1;
let (suf, n2) = read_uvarint(bytes.get(p..)?)?;
let start = p + n2;
let end = start
.checked_add(suf as usize)
.filter(|&e| e <= bytes.len())?;
if shared as usize > buf.len() {
return None;
}
buf.truncate(shared as usize);
buf.extend_from_slice(&bytes[start..end]);
Some(end)
}
fn run_entry_into(bytes: &[u8], off: usize, buf: &mut Vec<u8>) -> Option<usize> {
buf.clear(); entry_into(bytes, off, buf)
}
pub fn section_term(bytes: &[u8], meta: &SectionMeta, id: u32) -> Option<String> {
if id == ABSENT || id > meta.term_count {
return None;
}
let idx = (id - 1) as usize;
let run = idx / meta.restart_interval as usize;
let steps = idx % meta.restart_interval as usize;
let mut buf = Vec::new();
let mut pos = run_entry_into(bytes, *meta.restart_offsets.get(run)? as usize, &mut buf)?;
for _ in 0..steps {
pos = entry_into(bytes, pos, &mut buf)?;
}
Some(String::from_utf8_lossy(&buf).into_owned())
}
pub fn section_id(bytes: &[u8], meta: &SectionMeta, term: &str) -> Option<u32> {
let mut buf = Vec::new();
let mut lo = 0usize;
let mut hi = meta.restart_offsets.len();
while lo < hi {
let mid = (lo + hi) / 2;
run_entry_into(bytes, meta.restart_offsets[mid] as usize, &mut buf)?;
if buf.as_slice() <= term.as_bytes() {
lo = mid + 1;
} else {
hi = mid;
}
}
if lo == 0 {
return None; }
let run = lo - 1;
let mut pos = run_entry_into(bytes, meta.restart_offsets[run] as usize, &mut buf)?;
let base_id = (run * meta.restart_interval as usize) as u32 + 1;
let run_len = meta.restart_interval.min(
meta.term_count
.saturating_sub(run as u32 * meta.restart_interval),
);
for step in 0..run_len {
if buf.as_slice() == term.as_bytes() {
return Some(base_id + step);
}
if buf.as_slice() > term.as_bytes() {
return None;
}
if step + 1 < run_len {
pos = entry_into(bytes, pos, &mut buf)?;
}
}
None
}
pub type ChunkLoader = Box<dyn Fn(usize) -> Option<Vec<u8>> + Send + Sync>;
pub type ChunkBulkLoader = Box<dyn Fn(&[usize]) -> Option<Vec<Vec<u8>>> + Send + Sync>;
pub struct SectionChunk {
first_run: usize,
first_term: Vec<u8>,
body_start: u64,
data: OnceLock<Vec<u8>>,
runs: OnceLock<Vec<usize>>,
}
impl SectionChunk {
pub fn remote(first_run: usize, first_term: Vec<u8>, body_start: u64) -> Self {
SectionChunk {
first_run,
first_term,
body_start,
data: OnceLock::new(),
runs: OnceLock::new(),
}
}
pub fn resident(first_run: usize, first_term: Vec<u8>, body_start: u64, data: Vec<u8>) -> Self {
let cell = OnceLock::new();
let _ = cell.set(data);
SectionChunk {
first_run,
first_term,
body_start,
data: cell,
runs: OnceLock::new(),
}
}
fn run_offsets(&self, data: &[u8], restart_interval: usize) -> &[usize] {
if let Some(r) = self.runs.get() {
return r;
}
if data.is_empty() {
return &[];
}
self.runs
.get_or_init(|| chunk_run_offsets(data, restart_interval))
}
}
fn chunk_run_offsets(data: &[u8], restart_interval: usize) -> Vec<usize> {
let mut offs = vec![0usize];
let mut pos = 0usize;
let mut buf = Vec::new();
let mut count = 0usize;
while pos < data.len() {
let Some(next) = entry_into(data, pos, &mut buf) else {
break;
};
count += 1;
pos = next;
if pos < data.len() && count.is_multiple_of(restart_interval) {
offs.push(pos);
}
}
offs
}
pub fn run_first_term(bytes: &[u8], off: usize) -> Option<Vec<u8>> {
let mut buf = Vec::new();
run_entry_into(bytes, off, &mut buf)?;
Some(buf)
}
pub struct ChunkedSection {
meta: SectionMeta,
chunks: Vec<SectionChunk>,
loader: Option<ChunkLoader>,
bulk: Option<ChunkBulkLoader>,
failed: AtomicBool,
}
impl ChunkedSection {
pub fn local(section_bytes: Vec<u8>) -> Self {
let meta = parse_meta(§ion_bytes).unwrap_or(SectionMeta {
term_count: 0,
restart_interval: 1,
restart_offsets: Vec::new(),
});
let data = OnceLock::new();
let _ = data.set(section_bytes);
ChunkedSection {
meta,
chunks: vec![SectionChunk {
first_run: 0,
first_term: Vec::new(),
body_start: 0,
data,
runs: OnceLock::new(),
}],
loader: None,
bulk: None,
failed: AtomicBool::new(false),
}
}
pub fn from_parts(
meta: SectionMeta,
chunks: Vec<SectionChunk>,
loader: Option<ChunkLoader>,
) -> Self {
ChunkedSection {
meta,
chunks,
loader,
bulk: None,
failed: AtomicBool::new(false),
}
}
pub fn with_bulk_loader(mut self, bulk: ChunkBulkLoader) -> Self {
self.bulk = Some(bulk);
self
}
pub fn prefetch_all(&self) {
self.prefetch_chunks(&(0..self.chunks.len()).collect::<Vec<_>>());
}
pub fn prefetch_chunks(&self, cis: &[usize]) {
let Some(bulk) = &self.bulk else { return };
let missing: Vec<usize> = cis
.iter()
.copied()
.filter(|&ci| self.chunks.get(ci).is_some_and(|c| c.data.get().is_none()))
.collect();
if missing.len() < 2 {
return;
}
if let Some(bodies) = bulk(&missing) {
if bodies.len() == missing.len() {
for (&ci, body) in missing.iter().zip(bodies) {
let _ = self.chunks[ci].data.set(body);
}
}
}
}
pub fn meta(&self) -> &SectionMeta {
&self.meta
}
pub fn term_count(&self) -> u32 {
self.meta.term_count
}
pub fn load_incomplete(&self) -> bool {
self.failed.load(Ordering::Relaxed)
}
pub fn reset_load_failure(&self) {
self.failed.store(false, Ordering::Relaxed);
}
fn chunk_data(&self, ci: usize) -> &[u8] {
let cell = &self.chunks[ci].data;
if let Some(d) = cell.get() {
return d;
}
match &self.loader {
Some(load) => match load(ci) {
Some(bytes) => cell.get_or_init(|| bytes),
None => {
self.failed.store(true, Ordering::Relaxed);
&[]
}
},
None => cell.get_or_init(Vec::new),
}
}
fn chunk_of_run(&self, run: usize) -> Option<usize> {
let i = self.chunks.partition_point(|c| c.first_run <= run);
i.checked_sub(1)
}
pub fn chunk_of_id(&self, id: u32) -> Option<usize> {
if id == ABSENT || id > self.meta.term_count {
return None;
}
let run = (id - 1) as usize / self.meta.restart_interval as usize;
self.chunk_of_run(run)
}
fn run_off_in_chunk(&self, ci: usize, run: usize, bytes: &[u8], ri: usize) -> Option<usize> {
let chunk = &self.chunks[ci];
if self.meta.restart_offsets.is_empty() {
chunk
.run_offsets(bytes, ri)
.get(run.checked_sub(chunk.first_run)?)
.copied()
} else {
self.meta
.restart_offsets
.get(run)?
.checked_sub(chunk.body_start)
.map(|o| o as usize)
}
}
fn run_end_of_chunk(&self, ci: usize, bytes: &[u8], ri: usize) -> usize {
if let Some(next) = self.chunks.get(ci + 1) {
return next.first_run;
}
let chunk = &self.chunks[ci];
if self.meta.restart_offsets.is_empty() {
chunk.first_run + chunk.run_offsets(bytes, ri).len()
} else {
self.meta.restart_offsets.len()
}
}
pub(crate) fn chunk_of_term(&self, id: u32) -> Option<usize> {
if id == ABSENT || id > self.meta.term_count {
return None;
}
self.chunk_of_run((id - 1) as usize / self.meta.restart_interval as usize)
}
pub fn term(&self, id: u32) -> Option<String> {
if id == ABSENT || id > self.meta.term_count {
return None;
}
let idx = (id - 1) as usize;
let ri = self.meta.restart_interval as usize;
let run = idx / ri;
let steps = idx % ri;
let ci = self.chunk_of_run(run)?;
let bytes = self.chunk_data(ci);
let off = self.run_off_in_chunk(ci, run, bytes, ri)?;
let mut buf = Vec::new();
let mut pos = run_entry_into(bytes, off, &mut buf)?;
for _ in 0..steps {
pos = entry_into(bytes, pos, &mut buf)?;
}
Some(String::from_utf8_lossy(&buf).into_owned())
}
pub fn id(&self, term: &str) -> Option<u32> {
if self.chunks.is_empty() {
return None;
}
let ri = self.meta.restart_interval as usize;
let ci = if self.chunks.len() == 1 {
0
} else {
let i = self
.chunks
.partition_point(|c| c.first_term.as_slice() <= term.as_bytes());
i.checked_sub(1)?
};
let first_run = self.chunks[ci].first_run;
let bytes = self.chunk_data(ci);
let run_end = self.run_end_of_chunk(ci, bytes, ri);
let mut buf = Vec::new();
let mut lo = first_run;
let mut hi = run_end;
while lo < hi {
let mid = (lo + hi) / 2;
let off = self.run_off_in_chunk(ci, mid, bytes, ri)?;
run_entry_into(bytes, off, &mut buf)?;
if buf.as_slice() <= term.as_bytes() {
lo = mid + 1;
} else {
hi = mid;
}
}
if lo == first_run {
return None; }
let run = lo - 1;
let off = self.run_off_in_chunk(ci, run, bytes, ri)?;
let mut pos = run_entry_into(bytes, off, &mut buf)?;
let base_id = (run * ri) as u32 + 1;
let run_len = self.meta.restart_interval.min(
self.meta
.term_count
.saturating_sub(run as u32 * self.meta.restart_interval),
);
for step in 0..run_len {
if buf.as_slice() == term.as_bytes() {
return Some(base_id + step);
}
if buf.as_slice() > term.as_bytes() {
return None;
}
if step + 1 < run_len {
pos = entry_into(bytes, pos, &mut buf)?;
}
}
None
}
pub fn raw_section_bytes(&self) -> Vec<u8> {
if self.chunks.len() == 1 && self.chunks[0].body_start == 0 {
if let Some(bytes) = self.chunks[0].data.get() {
return bytes.clone();
}
}
let mut out = encode_section_header(&self.meta);
for ci in 0..self.chunks.len() {
out.extend_from_slice(self.chunk_data(ci));
}
out
}
}
pub fn encode_section_header(meta: &SectionMeta) -> Vec<u8> {
let body_start = meta.restart_offsets.first().copied().unwrap_or(0);
let mut out = Vec::new();
write_uvarint(&mut out, meta.term_count as u64);
write_uvarint(&mut out, meta.restart_interval as u64);
write_uvarint(&mut out, meta.restart_offsets.len() as u64);
for off in &meta.restart_offsets {
write_uvarint(&mut out, off.saturating_sub(body_start));
}
out
}
pub struct DictSection<'a> {
bytes: &'a [u8],
meta: SectionMeta,
}
impl<'a> DictSection<'a> {
pub fn parse(bytes: &'a [u8]) -> Result<Self, DictError> {
Ok(Self {
bytes,
meta: parse_meta(bytes)?,
})
}
pub fn len(&self) -> u32 {
self.meta.term_count
}
pub fn is_empty(&self) -> bool {
self.meta.term_count == 0
}
pub fn term(&self, id: u32) -> Option<String> {
section_term(self.bytes, &self.meta, id)
}
pub fn id(&self, term: &str) -> Option<u32> {
section_id(self.bytes, &self.meta, term)
}
}
fn common_prefix_len(a: &str, b: &str) -> usize {
a.bytes().zip(b.bytes()).take_while(|(x, y)| x == y).count()
}
#[cfg(test)]
mod tests {
use super::*;
fn sample() -> Vec<String> {
[
"http://ex.org/Alice",
"http://ex.org/Bob",
"http://ex.org/Alan",
"http://ex.org/knows",
"http://ex.org/Alice", "zeta",
]
.iter()
.map(|s| s.to_string())
.collect()
}
#[test]
fn round_trip_all_ids_and_terms() {
for r in [1u32, 2, 16, 1000] {
let mut b = DictSectionBuilder::new().with_restart_interval(r);
for t in sample() {
b.push(t);
}
let bytes = b.build();
let sec = DictSection::parse(&bytes).unwrap();
let mut expected = sample();
expected.sort();
expected.dedup();
assert_eq!(sec.len() as usize, expected.len());
for (i, term) in expected.iter().enumerate() {
let id = (i + 1) as u32;
assert_eq!(
sec.term(id).as_deref(),
Some(term.as_str()),
"term({id}) r={r}"
);
assert_eq!(sec.id(term), Some(id), "id({term}) r={r}");
}
}
}
#[test]
fn lookups_for_absent_terms() {
let mut b = DictSectionBuilder::new();
for t in sample() {
b.push(t);
}
let bytes = b.build();
let sec = DictSection::parse(&bytes).unwrap();
assert_eq!(sec.id("aaa-before-everything"), None);
assert_eq!(sec.id("zzz-after-everything"), None);
assert_eq!(sec.id("http://ex.org/Alic"), None); assert_eq!(sec.term(0), None);
assert_eq!(sec.term(9999), None);
}
#[test]
fn randomized_round_trip_and_near_misses_across_restart_intervals() {
let mut state = 0x9E37_79B9_7F4A_7C15u64;
let mut next = move || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
state
};
let pool: Vec<String> = (0..1500)
.map(|i| {
let n = next();
match n % 4 {
0 => format!("<http://example.org/entity/{n:x}>"),
1 => format!("\"literal value {} with spaces\"", n % 300), 2 => format!("<http://example.org/entity/{}/sub/{i}>", n % 64), _ => format!("_:b{}", n % 256),
}
})
.collect();
let mut expected = pool.clone();
expected.sort();
expected.dedup();
for r in [1u32, 3, 16, 64] {
let mut b = DictSectionBuilder::new().with_restart_interval(r);
for t in &pool {
b.push(t.clone());
}
let bytes = b.build();
let sec = DictSection::parse(&bytes).unwrap();
assert_eq!(sec.len() as usize, expected.len(), "r={r}");
for (i, term) in expected.iter().enumerate() {
let id = (i + 1) as u32;
assert_eq!(
sec.term(id).as_deref(),
Some(term.as_str()),
"term({id}) r={r}"
);
assert_eq!(sec.id(term), Some(id), "id({term}) r={r}");
}
for (i, term) in expected.iter().enumerate() {
let near_boundary = (i as u32) % r <= 1;
if !near_boundary && i % 37 != 0 {
continue;
}
let longer = format!("{term}\u{1}");
assert_eq!(sec.id(&longer), None, "near-miss long r={r}");
let mut shorter = term.clone();
shorter.pop();
if !shorter.is_empty() && expected.binary_search(&shorter).is_err() {
assert_eq!(sec.id(&shorter), None, "near-miss short {shorter:?} r={r}");
}
}
assert_eq!(sec.term(expected.len() as u32 + 1), None, "past-end r={r}");
}
}
}