1use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::OnceLock;
9
10use crate::varint::{read_uvarint, write_uvarint};
11
12pub const DEFAULT_RESTART_INTERVAL: u32 = 16;
14
15pub const ABSENT: u32 = 0;
17
18#[derive(Debug, thiserror::Error)]
19pub enum DictError {
20 #[error("malformed dictionary section: {0}")]
21 Malformed(&'static str),
22}
23
24#[derive(Default)]
26pub struct DictSectionBuilder {
27 terms: Vec<String>,
28 restart_interval: u32,
29}
30
31pub fn env_restart_interval() -> u32 {
40 static R: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
41 *R.get_or_init(|| {
42 std::env::var("RETE_DICT_RESTART_INTERVAL")
43 .ok()
44 .and_then(|v| v.parse::<u32>().ok())
45 .filter(|&v| v >= 1)
46 .unwrap_or(DEFAULT_RESTART_INTERVAL)
47 })
48}
49
50impl DictSectionBuilder {
51 pub fn new() -> Self {
52 Self {
53 terms: Vec::new(),
54 restart_interval: env_restart_interval(),
55 }
56 }
57
58 pub fn with_restart_interval(mut self, r: u32) -> Self {
59 assert!(r >= 1, "restart interval must be >= 1");
60 self.restart_interval = r;
61 self
62 }
63
64 pub fn push(&mut self, term: impl Into<String>) {
65 self.terms.push(term.into());
66 }
67
68 pub fn build(mut self) -> Vec<u8> {
71 self.terms.sort_unstable();
72 self.terms.dedup();
73 let r = self.restart_interval as usize;
74 let n = self.terms.len();
75 let num_restarts = n.div_ceil(r);
76
77 let mut body = Vec::new();
79 let mut restart_offsets = Vec::with_capacity(num_restarts);
80 let mut prev = "";
81 for (i, term) in self.terms.iter().enumerate() {
82 if i % r == 0 {
83 restart_offsets.push(body.len() as u64);
84 write_uvarint(&mut body, 0);
86 write_uvarint(&mut body, term.len() as u64);
87 body.extend_from_slice(term.as_bytes());
88 } else {
89 let shared = common_prefix_len(prev, term);
90 let suffix = &term.as_bytes()[shared..];
91 write_uvarint(&mut body, shared as u64);
92 write_uvarint(&mut body, suffix.len() as u64);
93 body.extend_from_slice(suffix);
94 }
95 prev = term;
96 }
97
98 let mut out = Vec::new();
100 write_uvarint(&mut out, n as u64);
101 write_uvarint(&mut out, self.restart_interval as u64);
102 write_uvarint(&mut out, num_restarts as u64);
103 for off in &restart_offsets {
104 write_uvarint(&mut out, *off);
105 }
106 out.extend_from_slice(&body);
107 out
108 }
109}
110
111#[derive(Debug, Clone)]
114pub struct SectionMeta {
115 pub term_count: u32,
116 pub restart_interval: u32,
117 pub restart_offsets: Vec<u64>,
124}
125
126pub fn parse_meta(bytes: &[u8]) -> Result<SectionMeta, DictError> {
128 let mut pos = 0;
129 let take = |pos: &mut usize| -> Result<u64, DictError> {
130 let (v, n) =
131 read_uvarint(&bytes[*pos..]).ok_or(DictError::Malformed("truncated header"))?;
132 *pos += n;
133 Ok(v)
134 };
135 let term_count = take(&mut pos)? as u32;
136 let restart_interval = take(&mut pos)? as u32;
137 let num_restarts = take(&mut pos)? as usize;
138 if restart_interval == 0 {
139 return Err(DictError::Malformed("zero restart interval"));
140 }
141 let mut rel = Vec::with_capacity(num_restarts.min(bytes.len()));
144 for _ in 0..num_restarts {
145 rel.push(take(&mut pos)?);
146 }
147 let body_start = pos as u64;
148 Ok(SectionMeta {
149 term_count,
150 restart_interval,
151 restart_offsets: rel.into_iter().map(|o| body_start + o).collect(),
152 })
153}
154
155#[inline]
162fn entry_into(bytes: &[u8], pos: usize, buf: &mut Vec<u8>) -> Option<usize> {
163 let (shared, n1) = read_uvarint(bytes.get(pos..)?)?;
164 let p = pos + n1;
165 let (suf, n2) = read_uvarint(bytes.get(p..)?)?;
166 let start = p + n2;
167 let end = start
168 .checked_add(suf as usize)
169 .filter(|&e| e <= bytes.len())?;
170 if shared as usize > buf.len() {
171 return None;
172 }
173 buf.truncate(shared as usize);
174 buf.extend_from_slice(&bytes[start..end]);
175 Some(end)
176}
177
178fn run_entry_into(bytes: &[u8], off: usize, buf: &mut Vec<u8>) -> Option<usize> {
180 buf.clear(); entry_into(bytes, off, buf)
182}
183
184pub fn section_term(bytes: &[u8], meta: &SectionMeta, id: u32) -> Option<String> {
187 if id == ABSENT || id > meta.term_count {
188 return None;
189 }
190 let idx = (id - 1) as usize;
191 let run = idx / meta.restart_interval as usize;
192 let steps = idx % meta.restart_interval as usize;
193 let mut buf = Vec::new();
194 let mut pos = run_entry_into(bytes, *meta.restart_offsets.get(run)? as usize, &mut buf)?;
195 for _ in 0..steps {
196 pos = entry_into(bytes, pos, &mut buf)?;
197 }
198 Some(String::from_utf8_lossy(&buf).into_owned())
199}
200
201pub fn section_id(bytes: &[u8], meta: &SectionMeta, term: &str) -> Option<u32> {
203 let mut buf = Vec::new();
204 let mut lo = 0usize;
206 let mut hi = meta.restart_offsets.len();
207 while lo < hi {
208 let mid = (lo + hi) / 2;
209 run_entry_into(bytes, meta.restart_offsets[mid] as usize, &mut buf)?;
210 if buf.as_slice() <= term.as_bytes() {
211 lo = mid + 1;
212 } else {
213 hi = mid;
214 }
215 }
216 if lo == 0 {
217 return None; }
219 let run = lo - 1;
220 let mut pos = run_entry_into(bytes, meta.restart_offsets[run] as usize, &mut buf)?;
221 let base_id = (run * meta.restart_interval as usize) as u32 + 1;
222 let run_len = meta.restart_interval.min(
225 meta.term_count
226 .saturating_sub(run as u32 * meta.restart_interval),
227 );
228 for step in 0..run_len {
229 if buf.as_slice() == term.as_bytes() {
230 return Some(base_id + step);
231 }
232 if buf.as_slice() > term.as_bytes() {
233 return None;
234 }
235 if step + 1 < run_len {
236 pos = entry_into(bytes, pos, &mut buf)?;
237 }
238 }
239 None
240}
241
242pub type ChunkLoader = Box<dyn Fn(usize) -> Option<Vec<u8>> + Send + Sync>;
254
255pub type ChunkBulkLoader = Box<dyn Fn(&[usize]) -> Option<Vec<Vec<u8>>> + Send + Sync>;
262
263pub struct SectionChunk {
267 first_run: usize,
268 first_term: Vec<u8>,
271 body_start: u64,
272 data: OnceLock<Vec<u8>>,
273 runs: OnceLock<Vec<usize>>,
279}
280
281impl SectionChunk {
282 pub fn remote(first_run: usize, first_term: Vec<u8>, body_start: u64) -> Self {
284 SectionChunk {
285 first_run,
286 first_term,
287 body_start,
288 data: OnceLock::new(),
289 runs: OnceLock::new(),
290 }
291 }
292
293 pub fn resident(first_run: usize, first_term: Vec<u8>, body_start: u64, data: Vec<u8>) -> Self {
295 let cell = OnceLock::new();
296 let _ = cell.set(data);
297 SectionChunk {
298 first_run,
299 first_term,
300 body_start,
301 data: cell,
302 runs: OnceLock::new(),
303 }
304 }
305
306 fn run_offsets(&self, data: &[u8], restart_interval: usize) -> &[usize] {
313 if let Some(r) = self.runs.get() {
314 return r;
315 }
316 if data.is_empty() {
317 return &[];
318 }
319 self.runs
320 .get_or_init(|| chunk_run_offsets(data, restart_interval))
321 }
322}
323
324fn chunk_run_offsets(data: &[u8], restart_interval: usize) -> Vec<usize> {
327 let mut offs = vec![0usize];
328 let mut pos = 0usize;
329 let mut buf = Vec::new();
330 let mut count = 0usize;
331 while pos < data.len() {
332 let Some(next) = entry_into(data, pos, &mut buf) else {
333 break;
334 };
335 count += 1;
336 pos = next;
337 if pos < data.len() && count.is_multiple_of(restart_interval) {
338 offs.push(pos);
339 }
340 }
341 offs
342}
343
344pub fn run_first_term(bytes: &[u8], off: usize) -> Option<Vec<u8>> {
347 let mut buf = Vec::new();
348 run_entry_into(bytes, off, &mut buf)?;
349 Some(buf)
350}
351
352pub struct ChunkedSection {
355 meta: SectionMeta,
356 chunks: Vec<SectionChunk>,
357 loader: Option<ChunkLoader>,
358 bulk: Option<ChunkBulkLoader>,
359 failed: AtomicBool,
360}
361
362impl ChunkedSection {
363 pub fn local(section_bytes: Vec<u8>) -> Self {
368 let meta = parse_meta(§ion_bytes).unwrap_or(SectionMeta {
369 term_count: 0,
370 restart_interval: 1,
371 restart_offsets: Vec::new(),
372 });
373 let data = OnceLock::new();
374 let _ = data.set(section_bytes);
375 ChunkedSection {
376 meta,
377 chunks: vec![SectionChunk {
378 first_run: 0,
379 first_term: Vec::new(),
380 body_start: 0,
381 data,
382 runs: OnceLock::new(),
383 }],
384 loader: None,
385 bulk: None,
386 failed: AtomicBool::new(false),
387 }
388 }
389
390 pub fn from_parts(
394 meta: SectionMeta,
395 chunks: Vec<SectionChunk>,
396 loader: Option<ChunkLoader>,
397 ) -> Self {
398 ChunkedSection {
399 meta,
400 chunks,
401 loader,
402 bulk: None,
403 failed: AtomicBool::new(false),
404 }
405 }
406
407 pub fn with_bulk_loader(mut self, bulk: ChunkBulkLoader) -> Self {
411 self.bulk = Some(bulk);
412 self
413 }
414
415 pub fn prefetch_all(&self) {
421 self.prefetch_chunks(&(0..self.chunks.len()).collect::<Vec<_>>());
422 }
423
424 pub fn prefetch_chunks(&self, cis: &[usize]) {
431 let Some(bulk) = &self.bulk else { return };
432 let missing: Vec<usize> = cis
433 .iter()
434 .copied()
435 .filter(|&ci| self.chunks.get(ci).is_some_and(|c| c.data.get().is_none()))
436 .collect();
437 if missing.len() < 2 {
438 return;
439 }
440 if let Some(bodies) = bulk(&missing) {
441 if bodies.len() == missing.len() {
442 for (&ci, body) in missing.iter().zip(bodies) {
443 let _ = self.chunks[ci].data.set(body);
444 }
445 }
446 }
447 }
448
449 pub fn meta(&self) -> &SectionMeta {
450 &self.meta
451 }
452
453 pub fn term_count(&self) -> u32 {
454 self.meta.term_count
455 }
456
457 pub fn load_incomplete(&self) -> bool {
460 self.failed.load(Ordering::Relaxed)
461 }
462
463 pub fn reset_load_failure(&self) {
467 self.failed.store(false, Ordering::Relaxed);
468 }
469
470 fn chunk_data(&self, ci: usize) -> &[u8] {
474 let cell = &self.chunks[ci].data;
475 if let Some(d) = cell.get() {
476 return d;
477 }
478 match &self.loader {
479 Some(load) => match load(ci) {
480 Some(bytes) => cell.get_or_init(|| bytes),
481 None => {
482 self.failed.store(true, Ordering::Relaxed);
483 &[]
484 }
485 },
486 None => cell.get_or_init(Vec::new),
487 }
488 }
489
490 fn chunk_of_run(&self, run: usize) -> Option<usize> {
493 let i = self.chunks.partition_point(|c| c.first_run <= run);
494 i.checked_sub(1)
495 }
496
497 pub fn chunk_of_id(&self, id: u32) -> Option<usize> {
501 if id == ABSENT || id > self.meta.term_count {
502 return None;
503 }
504 let run = (id - 1) as usize / self.meta.restart_interval as usize;
505 self.chunk_of_run(run)
506 }
507
508 fn run_off_in_chunk(&self, ci: usize, run: usize, bytes: &[u8], ri: usize) -> Option<usize> {
513 let chunk = &self.chunks[ci];
514 if self.meta.restart_offsets.is_empty() {
515 chunk
516 .run_offsets(bytes, ri)
517 .get(run.checked_sub(chunk.first_run)?)
518 .copied()
519 } else {
520 self.meta
521 .restart_offsets
522 .get(run)?
523 .checked_sub(chunk.body_start)
524 .map(|o| o as usize)
525 }
526 }
527
528 fn run_end_of_chunk(&self, ci: usize, bytes: &[u8], ri: usize) -> usize {
531 if let Some(next) = self.chunks.get(ci + 1) {
532 return next.first_run;
533 }
534 let chunk = &self.chunks[ci];
535 if self.meta.restart_offsets.is_empty() {
536 chunk.first_run + chunk.run_offsets(bytes, ri).len()
537 } else {
538 self.meta.restart_offsets.len()
539 }
540 }
541
542 pub(crate) fn chunk_of_term(&self, id: u32) -> Option<usize> {
546 if id == ABSENT || id > self.meta.term_count {
547 return None;
548 }
549 self.chunk_of_run((id - 1) as usize / self.meta.restart_interval as usize)
550 }
551
552 pub fn term(&self, id: u32) -> Option<String> {
553 if id == ABSENT || id > self.meta.term_count {
554 return None;
555 }
556 let idx = (id - 1) as usize;
557 let ri = self.meta.restart_interval as usize;
558 let run = idx / ri;
559 let steps = idx % ri;
560 let ci = self.chunk_of_run(run)?;
561 let bytes = self.chunk_data(ci);
562 let off = self.run_off_in_chunk(ci, run, bytes, ri)?;
563 let mut buf = Vec::new();
564 let mut pos = run_entry_into(bytes, off, &mut buf)?;
565 for _ in 0..steps {
566 pos = entry_into(bytes, pos, &mut buf)?;
567 }
568 Some(String::from_utf8_lossy(&buf).into_owned())
569 }
570
571 pub fn id(&self, term: &str) -> Option<u32> {
574 if self.chunks.is_empty() {
575 return None;
576 }
577 let ri = self.meta.restart_interval as usize;
578 let ci = if self.chunks.len() == 1 {
581 0
582 } else {
583 let i = self
584 .chunks
585 .partition_point(|c| c.first_term.as_slice() <= term.as_bytes());
586 i.checked_sub(1)?
587 };
588 let first_run = self.chunks[ci].first_run;
589 let bytes = self.chunk_data(ci);
590
591 let run_end = self.run_end_of_chunk(ci, bytes, ri);
593 let mut buf = Vec::new();
594 let mut lo = first_run;
595 let mut hi = run_end;
596 while lo < hi {
597 let mid = (lo + hi) / 2;
598 let off = self.run_off_in_chunk(ci, mid, bytes, ri)?;
599 run_entry_into(bytes, off, &mut buf)?;
600 if buf.as_slice() <= term.as_bytes() {
601 lo = mid + 1;
602 } else {
603 hi = mid;
604 }
605 }
606 if lo == first_run {
607 return None; }
609 let run = lo - 1;
610 let off = self.run_off_in_chunk(ci, run, bytes, ri)?;
611 let mut pos = run_entry_into(bytes, off, &mut buf)?;
612 let base_id = (run * ri) as u32 + 1;
613 let run_len = self.meta.restart_interval.min(
615 self.meta
616 .term_count
617 .saturating_sub(run as u32 * self.meta.restart_interval),
618 );
619 for step in 0..run_len {
620 if buf.as_slice() == term.as_bytes() {
621 return Some(base_id + step);
622 }
623 if buf.as_slice() > term.as_bytes() {
624 return None;
625 }
626 if step + 1 < run_len {
627 pos = entry_into(bytes, pos, &mut buf)?;
628 }
629 }
630 None
631 }
632
633 pub fn raw_section_bytes(&self) -> Vec<u8> {
637 if self.chunks.len() == 1 && self.chunks[0].body_start == 0 {
638 if let Some(bytes) = self.chunks[0].data.get() {
639 return bytes.clone();
640 }
641 }
642 let mut out = encode_section_header(&self.meta);
643 for ci in 0..self.chunks.len() {
644 out.extend_from_slice(self.chunk_data(ci));
645 }
646 out
647 }
648}
649
650pub fn encode_section_header(meta: &SectionMeta) -> Vec<u8> {
655 let body_start = meta.restart_offsets.first().copied().unwrap_or(0);
656 let mut out = Vec::new();
657 write_uvarint(&mut out, meta.term_count as u64);
658 write_uvarint(&mut out, meta.restart_interval as u64);
659 write_uvarint(&mut out, meta.restart_offsets.len() as u64);
660 for off in &meta.restart_offsets {
661 write_uvarint(&mut out, off.saturating_sub(body_start));
662 }
663 out
664}
665
666pub struct DictSection<'a> {
668 bytes: &'a [u8],
669 meta: SectionMeta,
670}
671
672impl<'a> DictSection<'a> {
673 pub fn parse(bytes: &'a [u8]) -> Result<Self, DictError> {
674 Ok(Self {
675 bytes,
676 meta: parse_meta(bytes)?,
677 })
678 }
679
680 pub fn len(&self) -> u32 {
681 self.meta.term_count
682 }
683
684 pub fn is_empty(&self) -> bool {
685 self.meta.term_count == 0
686 }
687
688 pub fn term(&self, id: u32) -> Option<String> {
690 section_term(self.bytes, &self.meta, id)
691 }
692
693 pub fn id(&self, term: &str) -> Option<u32> {
695 section_id(self.bytes, &self.meta, term)
696 }
697}
698
699fn common_prefix_len(a: &str, b: &str) -> usize {
700 a.bytes().zip(b.bytes()).take_while(|(x, y)| x == y).count()
701}
702
703#[cfg(test)]
704mod tests {
705 use super::*;
706
707 fn sample() -> Vec<String> {
708 [
710 "http://ex.org/Alice",
711 "http://ex.org/Bob",
712 "http://ex.org/Alan",
713 "http://ex.org/knows",
714 "http://ex.org/Alice", "zeta",
716 ]
717 .iter()
718 .map(|s| s.to_string())
719 .collect()
720 }
721
722 #[test]
723 fn round_trip_all_ids_and_terms() {
724 for r in [1u32, 2, 16, 1000] {
725 let mut b = DictSectionBuilder::new().with_restart_interval(r);
726 for t in sample() {
727 b.push(t);
728 }
729 let bytes = b.build();
730 let sec = DictSection::parse(&bytes).unwrap();
731
732 let mut expected = sample();
734 expected.sort();
735 expected.dedup();
736 assert_eq!(sec.len() as usize, expected.len());
737
738 for (i, term) in expected.iter().enumerate() {
739 let id = (i + 1) as u32;
740 assert_eq!(
741 sec.term(id).as_deref(),
742 Some(term.as_str()),
743 "term({id}) r={r}"
744 );
745 assert_eq!(sec.id(term), Some(id), "id({term}) r={r}");
746 }
747 }
748 }
749
750 #[test]
751 fn lookups_for_absent_terms() {
752 let mut b = DictSectionBuilder::new();
753 for t in sample() {
754 b.push(t);
755 }
756 let bytes = b.build();
757 let sec = DictSection::parse(&bytes).unwrap();
758 assert_eq!(sec.id("aaa-before-everything"), None);
759 assert_eq!(sec.id("zzz-after-everything"), None);
760 assert_eq!(sec.id("http://ex.org/Alic"), None); assert_eq!(sec.term(0), None);
762 assert_eq!(sec.term(9999), None);
763 }
764
765 #[test]
773 fn randomized_round_trip_and_near_misses_across_restart_intervals() {
774 let mut state = 0x9E37_79B9_7F4A_7C15u64;
775 let mut next = move || {
776 state ^= state << 13;
777 state ^= state >> 7;
778 state ^= state << 17;
779 state
780 };
781 let pool: Vec<String> = (0..1500)
782 .map(|i| {
783 let n = next();
784 match n % 4 {
785 0 => format!("<http://example.org/entity/{n:x}>"),
786 1 => format!("\"literal value {} with spaces\"", n % 300), 2 => format!("<http://example.org/entity/{}/sub/{i}>", n % 64), _ => format!("_:b{}", n % 256),
789 }
790 })
791 .collect();
792 let mut expected = pool.clone();
793 expected.sort();
794 expected.dedup();
795
796 for r in [1u32, 3, 16, 64] {
797 let mut b = DictSectionBuilder::new().with_restart_interval(r);
798 for t in &pool {
799 b.push(t.clone());
800 }
801 let bytes = b.build();
802 let sec = DictSection::parse(&bytes).unwrap();
803 assert_eq!(sec.len() as usize, expected.len(), "r={r}");
804
805 for (i, term) in expected.iter().enumerate() {
806 let id = (i + 1) as u32;
807 assert_eq!(
808 sec.term(id).as_deref(),
809 Some(term.as_str()),
810 "term({id}) r={r}"
811 );
812 assert_eq!(sec.id(term), Some(id), "id({term}) r={r}");
813 }
814 for (i, term) in expected.iter().enumerate() {
818 let near_boundary = (i as u32) % r <= 1;
819 if !near_boundary && i % 37 != 0 {
820 continue;
821 }
822 let longer = format!("{term}\u{1}");
823 assert_eq!(sec.id(&longer), None, "near-miss long r={r}");
824 let mut shorter = term.clone();
825 shorter.pop();
826 if !shorter.is_empty() && expected.binary_search(&shorter).is_err() {
827 assert_eq!(sec.id(&shorter), None, "near-miss short {shorter:?} r={r}");
828 }
829 }
830 assert_eq!(sec.term(expected.len() as u32 + 1), None, "past-end r={r}");
831 }
832 }
833}