1use std::io::BufWriter;
27use std::ops::Deref;
28use std::sync::Arc;
29
30use memmap2::Mmap;
31use roaring::RoaringBitmap;
32use serde::{Deserialize, Serialize};
33
34use crate::error::{Error, Result};
35use crate::fsutil::{write_atomic, AtomicFile};
36use crate::paths::Paths;
37use crate::trigram::{self, Trigram, TrigramDnf, TrigramQuery};
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct DocMeta {
42 pub path: String,
44 pub lang: String,
46 pub size: u64,
47 pub hash: u64,
49 pub lines: u32,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct SymbolEntry {
55 pub doc_id: u32,
56 pub name: String,
57 pub kind: String,
59 pub line_start: u32,
60 pub line_end: u32,
61 pub container: Option<String>,
67 pub signature: Option<String>,
69}
70
71#[derive(Debug, Clone)]
73pub struct RawSymbol {
74 pub name: String,
75 pub kind: String,
76 pub line_start: u32,
77 pub line_end: u32,
78 pub container: Option<String>,
79 pub signature: Option<String>,
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
86#[serde(rename_all = "lowercase")]
87pub enum RefKind {
88 Call,
89 Import,
90}
91
92impl RefKind {
93 pub fn as_str(self) -> &'static str {
94 match self {
95 RefKind::Call => "call",
96 RefKind::Import => "import",
97 }
98 }
99}
100
101#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct RefEntry {
104 pub doc_id: u32,
105 pub name: String,
106 pub kind: RefKind,
107 pub line: u32,
108 pub column: u32,
109}
110
111#[derive(Debug, Clone)]
113pub struct RawRef {
114 pub name: String,
115 pub kind: RefKind,
116 pub line: u32,
117 pub column: u32,
118}
119
120const OFFSET_BITS: u32 = 40;
127const OFFSET_MASK: u64 = (1 << OFFSET_BITS) - 1;
128const CARD_CAP: u64 = (1 << (64 - OFFSET_BITS)) - 1;
131
132const MAX_GROUP_TRIGRAMS: usize = 4;
137
138fn pack_entry(offset: u64, cardinality: u64) -> Result<u64> {
140 if offset > OFFSET_MASK {
141 return Err(Error::other(format!(
142 "postings blob offset {offset} exceeds the packable maximum"
143 )));
144 }
145 Ok((cardinality.min(CARD_CAP) << OFFSET_BITS) | offset)
146}
147
148fn unpack_offset(value: u64) -> u64 {
149 value & OFFSET_MASK
150}
151
152fn unpack_card(value: u64) -> u64 {
153 value >> OFFSET_BITS
154}
155
156pub struct SegmentWriter {
166 docs: Vec<DocMeta>,
167 syms: crate::table::SymTableBuilder,
168 refs: crate::table::RefTableBuilder,
169 pairs: Vec<u64>,
173}
174
175impl Default for SegmentWriter {
176 fn default() -> Self {
177 Self::new()
178 }
179}
180
181impl SegmentWriter {
182 pub fn new() -> Self {
183 SegmentWriter {
184 docs: Vec::new(),
185 syms: crate::table::SymTableBuilder::new(),
186 refs: crate::table::RefTableBuilder::new(),
187 pairs: Vec::new(),
188 }
189 }
190
191 pub fn is_empty(&self) -> bool {
192 self.docs.is_empty()
193 }
194
195 pub fn doc_count(&self) -> usize {
196 self.docs.len()
197 }
198
199 pub fn symbol_count(&self) -> usize {
200 self.syms.len()
201 }
202
203 pub fn add_doc(
207 &mut self,
208 meta: DocMeta,
209 trigrams: &[Trigram],
210 symbols: Vec<RawSymbol>,
211 refs: Vec<RawRef>,
212 ) -> u32 {
213 let doc_id = self.docs.len() as u32;
214 self.docs.push(meta);
215 self.pairs.extend(
216 trigrams
217 .iter()
218 .map(|t| (u64::from(trigram::key_of(t)) << 32) | u64::from(doc_id)),
219 );
220 for s in &symbols {
221 self.syms
222 .push(
223 doc_id,
224 &s.name,
225 &s.kind,
226 s.line_start,
227 s.line_end,
228 s.container.as_deref(),
229 s.signature.as_deref(),
230 )
231 .expect("writer doc ids are ascending");
232 }
233 for r in &refs {
234 self.refs
235 .push(doc_id, &r.name, r.kind, r.line, r.column)
236 .expect("writer doc ids are ascending");
237 }
238 doc_id
239 }
240
241 pub fn write(self, paths: &Paths, seg_id: u64) -> Result<()> {
243 let SegmentWriter {
244 docs,
245 syms,
246 refs,
247 pairs,
248 } = self;
249 let (postings, tables) = rayon::join(
252 || build_postings_blob(pairs),
253 || rayon::join(|| syms.finish(docs.len()), || refs.finish(docs.len())),
254 );
255 let (post_blob, fst_entries) = postings?;
256 let (syms_enc, refs_enc) = tables;
257 write_segment_files(
258 paths,
259 seg_id,
260 &docs,
261 syms_enc?,
262 refs_enc?,
263 &fst_entries,
264 post_blob,
265 )
266 }
267}
268
269fn append_checksum(buf: &mut Vec<u8>) {
271 let h = xxhash_rust::xxh3::xxh3_64(buf);
272 buf.extend_from_slice(&h.to_le_bytes());
273}
274
275fn verify_checksum<'a>(bytes: &'a [u8], what: &str) -> Result<&'a [u8]> {
277 if bytes.len() < 8 {
278 return Err(Error::Corrupt(format!(
279 "{what}: too short for checksum footer"
280 )));
281 }
282 let (payload, footer) = bytes.split_at(bytes.len() - 8);
283 let want = u64::from_le_bytes(footer.try_into().expect("8-byte footer"));
284 if xxhash_rust::xxh3::xxh3_64(payload) != want {
285 return Err(Error::Corrupt(format!("{what}: checksum mismatch")));
286 }
287 Ok(payload)
288}
289
290fn encode_table<T: Serialize>(rows: &[T]) -> Result<Vec<u8>> {
292 let mut buf = postcard::to_allocvec(rows)?;
293 append_checksum(&mut buf);
294 Ok(buf)
295}
296
297fn read_table<T: serde::de::DeserializeOwned>(
299 path: &std::path::Path,
300 what: &str,
301) -> Result<Vec<T>> {
302 let bytes = std::fs::read(path).map_err(|e| Error::io(path, e))?;
303 Ok(postcard::from_bytes(verify_checksum(&bytes, what)?)?)
304}
305
306pub(crate) fn write_segment_files(
311 paths: &Paths,
312 seg_id: u64,
313 docs: &[DocMeta],
314 syms: crate::table::EncodedTable,
315 refs: crate::table::EncodedTable,
316 fst_entries: &[(Trigram, u64)],
317 mut post_blob: Vec<u8>,
318) -> Result<()> {
319 std::fs::create_dir_all(paths.segments_dir())
320 .map_err(|e| Error::io(paths.segments_dir(), e))?;
321 let fst_path = paths.fst_file(seg_id);
325 let mut fst_out = AtomicFile::create(&fst_path)?;
326 let mut builder = fst::MapBuilder::new(BufWriter::new(fst_out.file()))?;
327 for (tri, value) in fst_entries {
328 builder.insert(tri, *value)?;
329 }
330 builder.finish()?;
331 fst_out.commit()?;
332
333 append_checksum(&mut post_blob);
334 write_atomic(&paths.post_file(seg_id), &post_blob)?;
335 write_atomic(&paths.docs_file(seg_id), &encode_table(docs)?)?;
337 syms.write_atomic(&paths.syms_file(seg_id))?;
338 refs.write_atomic(&paths.refs_file(seg_id))?;
339
340 let mut live = RoaringBitmap::new();
342 live.insert_range(0..docs.len() as u32);
343 write_bitmap(&paths.live_file(seg_id), &live)?;
344
345 Ok(())
346}
347
348pub(crate) type PostingsBlob = (Vec<u8>, Vec<(Trigram, u64)>);
351
352fn build_postings_blob(mut pairs: Vec<u64>) -> Result<PostingsBlob> {
358 use rayon::prelude::*;
359 pairs.par_sort_unstable();
360 if pairs.is_empty() {
361 return Ok((Vec::new(), Vec::new()));
362 }
363
364 let n = pairs.len();
367 let parts = rayon::current_num_threads().clamp(1, 64);
368 let mut bounds: Vec<usize> = vec![0];
369 for p in 1..parts {
370 let mut at = (n * p / parts).max(1);
372 while at < n && (pairs[at - 1] >> 32) == (pairs[at] >> 32) {
373 at += 1;
374 }
375 if at > *bounds.last().expect("non-empty") && at < n {
376 bounds.push(at);
377 }
378 }
379 bounds.push(n);
380
381 type Chunk = (Vec<u8>, Vec<(u32, u64, u64)>);
383 let chunks: Vec<Chunk> = bounds
384 .par_windows(2)
385 .map(|w| {
386 let span = &pairs[w[0]..w[1]];
387 let mut blob: Vec<u8> = Vec::new();
388 let mut entries: Vec<(u32, u64, u64)> = Vec::new();
389 let mut i = 0usize;
390 while i < span.len() {
391 let key = (span[i] >> 32) as u32;
392 let start = i;
393 while i < span.len() && (span[i] >> 32) as u32 == key {
394 i += 1;
395 }
396 let mut bm =
400 RoaringBitmap::from_sorted_iter(span[start..i].iter().map(|&p| p as u32))
401 .map_err(|e| Error::other(format!("postings pairs not sorted: {e}")))?;
402 bm.optimize();
403 let offset = blob.len() as u64;
404 bm.serialize_into(&mut blob)
405 .map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
406 entries.push((key, offset, bm.len()));
407 }
408 Ok((blob, entries))
409 })
410 .collect::<Result<_>>()?;
411 drop(pairs);
412
413 let total: usize = chunks.iter().map(|(b, _)| b.len()).sum();
415 let mut post_blob: Vec<u8> = Vec::with_capacity(total);
416 let mut fst_entries: Vec<(Trigram, u64)> =
417 Vec::with_capacity(chunks.iter().map(|(_, e)| e.len()).sum());
418 for (blob, entries) in chunks {
419 let base = post_blob.len() as u64;
420 post_blob.extend_from_slice(&blob);
421 for (key, offset, card) in entries {
422 fst_entries.push((trigram::tri_of(key), pack_entry(base + offset, card)?));
423 }
424 }
425 Ok((post_blob, fst_entries))
426}
427
428pub(crate) fn merge_postings(segments: &[Segment], remaps: &[Vec<u32>]) -> Result<PostingsBlob> {
436 use fst::Streamer;
437 let mut op = fst::map::OpBuilder::new();
438 for seg in segments {
439 op.push(seg.data.fst.stream());
440 }
441 let mut union = op.union();
442 let mut post_blob: Vec<u8> = Vec::new();
443 let mut fst_entries: Vec<(Trigram, u64)> = Vec::new();
444 while let Some((key, vals)) = union.next() {
445 if key.len() != 3 {
446 continue;
447 }
448 let tri: Trigram = [key[0], key[1], key[2]];
449 let mut out = RoaringBitmap::new();
450 for iv in vals {
451 let seg = &segments[iv.index];
452 let remap = &remaps[iv.index];
453 let bm = seg.data.posting_at(unpack_offset(iv.value))?;
454 for old in bm {
455 if let Some(&new_id) = remap.get(old as usize) {
456 if new_id != u32::MAX {
457 out.insert(new_id);
458 }
459 }
460 }
461 }
462 if out.is_empty() {
463 continue;
464 }
465 out.optimize();
466 let offset = post_blob.len() as u64;
467 out.serialize_into(&mut post_blob)
468 .map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
469 fst_entries.push((tri, pack_entry(offset, out.len())?));
470 }
471 Ok((post_blob, fst_entries))
472}
473
474pub struct SegmentData {
483 fst: fst::Map<Mmap>,
484 post: Mmap,
485 post_len: usize,
488 pub docs: Vec<DocMeta>,
489 syms: crate::table::SymTable,
493 refs: Option<crate::table::RefTable>,
495}
496
497pub struct Segment {
500 pub id: u64,
501 data: Arc<SegmentData>,
502 live: RoaringBitmap,
503}
504
505impl Deref for Segment {
506 type Target = SegmentData;
507 fn deref(&self) -> &SegmentData {
508 &self.data
509 }
510}
511
512impl Segment {
513 pub fn open(paths: &Paths, seg_id: u64) -> Result<Segment> {
514 let fst_path = paths.fst_file(seg_id);
515 let fst_file = std::fs::File::open(&fst_path).map_err(|e| Error::io(&fst_path, e))?;
516 let fst_mmap = unsafe { Mmap::map(&fst_file).map_err(|e| Error::io(&fst_path, e))? };
517 let fst = fst::Map::new(fst_mmap)?;
518 fst.as_fst()
525 .verify()
526 .map_err(|e| Error::Corrupt(format!("fst checksum: {e}")))?;
527
528 let post_path = paths.post_file(seg_id);
529 let post_file = std::fs::File::open(&post_path).map_err(|e| Error::io(&post_path, e))?;
530 let post = unsafe { Mmap::map(&post_file).map_err(|e| Error::io(&post_path, e))? };
531 let post_len = verify_checksum(&post, "postings blob")?.len();
535
536 let docs: Vec<DocMeta> = read_table(&paths.docs_file(seg_id), "docs table")?;
537 let syms = crate::table::SymTable::open(&paths.syms_file(seg_id))?;
538
539 let refs_path = paths.refs_file(seg_id);
542 let refs = match crate::table::RefTable::open(&refs_path) {
543 Ok(t) => Some(t),
544 Err(Error::Io { source, .. }) if source.kind() == std::io::ErrorKind::NotFound => None,
545 Err(e) => return Err(e),
546 };
547
548 if syms.doc_count() != docs.len()
552 || refs.as_ref().is_some_and(|r| r.doc_count() != docs.len())
553 {
554 return Err(Error::Corrupt(format!(
555 "segment {seg_id}: side-table doc count does not match docs table"
556 )));
557 }
558
559 let live = read_bitmap(&paths.live_file(seg_id))?;
560
561 Ok(Segment {
562 id: seg_id,
563 data: Arc::new(SegmentData {
564 fst,
565 post,
566 post_len,
567 docs,
568 syms,
569 refs,
570 }),
571 live,
572 })
573 }
574
575 pub fn reopen(&self, paths: &Paths) -> Result<Segment> {
579 let live = read_bitmap(&paths.live_file(self.id))?;
580 Ok(Segment {
581 id: self.id,
582 data: self.data.clone(),
583 live,
584 })
585 }
586
587 pub fn is_live(&self, doc_id: u32) -> bool {
588 self.live.contains(doc_id)
589 }
590
591 pub fn subtract_live(&mut self, doc_ids: &[u32]) {
596 for &id in doc_ids {
597 self.live.remove(id);
598 }
599 }
600
601 pub fn live_count(&self) -> u64 {
602 self.live.len()
603 }
604
605 pub fn all_live(&self) -> RoaringBitmap {
607 self.live.clone()
608 }
609
610 pub fn candidates(&self, query: &TrigramQuery) -> Result<RoaringBitmap> {
613 let mut filtering = query
614 .dnfs
615 .iter()
616 .filter(|d| trigram::dnf_filters(d))
617 .peekable();
618 if filtering.peek().is_none() {
619 return Ok(self.all_live());
620 }
621 let mut result: Option<RoaringBitmap> = None;
622 for dnf in filtering {
623 let bm = self.data.dnf_bitmap(dnf)?;
624 result = Some(match result.take() {
625 None => bm,
626 Some(a) => a & bm,
627 });
628 if result.as_ref().is_some_and(|b| b.is_empty()) {
629 break;
630 }
631 }
632 let mut out = result.unwrap_or_default();
633 out &= &self.live;
634 Ok(out)
635 }
636}
637
638impl SegmentData {
639 pub fn doc(&self, doc_id: u32) -> Option<&DocMeta> {
640 self.docs.get(doc_id as usize)
641 }
642
643 pub fn sym_count(&self) -> usize {
645 self.syms.len()
646 }
647
648 pub fn sym(&self, i: u32) -> Option<SymbolEntry> {
650 self.syms.get(i)
651 }
652
653 pub fn sym_names(&self) -> impl Iterator<Item = (u32, &str, &str)> {
657 self.syms.names()
658 }
659
660 pub fn doc_syms(&self, doc_id: u32) -> impl Iterator<Item = SymbolEntry> + '_ {
662 self.syms
663 .doc_range(doc_id)
664 .filter_map(move |i| self.syms.get(i))
665 }
666
667 pub fn doc_sym_count(&self, doc_id: u32) -> u32 {
669 let r = self.syms.doc_range(doc_id);
670 r.end - r.start
671 }
672
673 pub fn doc_refs(&self, doc_id: u32) -> impl Iterator<Item = RefEntry> + '_ {
675 self.refs
676 .iter()
677 .flat_map(move |t| t.doc_range(doc_id).filter_map(move |i| t.get(i)))
678 }
679
680 pub fn refs_named<'s>(&'s self, name: &'s str) -> impl Iterator<Item = RefEntry> + 's {
684 self.refs
685 .iter()
686 .flat_map(move |t| t.rows_named(name).filter_map(move |i| t.get(i)))
687 }
688
689 pub fn calls_to<'s>(&'s self, name: &'s str) -> impl Iterator<Item = RefEntry> + 's {
691 self.refs_named(name).filter(|r| r.kind == RefKind::Call)
692 }
693
694 pub fn syms_by_lower<'s>(&'s self, lower: &'s str) -> impl Iterator<Item = u32> + 's {
697 self.syms.rows_named(lower)
698 }
699
700 pub fn sym_name(&self, i: u32) -> &str {
702 self.syms.name(i)
703 }
704
705 pub fn sym_name_lower(&self, i: u32) -> &str {
707 self.syms.name_lower(i)
708 }
709
710 fn posting_entry(&self, tri: Trigram) -> Option<u64> {
712 self.fst.get(tri)
713 }
714
715 fn posting_at(&self, offset: u64) -> Result<RoaringBitmap> {
722 let start = offset as usize;
723 let slice = self.post.get(start..self.post_len).ok_or_else(|| {
724 Error::Corrupt(format!(
725 "posting offset {start} out of range for postings blob of length {}",
726 self.post_len
727 ))
728 })?;
729 RoaringBitmap::deserialize_from(slice)
730 .map_err(|e| Error::Corrupt(format!("posting list: {e}")))
731 }
732
733 fn dnf_bitmap(&self, dnf: &TrigramDnf) -> Result<RoaringBitmap> {
735 let mut acc = RoaringBitmap::new();
736 for group in dnf {
737 acc |= self.group_bitmap(group)?;
738 }
739 Ok(acc)
740 }
741
742 fn group_bitmap(&self, group: &[Trigram]) -> Result<RoaringBitmap> {
749 let mut entries: Vec<u64> = Vec::with_capacity(group.len());
750 for tri in group {
751 match self.posting_entry(*tri) {
752 Some(v) => entries.push(v),
753 None => return Ok(RoaringBitmap::new()),
755 }
756 }
757 entries.sort_unstable_by_key(|&v| unpack_card(v));
758 entries.truncate(MAX_GROUP_TRIGRAMS);
759
760 let mut acc: Option<RoaringBitmap> = None;
761 for v in entries {
762 let bm = self.posting_at(unpack_offset(v))?;
763 acc = Some(match acc.take() {
764 None => bm,
765 Some(a) => a & bm,
766 });
767 if acc.as_ref().is_some_and(|b| b.is_empty()) {
768 break;
769 }
770 }
771 Ok(acc.unwrap_or_default())
772 }
773}
774
775pub(crate) fn write_bitmap(path: &std::path::Path, bm: &RoaringBitmap) -> Result<()> {
776 let mut buf = Vec::with_capacity(bm.serialized_size() + 8);
777 bm.serialize_into(&mut buf)
778 .map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
779 append_checksum(&mut buf);
780 write_atomic(path, &buf)
781}
782
783pub(crate) fn read_bitmap(path: &std::path::Path) -> Result<RoaringBitmap> {
784 let bytes = std::fs::read(path).map_err(|e| Error::io(path, e))?;
785 RoaringBitmap::deserialize_from(verify_checksum(&bytes, "live bitmap")?)
786 .map_err(|e| Error::Corrupt(format!("live bitmap: {e}")))
787}