use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use crate::chunk_cache::{CacheEntry, ChunkCache};
use crate::varint::{read_uvarint, write_uvarint};
pub const DEFAULT_RESTART_INTERVAL: u32 = 16;
pub const ABSENT: u32 = 0;
#[cfg(feature = "parallel")]
const PAR_DECODE_FANOUT: usize = 32;
#[cfg(feature = "parallel")]
fn decode_pool() -> &'static rayon::ThreadPool {
use std::sync::OnceLock;
static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
POOL.get_or_init(|| {
let n = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(4)
.clamp(1, 8);
rayon::ThreadPoolBuilder::new()
.num_threads(n)
.thread_name(|i| format!("rete-decode-{i}"))
.build()
.expect("build dedicated chunk-decode thread pool")
})
}
#[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 = u32::try_from(take(&mut pos)?)
.map_err(|_| DictError::Malformed("term count out of range"))?;
let restart_interval = u32::try_from(take(&mut pos)?)
.map_err(|_| DictError::Malformed("restart interval out of range"))?;
let num_restarts = usize::try_from(take(&mut pos)?)
.map_err(|_| DictError::Malformed("restart count out of range"))?;
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;
let restart_offsets = rel
.into_iter()
.map(|o| body_start.checked_add(o))
.collect::<Option<Vec<u64>>>()
.ok_or(DictError::Malformed("restart offset out of range"))?;
Ok(SectionMeta {
term_count,
restart_interval,
restart_offsets,
})
}
#[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 type MultiChunkBulkLoader = Box<dyn Fn(&[(u8, usize)]) -> Option<Vec<Vec<u8>>> + Send + Sync>;
pub struct SectionChunk {
first_run: usize,
key: Vec<u8>,
body_start: u64,
}
impl SectionChunk {
pub fn new(first_run: usize, key: Vec<u8>, body_start: u64) -> Self {
SectionChunk {
first_run,
key,
body_start,
}
}
}
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 fn run_last_term(bytes: &[u8], run_off: usize, end: usize) -> Option<Vec<u8>> {
let lim = bytes.get(..end)?;
let mut buf = Vec::new();
let mut pos = run_entry_into(lim, run_off, &mut buf)?;
while pos < end {
pos = entry_into(lim, pos, &mut buf)?;
}
Some(buf)
}
pub fn shortest_separator(prev_last: &[u8], first: &[u8]) -> Vec<u8> {
let shared = prev_last
.iter()
.zip(first)
.take_while(|(a, b)| a == b)
.count();
if shared >= first.len() {
return first.to_vec();
}
first[..shared + 1].to_vec()
}
#[derive(Debug, Clone)]
pub struct ChunkLeaf {
pub first: usize,
pub n: usize,
pub first_run: usize,
pub key: Vec<u8>,
pub leaf: (u64, u64),
pub chunk_base: u64,
}
type ChunkCell = std::sync::OnceLock<Box<[(SectionChunk, u64, u64)]>>;
pub struct LeafedChunks {
leaves: Vec<ChunkLeaf>,
cells: Vec<ChunkCell>,
total: usize,
fetch: crate::index::LeafFetch,
codec: u8,
max_block: u64,
failed: AtomicBool,
}
impl LeafedChunks {
pub fn with_max_block(mut self, max: u64) -> Self {
self.max_block = max;
self
}
pub fn new(
leaves: Vec<ChunkLeaf>,
total: usize,
fetch: crate::index::LeafFetch,
codec: u8,
) -> Self {
let cells = leaves.iter().map(|_| std::sync::OnceLock::new()).collect();
LeafedChunks {
leaves,
cells,
total,
fetch,
codec,
max_block: u64::MAX,
failed: AtomicBool::new(false),
}
}
fn leaf_of(&self, ci: usize) -> usize {
self.leaves
.partition_point(|l| l.first <= ci)
.saturating_sub(1)
}
fn decode_leaf(&self, li: usize, stored: &[u8]) -> Option<Box<[(SectionChunk, u64, u64)]>> {
let l = &self.leaves[li];
let raw = crate::file::decompress_bounded(self.codec, stored, self.max_block).ok()?;
let ents = crate::v06::decode_chunk_leaf(&raw, l.n).ok()?;
let mut at = l.chunk_base;
let mut out = Vec::with_capacity(ents.len());
for e in ents {
out.push((SectionChunk::new(e.first_run, e.key, 0), at, e.clen));
at = at.checked_add(e.clen)?;
}
Some(out.into_boxed_slice())
}
pub fn ensure_leaves(&self, lis: &[usize]) {
let mut missing: Vec<usize> = lis
.iter()
.copied()
.filter(|&li| li < self.cells.len() && self.cells[li].get().is_none())
.collect();
missing.sort_unstable();
missing.dedup();
if missing.is_empty() {
return;
}
let ranges: Vec<(u64, u64)> = missing.iter().map(|&li| self.leaves[li].leaf).collect();
match (self.fetch)(&ranges) {
Some(blobs) if blobs.len() == missing.len() => {
for (&li, b) in missing.iter().zip(blobs) {
match self.decode_leaf(li, &b) {
Some(v) => {
let _ = self.cells[li].set(v);
}
None => self.failed.store(true, Ordering::Relaxed),
}
}
}
_ => self.failed.store(true, Ordering::Relaxed),
}
}
pub fn leaves_of(&self, cis: &[usize]) -> Vec<usize> {
let mut v: Vec<usize> = cis
.iter()
.filter(|&&ci| ci < self.total)
.map(|&ci| self.leaf_of(ci))
.collect();
v.sort_unstable();
v.dedup();
v
}
fn entry(&self, ci: usize) -> Option<&(SectionChunk, u64, u64)> {
if ci >= self.total {
return None;
}
let li = self.leaf_of(ci);
if self.cells[li].get().is_none() {
self.ensure_leaves(&[li]);
}
self.cells[li].get()?.get(ci - self.leaves[li].first)
}
pub fn range(&self, ci: usize) -> Option<(u64, u64)> {
self.entry(ci).map(|(_, off, len)| (*off, *len))
}
pub fn leaf_of_run(&self, run: usize) -> Option<usize> {
self.leaves
.partition_point(|l| l.first_run <= run)
.checked_sub(1)
}
pub fn leaf_of_term(&self, term: &[u8]) -> Option<usize> {
self.leaves
.partition_point(|l| l.key.as_slice() <= term)
.checked_sub(1)
}
pub fn load_failed(&self) -> bool {
self.failed.load(Ordering::Relaxed)
}
fn resident_bytes(&self) -> u64 {
self.cells
.iter()
.filter_map(|c| c.get())
.flat_map(|v| v.iter())
.map(|(c, _, _)| (std::mem::size_of::<SectionChunk>() + 16 + c.key.capacity()) as u64)
.sum::<u64>()
+ self
.leaves
.iter()
.map(|l| (std::mem::size_of::<ChunkLeaf>() + l.key.capacity()) as u64)
.sum::<u64>()
}
}
pub enum ChunkDir {
Full(Vec<SectionChunk>),
Leafed(Arc<LeafedChunks>),
}
impl ChunkDir {
pub fn len(&self) -> usize {
match self {
ChunkDir::Full(v) => v.len(),
ChunkDir::Leafed(l) => l.total,
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn get(&self, ci: usize) -> Option<&SectionChunk> {
match self {
ChunkDir::Full(v) => v.get(ci),
ChunkDir::Leafed(l) => l.entry(ci).map(|(c, _, _)| c),
}
}
fn next_first_run(&self, ci: usize) -> Option<usize> {
match self {
ChunkDir::Full(v) => v.get(ci + 1).map(|c| c.first_run),
ChunkDir::Leafed(l) => {
if ci + 1 >= l.total {
return None;
}
let li = l.leaf_of(ci + 1);
if l.leaves[li].first == ci + 1 {
Some(l.leaves[li].first_run)
} else {
l.entry(ci + 1).map(|(c, _, _)| c.first_run)
}
}
}
}
fn chunk_of_run(&self, run: usize) -> Option<usize> {
match self {
ChunkDir::Full(v) => v.partition_point(|c| c.first_run <= run).checked_sub(1),
ChunkDir::Leafed(l) => {
let li = l.leaf_of_run(run)?;
l.ensure_leaves(&[li]);
let cell = l.cells[li].get()?;
let j = cell.partition_point(|(c, _, _)| c.first_run <= run);
Some(l.leaves[li].first + j.checked_sub(1)?)
}
}
}
fn chunk_for_term(&self, term: &[u8]) -> Option<usize> {
match self {
ChunkDir::Full(v) => match v.len() {
0 => None,
1 => Some(0),
_ => v
.partition_point(|c| c.key.as_slice() <= term)
.checked_sub(1),
},
ChunkDir::Leafed(l) => {
if l.total == 0 {
return None;
}
let li = l.leaf_of_term(term)?;
l.ensure_leaves(&[li]);
let cell = l.cells[li].get()?;
let j = cell.partition_point(|(c, _, _)| c.key.as_slice() <= term);
Some(l.leaves[li].first + j.saturating_sub(1))
}
}
}
fn load_failed(&self) -> bool {
matches!(self, ChunkDir::Leafed(l) if l.load_failed())
}
fn directory_bytes(&self) -> u64 {
match self {
ChunkDir::Full(v) => {
v.len() as u64 * std::mem::size_of::<SectionChunk>() as u64
+ v.iter().map(|c| c.key.capacity() as u64).sum::<u64>()
}
ChunkDir::Leafed(l) => l.resident_bytes(),
}
}
}
pub struct ChunkedSection {
meta: SectionMeta,
chunks: ChunkDir,
loader: Option<ChunkLoader>,
bulk: Option<ChunkBulkLoader>,
failed: AtomicBool,
cache: Arc<ChunkCache>,
section_index: u8,
ra_next: AtomicU64,
ra_window: AtomicU32,
}
pub const READAHEAD_MAX_CHUNKS: u32 = 16;
impl ChunkedSection {
pub fn local(section_bytes: Vec<u8>, cache: Arc<ChunkCache>, section_index: u8) -> Self {
let meta = parse_meta(§ion_bytes).unwrap_or(SectionMeta {
term_count: 0,
restart_interval: 1,
restart_offsets: Vec::new(),
});
cache.insert((section_index, 0), Arc::from(section_bytes));
ChunkedSection {
meta,
chunks: ChunkDir::Full(vec![SectionChunk::new(0, Vec::new(), 0)]),
loader: None,
bulk: None,
failed: AtomicBool::new(false),
ra_next: AtomicU64::new(0),
ra_window: AtomicU32::new(1),
cache,
section_index,
}
}
pub fn from_parts(
meta: SectionMeta,
chunks: Vec<SectionChunk>,
loader: Option<ChunkLoader>,
cache: Arc<ChunkCache>,
section_index: u8,
) -> Self {
ChunkedSection {
meta,
chunks: ChunkDir::Full(chunks),
loader,
bulk: None,
failed: AtomicBool::new(false),
ra_next: AtomicU64::new(0),
ra_window: AtomicU32::new(1),
cache,
section_index,
}
}
pub fn resident(
meta: SectionMeta,
chunks: Vec<(SectionChunk, Vec<u8>)>,
cache: Arc<ChunkCache>,
section_index: u8,
) -> Self {
let descriptors = chunks
.into_iter()
.enumerate()
.map(|(ci, (chunk, body))| {
cache.insert((section_index, ci as u32), Arc::from(body));
chunk
})
.collect();
ChunkedSection {
meta,
chunks: ChunkDir::Full(descriptors),
loader: None,
bulk: None,
failed: AtomicBool::new(false),
ra_next: AtomicU64::new(0),
ra_window: AtomicU32::new(1),
cache,
section_index,
}
}
pub fn leafed(
meta: SectionMeta,
dir: Arc<LeafedChunks>,
loader: ChunkLoader,
cache: Arc<ChunkCache>,
section_index: u8,
) -> Self {
ChunkedSection {
meta,
chunks: ChunkDir::Leafed(dir),
loader: Some(loader),
bulk: None,
failed: AtomicBool::new(false),
ra_next: AtomicU64::new(0),
ra_window: AtomicU32::new(1),
cache,
section_index,
}
}
pub fn chunk_dir(&self) -> &ChunkDir {
&self.chunks
}
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| {
ci < self.chunks.len() && !self.cache.contains((self.section_index, ci as u32))
})
.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) {
self.cache
.insert((self.section_index, ci as u32), Arc::from(body));
}
}
}
}
pub(crate) fn missing_chunks(&self, cis: &[usize]) -> Vec<usize> {
if self.loader.is_none() {
return Vec::new();
}
cis.iter()
.copied()
.filter(|&ci| {
ci < self.chunks.len() && !self.cache.contains((self.section_index, ci as u32))
})
.collect()
}
pub(crate) fn seed_chunk(&self, ci: usize, body: Vec<u8>) {
self.cache
.insert((self.section_index, ci as u32), Arc::from(body));
}
pub fn meta(&self) -> &SectionMeta {
&self.meta
}
pub fn term_count(&self) -> u32 {
self.meta.term_count
}
pub(crate) fn cache(&self) -> &Arc<ChunkCache> {
&self.cache
}
pub fn directory_bytes(&self) -> u64 {
let n = self.chunks.len() as u64;
let ranges = if self.loader.is_some() && matches!(self.chunks, ChunkDir::Full(_)) {
16 * n
} else {
0
};
self.chunks.directory_bytes() + ranges + self.meta.restart_offsets.capacity() as u64 * 8
}
pub fn load_incomplete(&self) -> bool {
self.failed.load(Ordering::Relaxed) || self.chunks.load_failed()
}
pub fn reset_load_failure(&self) {
self.failed.store(false, Ordering::Relaxed);
}
fn chunk_entry(&self, ci: usize) -> Option<Arc<CacheEntry>> {
let key = (self.section_index, ci as u32);
if let Some(e) = self.cache.get(key) {
return Some(e);
}
match &self.loader {
Some(load) => {
if self.bulk.is_some() {
let seq = self.ra_next.load(Ordering::Relaxed) == ci as u64 + 1;
let w = if seq {
(self.ra_window.load(Ordering::Relaxed) * 2).min(READAHEAD_MAX_CHUNKS)
} else {
1
};
self.ra_window.store(w, Ordering::Relaxed);
let end = (ci + w as usize).min(self.chunks.len());
self.ra_next.store(end as u64 + 1, Ordering::Relaxed);
if end > ci + 1 {
self.prefetch_chunks(&(ci..end).collect::<Vec<_>>());
if let Some(e) = self.cache.get(key) {
return Some(e);
}
}
}
match load(ci) {
Some(bytes) => Some(self.cache.insert(key, Arc::from(bytes))),
None => {
self.failed.store(true, Ordering::Relaxed);
None
}
}
}
None => Some(self.cache.insert(key, Arc::from(Vec::new()))),
}
}
fn chunk_of_run(&self, run: usize) -> Option<usize> {
self.chunks.chunk_of_run(run)
}
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,
entry: &CacheEntry,
ri: usize,
) -> Option<usize> {
let chunk = self.chunks.get(ci)?;
if self.meta.restart_offsets.is_empty() {
entry
.runs(|b| chunk_run_offsets(b, 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, entry: &CacheEntry, ri: usize) -> usize {
if let Some(next) = self.chunks.next_first_run(ci) {
return next;
}
let first_run = self.chunks.get(ci).map_or(0, |c| c.first_run);
if self.meta.restart_offsets.is_empty() {
first_run + entry.runs(|b| chunk_run_offsets(b, 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 entry = self.chunk_entry(ci)?;
let bytes: &[u8] = entry.body();
let off = self.run_off_in_chunk(ci, run, &entry, 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 resolve_into(&self, jobs: &mut [(u32, usize)], out: &mut [Option<String>]) {
let ri = self.meta.restart_interval as usize;
if ri == 0 {
return;
}
let tc = self.meta.term_count;
jobs.sort_unstable_by_key(|&(id, _)| id);
let mut segments: Vec<(usize, usize, usize)> = Vec::new();
let mut i = 0;
while i < jobs.len() {
let id0 = jobs[i].0;
if id0 == ABSENT || id0 > tc {
i += 1;
continue;
}
let Some(ci) = self.chunk_of_run((id0 - 1) as usize / ri) else {
i += 1;
continue;
};
let start = i;
i += 1;
while i < jobs.len() {
let id = jobs[i].0;
if id == ABSENT || id > tc {
i += 1;
continue;
}
if self.chunk_of_run((id - 1) as usize / ri) != Some(ci) {
break; }
i += 1;
}
segments.push((ci, start, i));
}
if segments.is_empty() {
return;
}
self.decode_and_walk(&segments, jobs, out, ri, tc);
}
#[cfg(feature = "parallel")]
fn decode_and_walk(
&self,
segments: &[(usize, usize, usize)],
jobs: &[(u32, usize)],
out: &mut [Option<String>],
ri: usize,
tc: u32,
) {
use rayon::prelude::*;
for group in segments.chunks(PAR_DECODE_FANOUT) {
let has_miss = group
.iter()
.any(|&(ci, _, _)| !self.cache.contains((self.section_index, ci as u32)));
if has_miss {
let entries: Vec<Option<Arc<CacheEntry>>> = decode_pool().install(|| {
group
.par_iter()
.map(|&(ci, _, _)| self.chunk_entry(ci))
.collect()
});
for (&(ci, s, e), entry) in group.iter().zip(entries) {
if let Some(entry) = entry {
self.walk_chunk(ci, &entry, &jobs[s..e], out, ri, tc);
}
}
} else {
for &(ci, s, e) in group {
if let Some(entry) = self.chunk_entry(ci) {
self.walk_chunk(ci, &entry, &jobs[s..e], out, ri, tc);
}
}
}
}
}
#[cfg(not(feature = "parallel"))]
fn decode_and_walk(
&self,
segments: &[(usize, usize, usize)],
jobs: &[(u32, usize)],
out: &mut [Option<String>],
ri: usize,
tc: u32,
) {
for &(ci, s, e) in segments {
if let Some(entry) = self.chunk_entry(ci) {
self.walk_chunk(ci, &entry, &jobs[s..e], out, ri, tc);
}
}
}
fn walk_chunk(
&self,
ci: usize,
entry: &CacheEntry,
chunk_jobs: &[(u32, usize)],
out: &mut [Option<String>],
ri: usize,
tc: u32,
) {
let bytes: &[u8] = entry.body();
let mut k = 0;
while k < chunk_jobs.len() {
let id = chunk_jobs[k].0;
if id == ABSENT || id > tc {
k += 1;
continue;
}
let run = (id - 1) as usize / ri;
let run_start = self.run_off_in_chunk(ci, run, entry, ri).and_then(|off| {
let mut buf = Vec::new();
run_entry_into(bytes, off, &mut buf).map(|pos| (buf, pos))
});
let Some((mut buf, mut pos)) = run_start else {
while k < chunk_jobs.len() {
let id = chunk_jobs[k].0;
if id == ABSENT || id > tc {
k += 1;
continue;
}
if (id - 1) as usize / ri != run {
break;
}
k += 1;
}
continue;
};
let mut cur_step = 0usize;
let mut broken = false;
while k < chunk_jobs.len() {
let id = chunk_jobs[k].0;
if id == ABSENT || id > tc {
k += 1;
continue;
}
if (id - 1) as usize / ri != run {
break; }
let steps = (id - 1) as usize % ri;
if !broken {
while cur_step < steps {
match entry_into(bytes, pos, &mut buf) {
Some(next) => {
pos = next;
cur_step += 1;
}
None => {
broken = true;
break;
}
}
}
if !broken && cur_step == steps {
out[chunk_jobs[k].1] = Some(String::from_utf8_lossy(&buf).into_owned());
}
}
k += 1;
}
}
}
pub(crate) fn chunk_for_term(&self, term: &str) -> Option<usize> {
self.chunks.chunk_for_term(term.as_bytes())
}
pub(crate) fn rank(&self, bound: &str) -> Option<u32> {
let tc = self.meta.term_count;
let ri = self.meta.restart_interval as usize;
if tc == 0 || ri == 0 {
return Some(0);
}
let Some(ci) = self.chunk_for_term(bound) else {
return Some(0); };
let first = (self.chunks.get(ci)?.first_run.checked_mul(ri)? as u64 + 1).min(tc as u64 + 1);
let last = match self.chunks.next_first_run(ci) {
Some(next) => (next.checked_mul(ri)? as u64).min(tc as u64),
None => tc as u64,
};
let (mut lo, mut hi) = (first, last + 1);
while lo < hi {
let mid = lo + (hi - lo) / 2;
let term = self.term(u32::try_from(mid).ok()?)?;
if term.as_bytes() < bound.as_bytes() {
lo = mid + 1;
} else {
hi = mid;
}
}
u32::try_from(lo - 1).ok()
}
pub(crate) fn faults_chunks(&self) -> bool {
self.loader.is_some()
}
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 = self.chunk_for_term(term)?;
let first_run = self.chunks.get(ci)?.first_run;
let entry = self.chunk_entry(ci)?;
let bytes: &[u8] = entry.body();
let run_end = self.run_end_of_chunk(ci, &entry, 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, &entry, 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, &entry, 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.get(0).is_some_and(|c| c.body_start == 0) {
if let Some(entry) = self.cache.get((self.section_index, 0)) {
return entry.body().to_vec();
}
}
let mut out = encode_section_header(&self.meta);
for ci in 0..self.chunks.len() {
if let Some(entry) = self.chunk_entry(ci) {
out.extend_from_slice(entry.body());
}
}
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 lite_multichunk(terms: &[&str], ri: u32, runs_per_chunk: usize) -> (ChunkedSection, u32) {
let mut b = DictSectionBuilder::new().with_restart_interval(ri);
for t in terms {
b.push(*t);
}
let raw = b.build();
let full_meta = parse_meta(&raw).unwrap();
let n_runs = full_meta.restart_offsets.len();
let mut chunks = Vec::new();
let mut bodies: Vec<Vec<u8>> = Vec::new();
let mut r = 0;
while r < n_runs {
let first_run = r;
let start = full_meta.restart_offsets[r] as usize;
let r2 = (r + runs_per_chunk).min(n_runs);
let end = if r2 < n_runs {
full_meta.restart_offsets[r2] as usize
} else {
raw.len()
};
chunks.push(SectionChunk::new(first_run, Vec::new(), start as u64));
bodies.push(raw[start..end].to_vec());
r = r2;
}
assert!(chunks.len() > 1, "fixture must be multi-chunk");
let lite_meta = SectionMeta {
term_count: full_meta.term_count,
restart_interval: full_meta.restart_interval,
restart_offsets: Vec::new(),
};
let loader: ChunkLoader = Box::new(move |ci| bodies.get(ci).cloned());
let sec = ChunkedSection::from_parts(
lite_meta,
chunks,
Some(loader),
ChunkCache::unlimited_arc(),
0,
);
(sec, full_meta.term_count)
}
#[test]
fn resolve_into_matches_term_across_chunks_and_runs() {
let owned: Vec<String> = (0..240)
.map(|i| match i % 4 {
0 => format!("<http://example.org/entity/{i:05}>"),
1 => format!("<http://example.org/entity/{i:05}/sub/leaf>"),
2 => format!("\"literal value {} padded {}\"", i, "x".repeat(i % 50)),
_ => format!("_:b{i:05}"),
})
.collect();
let terms: Vec<&str> = owned.iter().map(String::as_str).collect();
for (ri, rpc) in [(1u32, 5usize), (4, 3), (16, 2)] {
let (sec, count) = lite_multichunk(&terms, ri, rpc);
let raw = {
let mut b = DictSectionBuilder::new().with_restart_interval(ri);
for t in &terms {
b.push(*t);
}
b.build()
};
let full_meta = parse_meta(&raw).unwrap();
let mut probe: Vec<u32> = (1..=count).collect();
probe.extend([1, 1, count, count, ABSENT, count + 3, count + 100]);
let mut state = 0xDEAD_BEEF_1234_5678u64;
for k in (1..probe.len()).rev() {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
probe.swap(k, (state % (k as u64 + 1)) as usize);
}
let mut jobs: Vec<(u32, usize)> = probe.iter().map(|&id| (id, 0)).collect();
for (slot, j) in jobs.iter_mut().enumerate() {
j.1 = slot;
}
let mut out: Vec<Option<String>> = vec![None; probe.len()];
sec.resolve_into(&mut jobs, &mut out);
for (slot, &id) in probe.iter().enumerate() {
let via_term = sec.term(id);
assert_eq!(
out[slot], via_term,
"resolve_into(id={id}) != term(id) at ri={ri} rpc={rpc}"
);
let truth = section_term(&raw, &full_meta, id);
assert_eq!(out[slot], truth, "resolve_into(id={id}) != section_term");
}
}
}
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}");
}
}
}