1use alloc::vec::Vec;
19
20use plugmem_arena::{
21 Arena, ArenaCfg, BlobHeap, BlobHeapCfg, ChunkPool, ChunkPoolCfg, Interner, ShardMode, Slot,
22};
23
24use crate::config::Config;
25use crate::error::Error;
26use crate::id::{FactId, NONE_U32};
27use crate::index::IdListIndex;
28use crate::index::bm25::Bm25Index;
29use crate::index::hnsw::HnswGraph;
30use crate::index::postings::PostingStore;
31use crate::index::varint::decode_u32;
32use crate::index::vecpool::VecPool;
33use crate::memory::FactFault;
34use crate::model::{EntityRecord, FactAux, FactRecord, TemporalSlot};
35use crate::snapshot::{Prefix, SectionMeta, Snapshot, SnapshotSink, build_prefix, pad_len};
36use xxhash_rust::xxh3::Xxh3;
37
38use super::Memory;
39
40mod kind {
43 pub const FACTS_META: u16 = 1;
44 pub const FACTS_POOL: u16 = 2;
45 pub const AUX_META: u16 = 3;
46 pub const AUX_POOL: u16 = 4;
47 pub const ENTITIES_META: u16 = 5;
48 pub const ENTITIES_POOL: u16 = 6;
49 pub const BY_NAME_META: u16 = 7;
50 pub const BY_NAME_POOL: u16 = 8;
51 pub const EDGES_OUT_META: u16 = 9;
52 pub const EDGES_OUT_POOL: u16 = 10;
53 pub const EDGES_IN_META: u16 = 11;
54 pub const EDGES_IN_POOL: u16 = 12;
55 pub const TEMPORAL_META: u16 = 13;
56 pub const TEMPORAL_POOL: u16 = 14;
57 pub const TEXTS_INDEX: u16 = 15;
58 pub const TEXTS_POOL: u16 = 16;
59 pub const TERMS_INDEX: u16 = 17;
60 pub const TERMS_POOL: u16 = 18;
61 pub const TERMS_TABLE: u16 = 19;
62 pub const TAG_LISTS_META: u16 = 20;
63 pub const TAG_LISTS_POOL: u16 = 21;
64 pub const BM25_HANDLES_META: u16 = 22;
65 pub const BM25_HANDLES_POOL: u16 = 23;
66 pub const BM25_CHUNKS_META: u16 = 24;
67 pub const BM25_CHUNKS_POOL: u16 = 25;
68 pub const BM25_DOCLEN_META: u16 = 26;
69 pub const BM25_DOCLEN_POOL: u16 = 27;
70 pub const TAGS_HANDLES_META: u16 = 28;
71 pub const TAGS_HANDLES_POOL: u16 = 29;
72 pub const TAGS_CHUNKS_META: u16 = 30;
73 pub const TAGS_CHUNKS_POOL: u16 = 31;
74 pub const ENTFACTS_HANDLES_META: u16 = 32;
75 pub const ENTFACTS_HANDLES_POOL: u16 = 33;
76 pub const ENTFACTS_CHUNKS_META: u16 = 34;
77 pub const ENTFACTS_CHUNKS_POOL: u16 = 35;
78 pub const ENGINE_STATE: u16 = 36;
79 pub const VEC_POOL: u16 = 37;
80 pub const HNSW_META: u16 = 38;
81 pub const HNSW_LEVEL0: u16 = 39;
82 pub const HNSW_UPPER_META: u16 = 40;
83 pub const HNSW_UPPER_POOL: u16 = 41;
84 pub const HNSW_LISTS_META: u16 = 42;
85 pub const HNSW_LISTS_POOL: u16 = 43;
86 pub const METAS_INDEX: u16 = 44;
87 pub const METAS_POOL: u16 = 45;
88}
89
90const STATE_LEN: usize = 24;
92
93type SectionFn<'f> = dyn FnMut(u16, &[&[u8]]) -> Result<(), Error> + 'f;
97
98pub(crate) struct Sections<'r, 'a> {
107 pub(crate) facts: &'r Arena<'a, FactRecord>,
108 pub(crate) fact_aux: &'r Arena<'a, FactAux>,
109 pub(crate) entities: &'r Arena<'a, EntityRecord>,
110 pub(crate) temporal: &'r Arena<'a, TemporalSlot>,
111 pub(crate) texts: &'r BlobHeap<'a>,
112 pub(crate) metas: &'r BlobHeap<'a>,
113 pub(crate) tag_lists: &'r ChunkPool<'a>,
114 pub(crate) bm25: &'r Bm25Index<'a>,
115 pub(crate) tags_idx: &'r IdListIndex<'a>,
116 pub(crate) entity_facts: &'r IdListIndex<'a>,
117 pub(crate) vecs: &'r VecPool<'a>,
118 pub(crate) hnsw: &'r HnswGraph<'a>,
119}
120
121fn arena_sections<T: Slot>(a: &Arena<'_, T>) -> (Vec<u8>, Vec<u8>) {
123 let (mut meta, mut pool) = (Vec::new(), Vec::new());
124 a.dump_meta(&mut meta);
125 a.dump_pool(&mut pool);
126 (meta, pool)
127}
128
129fn section<'a>(snap: &Snapshot<'a>, kind: u16) -> Result<&'a [u8], Error> {
131 snap.section(kind)
132 .ok_or(Error::Corrupt("snapshot is missing a required section"))
133}
134
135impl<'a, const TF: bool> PostingStore<'a, TF> {
136 pub(crate) fn dump_sections(&self) -> [Vec<u8>; 4] {
138 let (hm, hp) = (self.handles_meta(), self.handles_pool());
139 let (cm, cp) = (self.chunks_meta(), self.chunks_pool());
140 [hm, hp, cm, cp]
141 }
142
143 pub(crate) fn load_sections(
150 shards: usize,
151 max_bytes: usize,
152 hm: &[u8],
153 hp: &[u8],
154 cm: &[u8],
155 cp: &[u8],
156 ) -> Result<Self, Error> {
157 let handles = Arena::<crate::index::postings::IdListSlot>::load(
158 ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(max_bytes),
159 hm,
160 hp,
161 )?;
162 let pool = ChunkPool::load(ChunkPoolCfg::new().with_max_bytes(max_bytes), cm, cp)?;
163 Self::validate_lists(&handles, &pool)?;
164 Ok(Self::from_parts(handles, pool))
165 }
166
167 pub(crate) fn load_sections_borrowed(
172 shards: usize,
173 max_bytes: usize,
174 hm: &[u8],
175 hp: &'a [u8],
176 cm: &[u8],
177 cp: &'a [u8],
178 ) -> Result<Self, Error> {
179 let handles = Arena::<crate::index::postings::IdListSlot>::load_borrowed(
180 ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(max_bytes),
181 hm,
182 hp,
183 )?;
184 let pool = ChunkPool::load_borrowed(ChunkPoolCfg::new().with_max_bytes(max_bytes), cm, cp)?;
185 Self::validate_lists(&handles, &pool)?;
186 Ok(Self::from_parts(handles, pool))
187 }
188
189 fn validate_lists(
194 handles: &Arena<'_, crate::index::postings::IdListSlot>,
195 pool: &ChunkPool<'_>,
196 ) -> Result<(), Error> {
197 let mut visited = alloc::vec![false; pool.chunks()];
198 for slot in handles.iter() {
199 pool.validate_chain(&slot.handle, &mut visited)?;
200 let mut count = 0u32;
201 let mut last = 0u32;
202 let mut first = true;
203 for chunk in pool.iter(&slot.handle) {
204 let mut cur = chunk;
205 while !cur.is_empty() {
206 let Some((delta, used)) = decode_u32(cur) else {
207 return Err(Error::Corrupt("posting entry is malformed"));
208 };
209 let mut entry_len = used;
210 if TF {
211 if cur.len() < used + 1 {
212 return Err(Error::Corrupt("posting entry is malformed"));
213 }
214 entry_len += 1;
215 }
216 cur = &cur[entry_len..];
217 let id = if first {
218 first = false;
219 delta
220 } else {
221 if delta == 0 {
222 return Err(Error::Corrupt("posting ids are not ascending"));
223 }
224 last.checked_add(delta)
225 .ok_or(Error::Corrupt("posting id overflows"))?
226 };
227 last = id;
228 count += 1;
229 }
230 }
231 if count != slot.count || (count > 0 && last != slot.last) {
232 return Err(Error::Corrupt("posting list disagrees with its handle"));
233 }
234 }
235 if pool.orphan_count(&visited) != 0 {
236 return Err(Error::Corrupt("posting pool has orphan chunks"));
237 }
238 Ok(())
239 }
240}
241
242impl<'a> Bm25Index<'a> {
243 fn dump_pairs(&self) -> [(u16, Vec<u8>); 6] {
245 let [hm, hp, cm, cp] = self.postings().dump_sections();
246 let (dm, dp) = arena_sections(self.doc_len_arena());
247 [
248 (kind::BM25_HANDLES_META, hm),
249 (kind::BM25_HANDLES_POOL, hp),
250 (kind::BM25_CHUNKS_META, cm),
251 (kind::BM25_CHUNKS_POOL, cp),
252 (kind::BM25_DOCLEN_META, dm),
253 (kind::BM25_DOCLEN_POOL, dp),
254 ]
255 }
256
257 fn load_from(snap: &Snapshot<'_>, cfg: &Config) -> Result<Self, Error> {
260 let postings = PostingStore::<true>::load_sections(
261 cfg.shards_postings,
262 cfg.max_bytes,
263 section(snap, kind::BM25_HANDLES_META)?,
264 section(snap, kind::BM25_HANDLES_POOL)?,
265 section(snap, kind::BM25_CHUNKS_META)?,
266 section(snap, kind::BM25_CHUNKS_POOL)?,
267 )?;
268 let doc_len = Arena::load(
269 ArenaCfg::new(cfg.shards_postings, ShardMode::Uniform).with_max_bytes(cfg.max_bytes),
270 section(snap, kind::BM25_DOCLEN_META)?,
271 section(snap, kind::BM25_DOCLEN_POOL)?,
272 )?;
273 Self::assemble(postings, doc_len, snap)
274 }
275
276 fn load_from_borrowed(snap: &Snapshot<'a>, cfg: &Config) -> Result<Self, Error> {
279 let postings = PostingStore::<true>::load_sections_borrowed(
280 cfg.shards_postings,
281 cfg.max_bytes,
282 section(snap, kind::BM25_HANDLES_META)?,
283 section(snap, kind::BM25_HANDLES_POOL)?,
284 section(snap, kind::BM25_CHUNKS_META)?,
285 section(snap, kind::BM25_CHUNKS_POOL)?,
286 )?;
287 let doc_len = Arena::load_borrowed(
288 ArenaCfg::new(cfg.shards_postings, ShardMode::Uniform).with_max_bytes(cfg.max_bytes),
289 section(snap, kind::BM25_DOCLEN_META)?,
290 section(snap, kind::BM25_DOCLEN_POOL)?,
291 )?;
292 Self::assemble(postings, doc_len, snap)
293 }
294
295 fn assemble(
299 postings: PostingStore<'a, true>,
300 doc_len: Arena<'a, crate::index::bm25::DocLenSlot>,
301 snap: &Snapshot<'_>,
302 ) -> Result<Self, Error> {
303 let state = section(snap, kind::ENGINE_STATE)?;
304 if state.len() != STATE_LEN {
305 return Err(Error::Corrupt("engine state section has a wrong length"));
306 }
307 let total_docs = u64::from_le_bytes(state[8..16].try_into().unwrap());
308 let total_len = u64::from_le_bytes(state[16..24].try_into().unwrap());
309 if total_docs != doc_len.len() as u64 {
310 return Err(Error::Corrupt("bm25 document total disagrees with doc_len"));
311 }
312 Ok(Self::from_parts(postings, doc_len, total_docs, total_len))
313 }
314}
315
316impl<'a> Memory<'a> {
317 fn sections(&self) -> Sections<'_, 'a> {
320 Sections {
321 facts: &self.facts,
322 fact_aux: &self.fact_aux,
323 entities: &self.entities,
324 temporal: &self.temporal,
325 texts: &self.texts,
326 metas: &self.metas,
327 tag_lists: &self.tag_lists,
328 bm25: &self.bm25,
329 tags_idx: &self.tags_idx,
330 entity_facts: &self.entity_facts,
331 vecs: &self.vecs,
332 hnsw: &self.hnsw,
333 }
334 }
335
336 fn emit_sections_from(&self, s: &Sections<'_, '_>, f: &mut SectionFn<'_>) -> Result<(), Error> {
346 for (mk, pk, arena) in [
347 (kind::FACTS_META, kind::FACTS_POOL, arena_sections(s.facts)),
348 (kind::AUX_META, kind::AUX_POOL, arena_sections(s.fact_aux)),
349 (
350 kind::ENTITIES_META,
351 kind::ENTITIES_POOL,
352 arena_sections(s.entities),
353 ),
354 (
355 kind::BY_NAME_META,
356 kind::BY_NAME_POOL,
357 arena_sections(&self.by_name),
358 ),
359 (
360 kind::EDGES_OUT_META,
361 kind::EDGES_OUT_POOL,
362 arena_sections(&self.edges_out),
363 ),
364 (
365 kind::EDGES_IN_META,
366 kind::EDGES_IN_POOL,
367 arena_sections(&self.edges_in),
368 ),
369 (
370 kind::TEMPORAL_META,
371 kind::TEMPORAL_POOL,
372 arena_sections(s.temporal),
373 ),
374 ] {
375 let (m, p) = arena;
376 f(mk, &[&m])?;
377 f(pk, &[&p])?;
378 }
379 let (mut i, mut p) = (Vec::new(), Vec::new());
380 s.texts.dump_index(&mut i);
381 s.texts.dump_pool(&mut p);
382 f(kind::TEXTS_INDEX, &[&i])?;
383 f(kind::TEXTS_POOL, &[&p])?;
384 let (mut i, mut p) = (Vec::new(), Vec::new());
385 s.metas.dump_index(&mut i);
386 s.metas.dump_pool(&mut p);
387 f(kind::METAS_INDEX, &[&i])?;
388 f(kind::METAS_POOL, &[&p])?;
389 let (mut i, mut p, mut t) = (Vec::new(), Vec::new(), Vec::new());
390 self.terms.dump_index(&mut i);
391 self.terms.dump_pool(&mut p);
392 self.terms.dump_table(&mut t);
393 f(kind::TERMS_INDEX, &[&i])?;
394 f(kind::TERMS_POOL, &[&p])?;
395 f(kind::TERMS_TABLE, &[&t])?;
396 let (mut m, mut p) = (Vec::new(), Vec::new());
397 s.tag_lists.dump_meta(&mut m);
398 s.tag_lists.dump_pool(&mut p);
399 f(kind::TAG_LISTS_META, &[&m])?;
400 f(kind::TAG_LISTS_POOL, &[&p])?;
401 for (k, bytes) in s.bm25.dump_pairs() {
402 f(k, &[&bytes])?;
403 }
404 let [hm, hp, cm, cp] = s.tags_idx.dump_sections();
405 f(kind::TAGS_HANDLES_META, &[&hm])?;
406 f(kind::TAGS_HANDLES_POOL, &[&hp])?;
407 f(kind::TAGS_CHUNKS_META, &[&cm])?;
408 f(kind::TAGS_CHUNKS_POOL, &[&cp])?;
409 let [hm, hp, cm, cp] = s.entity_facts.dump_sections();
410 f(kind::ENTFACTS_HANDLES_META, &[&hm])?;
411 f(kind::ENTFACTS_HANDLES_POOL, &[&hp])?;
412 f(kind::ENTFACTS_CHUNKS_META, &[&cm])?;
413 f(kind::ENTFACTS_CHUNKS_POOL, &[&cp])?;
414 let mut state = Vec::with_capacity(STATE_LEN);
415 state.extend_from_slice(&self.next_fact.to_le_bytes());
416 state.extend_from_slice(&self.next_entity.to_le_bytes());
417 state.extend_from_slice(&s.bm25.docs().to_le_bytes());
418 state.extend_from_slice(&s.bm25.total_len().to_le_bytes());
419 f(kind::ENGINE_STATE, &[&state])?;
420 f(kind::VEC_POOL, &s.vecs.pieces())?;
423 f(kind::HNSW_META, &[&s.hnsw.dump_meta()])?;
426 f(kind::HNSW_LEVEL0, &[&s.hnsw.dump_level0()])?;
427 let [um, up, lm, lp] = s.hnsw.dump_upper();
428 f(kind::HNSW_UPPER_META, &[&um])?;
429 f(kind::HNSW_UPPER_POOL, &[&up])?;
430 f(kind::HNSW_LISTS_META, &[&lm])?;
431 f(kind::HNSW_LISTS_POOL, &[&lp])?;
432 Ok(())
433 }
434
435 pub fn write_snapshot_to(&self, created_at: u64, sink: impl SnapshotSink) -> Result<(), Error> {
447 self.write_snapshot_with(&self.sections(), created_at, sink)
448 }
449
450 pub(crate) fn write_snapshot_with(
457 &self,
458 s: &Sections<'_, '_>,
459 created_at: u64,
460 mut sink: impl SnapshotSink,
461 ) -> Result<(), Error> {
462 let mut cfg_bytes = Vec::new();
463 self.cfg.encode(&mut cfg_bytes);
464 let flags = if self.cfg.dim > 0 {
465 crate::snapshot::FLAG_VECTORS
466 } else {
467 0
468 };
469
470 let mut metas: Vec<SectionMeta> = Vec::new();
472 self.emit_sections_from(s, &mut |kind, pieces| {
473 let mut h = Xxh3::new();
474 let mut len = 0u64;
475 for p in pieces {
476 h.update(p);
477 len += p.len() as u64;
478 }
479 metas.push(SectionMeta {
480 kind,
481 len,
482 hash: h.digest(),
483 });
484 Ok(())
485 })?;
486
487 let Prefix {
488 bytes: prefix,
489 offsets,
490 file_len: _,
491 } = build_prefix(
492 &cfg_bytes,
493 flags,
494 created_at,
495 env!("CARGO_PKG_VERSION"),
496 &metas,
497 );
498 sink.write(&prefix)?;
499 let mut file_hash = Xxh3::new();
500 file_hash.update(&prefix);
501
502 let zero = [0u8; 64]; let mut idx = 0usize;
505 self.emit_sections_from(s, &mut |_, pieces| {
506 for p in pieces {
507 sink.write(p)?;
508 file_hash.update(p);
509 }
510 let n = pad_len(offsets[idx], metas[idx].len);
511 sink.write(&zero[..n])?;
512 file_hash.update(&zero[..n]);
513 idx += 1;
514 Ok(())
515 })?;
516
517 sink.patch(
518 crate::snapshot::FILE_HASH_OFFSET,
519 &file_hash.digest().to_le_bytes(),
520 )
521 }
522
523 pub fn snapshot_bytes(&self, created_at: u64) -> Vec<u8> {
528 let mut out = Vec::new();
529 self.write_snapshot_to(created_at, &mut out)
530 .expect("writing a snapshot into a Vec is infallible");
531 out
532 }
533
534 pub fn snapshot<S: crate::storage::Storage>(
536 &mut self,
537 store: &mut S,
538 now: u64,
539 ) -> Result<(), Error> {
540 let bytes = self.snapshot_bytes(now);
541 store
542 .write_snapshot(&bytes)
543 .map_err(|e| Error::Storage(alloc::format!("{e:?}")))?;
544 store
545 .clear_journal()
546 .map_err(|e| Error::Storage(alloc::format!("{e:?}")))?;
547 Ok(())
548 }
549
550 pub(super) fn load_snapshot(bytes: &[u8], cfg: Config) -> Result<Self, Error> {
555 cfg.validate()?;
556 let snap = Snapshot::parse(bytes)?;
557 let cfg = Self::reconcile_config(&snap, cfg)?;
558 let mut mem = Self::new(cfg)?;
559 let cfg = &mem.cfg;
560 let uni =
561 |shards: usize| ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(cfg.max_bytes);
562 let ord =
563 |shards: usize| ArenaCfg::new(shards, ShardMode::Ordered).with_max_bytes(cfg.max_bytes);
564 let blob = BlobHeapCfg::new()
565 .with_max_bytes(cfg.max_bytes)
566 .with_max_blob(cfg.max_blob);
567 mem.facts = Arena::load(
568 uni(cfg.shards_facts),
569 section(&snap, kind::FACTS_META)?,
570 section(&snap, kind::FACTS_POOL)?,
571 )?;
572 mem.fact_aux = Arena::load(
573 uni(cfg.shards_facts),
574 section(&snap, kind::AUX_META)?,
575 section(&snap, kind::AUX_POOL)?,
576 )?;
577 mem.entities = Arena::load(
578 uni(cfg.shards_entities),
579 section(&snap, kind::ENTITIES_META)?,
580 section(&snap, kind::ENTITIES_POOL)?,
581 )?;
582 mem.by_name = Arena::load(
583 ord(cfg.shards_entities),
584 section(&snap, kind::BY_NAME_META)?,
585 section(&snap, kind::BY_NAME_POOL)?,
586 )?;
587 mem.edges_out = Arena::load(
588 ord(cfg.shards_edges),
589 section(&snap, kind::EDGES_OUT_META)?,
590 section(&snap, kind::EDGES_OUT_POOL)?,
591 )?;
592 mem.edges_in = Arena::load(
593 ord(cfg.shards_edges),
594 section(&snap, kind::EDGES_IN_META)?,
595 section(&snap, kind::EDGES_IN_POOL)?,
596 )?;
597 mem.temporal = Arena::load(
598 ord(cfg.shards_temporal),
599 section(&snap, kind::TEMPORAL_META)?,
600 section(&snap, kind::TEMPORAL_POOL)?,
601 )?;
602 mem.texts = BlobHeap::load(
603 blob,
604 section(&snap, kind::TEXTS_INDEX)?,
605 section(&snap, kind::TEXTS_POOL)?,
606 )?;
607 mem.metas = BlobHeap::load(
608 blob,
609 section(&snap, kind::METAS_INDEX)?,
610 section(&snap, kind::METAS_POOL)?,
611 )?;
612 mem.terms = Interner::load(
613 blob,
614 section(&snap, kind::TERMS_INDEX)?,
615 section(&snap, kind::TERMS_POOL)?,
616 section(&snap, kind::TERMS_TABLE)?,
617 )?;
618 mem.tag_lists = ChunkPool::load(
619 ChunkPoolCfg::new().with_max_bytes(cfg.max_bytes),
620 section(&snap, kind::TAG_LISTS_META)?,
621 section(&snap, kind::TAG_LISTS_POOL)?,
622 )?;
623 mem.bm25 = Bm25Index::load_from(&snap, cfg)?;
624 mem.tags_idx = IdListIndex::load_sections(
625 cfg.shards_postings,
626 cfg.max_bytes,
627 section(&snap, kind::TAGS_HANDLES_META)?,
628 section(&snap, kind::TAGS_HANDLES_POOL)?,
629 section(&snap, kind::TAGS_CHUNKS_META)?,
630 section(&snap, kind::TAGS_CHUNKS_POOL)?,
631 )?;
632 mem.entity_facts = IdListIndex::load_sections(
633 cfg.shards_entities,
634 cfg.max_bytes,
635 section(&snap, kind::ENTFACTS_HANDLES_META)?,
636 section(&snap, kind::ENTFACTS_HANDLES_POOL)?,
637 section(&snap, kind::ENTFACTS_CHUNKS_META)?,
638 section(&snap, kind::ENTFACTS_CHUNKS_POOL)?,
639 )?;
640 mem.vecs = VecPool::from_parts(cfg.dim, cfg.max_bytes, section(&snap, kind::VEC_POOL)?)?;
641 mem.hnsw = crate::index::hnsw::HnswGraph::from_parts(
642 cfg.hnsw_m,
643 cfg.hnsw_m0,
644 cfg.max_bytes,
645 section(&snap, kind::HNSW_META)?,
646 section(&snap, kind::HNSW_LEVEL0)?,
647 section(&snap, kind::HNSW_UPPER_META)?,
648 section(&snap, kind::HNSW_UPPER_POOL)?,
649 section(&snap, kind::HNSW_LISTS_META)?,
650 section(&snap, kind::HNSW_LISTS_POOL)?,
651 )?;
652 Self::finish_load(mem, &snap)
653 }
654
655 pub(super) fn load_snapshot_borrowed(bytes: &'a [u8], cfg: Config) -> Result<Self, Error> {
663 cfg.validate()?;
664 let snap = Snapshot::parse(bytes)?;
665 let cfg = Self::reconcile_config(&snap, cfg)?;
666 let mut mem = Self::new(cfg)?;
667 let cfg = &mem.cfg;
668 let uni =
669 |shards: usize| ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(cfg.max_bytes);
670 let ord =
671 |shards: usize| ArenaCfg::new(shards, ShardMode::Ordered).with_max_bytes(cfg.max_bytes);
672 let blob = BlobHeapCfg::new()
673 .with_max_bytes(cfg.max_bytes)
674 .with_max_blob(cfg.max_blob);
675 mem.facts = Arena::load_borrowed(
676 uni(cfg.shards_facts),
677 section(&snap, kind::FACTS_META)?,
678 section(&snap, kind::FACTS_POOL)?,
679 )?;
680 mem.fact_aux = Arena::load_borrowed(
681 uni(cfg.shards_facts),
682 section(&snap, kind::AUX_META)?,
683 section(&snap, kind::AUX_POOL)?,
684 )?;
685 mem.entities = Arena::load_borrowed(
686 uni(cfg.shards_entities),
687 section(&snap, kind::ENTITIES_META)?,
688 section(&snap, kind::ENTITIES_POOL)?,
689 )?;
690 mem.by_name = Arena::load_borrowed(
691 ord(cfg.shards_entities),
692 section(&snap, kind::BY_NAME_META)?,
693 section(&snap, kind::BY_NAME_POOL)?,
694 )?;
695 mem.edges_out = Arena::load_borrowed(
696 ord(cfg.shards_edges),
697 section(&snap, kind::EDGES_OUT_META)?,
698 section(&snap, kind::EDGES_OUT_POOL)?,
699 )?;
700 mem.edges_in = Arena::load_borrowed(
701 ord(cfg.shards_edges),
702 section(&snap, kind::EDGES_IN_META)?,
703 section(&snap, kind::EDGES_IN_POOL)?,
704 )?;
705 mem.temporal = Arena::load_borrowed(
706 ord(cfg.shards_temporal),
707 section(&snap, kind::TEMPORAL_META)?,
708 section(&snap, kind::TEMPORAL_POOL)?,
709 )?;
710 mem.texts = BlobHeap::load_borrowed(
711 blob,
712 section(&snap, kind::TEXTS_INDEX)?,
713 section(&snap, kind::TEXTS_POOL)?,
714 )?;
715 mem.metas = BlobHeap::load_borrowed(
716 blob,
717 section(&snap, kind::METAS_INDEX)?,
718 section(&snap, kind::METAS_POOL)?,
719 )?;
720 mem.terms = Interner::load_borrowed(
721 blob,
722 section(&snap, kind::TERMS_INDEX)?,
723 section(&snap, kind::TERMS_POOL)?,
724 section(&snap, kind::TERMS_TABLE)?,
725 )?;
726 mem.tag_lists = ChunkPool::load_borrowed(
727 ChunkPoolCfg::new().with_max_bytes(cfg.max_bytes),
728 section(&snap, kind::TAG_LISTS_META)?,
729 section(&snap, kind::TAG_LISTS_POOL)?,
730 )?;
731 mem.bm25 = Bm25Index::load_from_borrowed(&snap, cfg)?;
732 mem.tags_idx = IdListIndex::load_sections_borrowed(
733 cfg.shards_postings,
734 cfg.max_bytes,
735 section(&snap, kind::TAGS_HANDLES_META)?,
736 section(&snap, kind::TAGS_HANDLES_POOL)?,
737 section(&snap, kind::TAGS_CHUNKS_META)?,
738 section(&snap, kind::TAGS_CHUNKS_POOL)?,
739 )?;
740 mem.entity_facts = IdListIndex::load_sections_borrowed(
741 cfg.shards_entities,
742 cfg.max_bytes,
743 section(&snap, kind::ENTFACTS_HANDLES_META)?,
744 section(&snap, kind::ENTFACTS_HANDLES_POOL)?,
745 section(&snap, kind::ENTFACTS_CHUNKS_META)?,
746 section(&snap, kind::ENTFACTS_CHUNKS_POOL)?,
747 )?;
748 mem.vecs =
749 VecPool::from_parts_borrowed(cfg.dim, cfg.max_bytes, section(&snap, kind::VEC_POOL)?)?;
750 mem.hnsw = crate::index::hnsw::HnswGraph::from_parts_borrowed(
751 cfg.hnsw_m,
752 cfg.hnsw_m0,
753 cfg.max_bytes,
754 section(&snap, kind::HNSW_META)?,
755 section(&snap, kind::HNSW_LEVEL0)?,
756 section(&snap, kind::HNSW_UPPER_META)?,
757 section(&snap, kind::HNSW_UPPER_POOL)?,
758 section(&snap, kind::HNSW_LISTS_META)?,
759 section(&snap, kind::HNSW_LISTS_POOL)?,
760 )?;
761 Self::finish_load(mem, &snap)
762 }
763
764 fn reconcile_config(snap: &Snapshot<'_>, mut cfg: Config) -> Result<Config, Error> {
768 let stored = Config::decode(snap.config())?;
769 if stored.dim != cfg.dim {
770 return Err(Error::ConfigMismatch("stored dim differs"));
771 }
772 if [
773 (stored.shards_facts, cfg.shards_facts),
774 (stored.shards_entities, cfg.shards_entities),
775 (stored.shards_edges, cfg.shards_edges),
776 (stored.shards_temporal, cfg.shards_temporal),
777 (stored.shards_postings, cfg.shards_postings),
778 ]
779 .iter()
780 .any(|&(a, b)| a != b)
781 {
782 return Err(Error::ConfigMismatch("stored shard counts differ"));
783 }
784 if stored.max_bytes != cfg.max_bytes
785 || stored.max_text != cfg.max_text
786 || stored.max_blob != cfg.max_blob
787 {
788 return Err(Error::ConfigMismatch("stored size limits differ"));
789 }
790 if cfg.db_uuid != 0 && stored.db_uuid != cfg.db_uuid {
794 return Err(Error::ConfigMismatch("stored db_uuid differs"));
795 }
796 cfg.db_uuid = stored.db_uuid;
797 Ok(cfg)
798 }
799
800 fn finish_load(mut mem: Self, snap: &Snapshot<'_>) -> Result<Self, Error> {
809 let state = section(snap, kind::ENGINE_STATE)?;
810 if state.len() != STATE_LEN {
811 return Err(Error::Corrupt("engine state section has a wrong length"));
812 }
813 mem.next_fact = u32::from_le_bytes(state[0..4].try_into().unwrap());
814 mem.next_entity = u32::from_le_bytes(state[4..8].try_into().unwrap());
815 if (mem.next_fact as usize) < mem.facts.len()
816 || (mem.next_entity as usize) < mem.entities.len()
817 {
818 return Err(Error::Corrupt("engine id counters below record counts"));
819 }
820 mem.validate_references()?;
821 Ok(mem)
822 }
823
824 fn validate_references(&self) -> Result<(), Error> {
833 let texts = self.texts.len() as u32;
834 let terms = self.terms.len() as u32;
835 self.hnsw.validate(&self.vecs)?;
839 for fact in self.facts.iter() {
840 if fact.id.0 >= self.next_fact
841 || fact.text.0 >= texts
842 || (fact.entity.0 != NONE_U32 && fact.entity.0 >= self.next_entity)
843 || (fact.revises.0 != NONE_U32 && fact.revises.0 >= self.next_fact)
844 || fact.kind != 0
845 {
846 return Err(Error::Corrupt("fact record references out of range"));
847 }
848 if !fact.has_vector() && fact.vector != NONE_U32 {
852 return Err(Error::Corrupt("fact without a vector flag carries a slot"));
853 }
854 }
855 let metas = self.metas.len() as u32;
856 let mut visited = alloc::vec![false; self.tag_lists.chunks()];
857 for aux in self.fact_aux.iter() {
858 if aux.id.0 >= self.next_fact || (aux.meta.0 != NONE_U32 && aux.meta.0 >= metas) {
859 return Err(Error::Corrupt("aux record references out of range"));
860 }
861 self.tag_lists.validate_chain(&aux.tags, &mut visited)?;
862 for chunk in self.tag_lists.iter(&aux.tags) {
863 if !chunk.len().is_multiple_of(4) {
864 return Err(Error::Corrupt("tag list is not a term-id sequence"));
865 }
866 for raw in chunk.chunks_exact(4) {
867 if u32::from_be_bytes(raw.try_into().unwrap()) >= terms {
868 return Err(Error::Corrupt("tag term out of range"));
869 }
870 }
871 }
872 }
873 if self.tag_lists.orphan_count(&visited) != 0 {
874 return Err(Error::Corrupt("tag pool has orphan chunks"));
875 }
876 for entity in self.entities.iter() {
877 if entity.id.0 >= self.next_entity
878 || entity.name.0 >= texts
879 || entity.name_term.0 >= terms
880 {
881 return Err(Error::Corrupt("entity record references out of range"));
882 }
883 }
884 for by_name in self.by_name.iter() {
885 if by_name.name_term.0 >= terms || !self.entities.contains(&by_name.id.0.to_be_bytes())
886 {
887 return Err(Error::Corrupt("by-name record references out of range"));
888 }
889 }
890 for arena in [&self.edges_out, &self.edges_in] {
891 for edge in arena.iter() {
892 if edge.a.0 >= self.next_entity
893 || edge.b.0 >= self.next_entity
894 || edge.rel.0 >= terms
895 || (edge.fact.0 != NONE_U32 && edge.fact.0 >= self.next_fact)
896 || !self.entities.contains(&edge.a.0.to_be_bytes())
897 || !self.entities.contains(&edge.b.0.to_be_bytes())
898 {
899 return Err(Error::Corrupt("edge record references out of range"));
900 }
901 }
902 }
903 for slot in self.temporal.iter() {
904 if slot.fact.0 >= self.next_fact {
905 return Err(Error::Corrupt("temporal record references out of range"));
906 }
907 }
908 Ok(())
909 }
910
911 pub fn verify(&self) -> Result<(), Error> {
932 for (_, text) in self.texts.iter() {
935 if core::str::from_utf8(text).is_err() {
936 return Err(Error::Corrupt("stored text is not valid UTF-8"));
937 }
938 }
939 let mut pairs = Vec::new();
944 for aux in self.fact_aux.iter() {
945 if aux.meta.0 != NONE_U32 {
946 crate::metadata::decode(self.metas.get(aux.meta), &mut pairs)?;
947 }
948 }
949 self.vecs.validate()?;
953 let vslots = self.vecs.len() as u32;
954 let mut with_vec = 0u32;
955 for fact in self.facts.iter() {
956 if fact.has_vector() {
957 if fact.vector >= vslots || self.vecs.slot_fact(fact.vector as usize) != fact.id.0 {
958 return Err(Error::Corrupt(
959 "fact vector slot is out of range or mismatched",
960 ));
961 }
962 with_vec += 1;
963 }
964 }
965 if with_vec != vslots {
966 return Err(Error::Corrupt("vector pool has orphan slots"));
967 }
968 Ok(())
969 }
970
971 pub fn faulty_facts(&self) -> Vec<(FactId, FactFault)> {
983 let vslots = self.vecs.len() as u32;
984 let metas = self.metas.len() as u32;
985 let mut pairs = Vec::new();
986 let mut out = Vec::new();
987 for i in 0..self.next_fact {
988 let id = FactId(i);
989 let Some(record) = self.fact(id) else {
990 continue; };
992 if record.is_tombstone() {
993 continue;
994 }
995 if core::str::from_utf8(self.texts.get(record.text)).is_err() {
996 out.push((id, FactFault::Text));
997 continue;
998 }
999 if record.has_vector()
1000 && (record.vector >= vslots || self.vecs.slot_fact(record.vector as usize) != id.0)
1001 {
1002 out.push((id, FactFault::Vector));
1003 continue;
1004 }
1005 if let Some(aux) = self.fact_aux.get(&id.0.to_be_bytes())
1009 && aux.meta.0 != NONE_U32
1010 && (aux.meta.0 >= metas
1011 || crate::metadata::decode(self.metas.get(aux.meta), &mut pairs).is_err())
1012 {
1013 out.push((id, FactFault::Metadata));
1014 }
1015 }
1016 out
1017 }
1018}