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 trigram_keys: &[u32],
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 trigram_keys
217 .iter()
218 .map(|&k| (u64::from(k) << 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_and_post, tables) = rayon::join(
522 || -> Result<(fst::Map<Mmap>, Mmap, usize)> {
523 let fst_path = paths.fst_file(seg_id);
524 let (fst, post_and_len) = rayon::join(
525 || -> Result<fst::Map<Mmap>> {
526 let fst_file =
527 std::fs::File::open(&fst_path).map_err(|e| Error::io(&fst_path, e))?;
528 let fst_mmap =
529 unsafe { Mmap::map(&fst_file).map_err(|e| Error::io(&fst_path, e))? };
530 let fst = fst::Map::new(fst_mmap)?;
531 fst.as_fst()
539 .verify()
540 .map_err(|e| Error::Corrupt(format!("fst checksum: {e}")))?;
541 Ok(fst)
542 },
543 || -> Result<(Mmap, usize)> {
544 let post_path = paths.post_file(seg_id);
545 let post_file =
546 std::fs::File::open(&post_path).map_err(|e| Error::io(&post_path, e))?;
547 let post = unsafe {
548 Mmap::map(&post_file).map_err(|e| Error::io(&post_path, e))?
549 };
550 let post_len = verify_checksum(&post, "postings blob")?.len();
555 Ok((post, post_len))
556 },
557 );
558 let (post, post_len) = post_and_len?;
559 Ok((fst?, post, post_len))
560 },
561 || -> Result<(
562 Vec<DocMeta>,
563 crate::table::SymTable,
564 Option<crate::table::RefTable>,
565 )> {
566 let (docs_and_syms, refs) = rayon::join(
567 || -> Result<(Vec<DocMeta>, crate::table::SymTable)> {
568 let docs: Vec<DocMeta> =
569 read_table(&paths.docs_file(seg_id), "docs table")?;
570 let syms = crate::table::SymTable::open(&paths.syms_file(seg_id))?;
571 Ok((docs, syms))
572 },
573 || -> Result<Option<crate::table::RefTable>> {
574 let refs_path = paths.refs_file(seg_id);
578 match crate::table::RefTable::open(&refs_path) {
579 Ok(t) => Ok(Some(t)),
580 Err(Error::Io { source, .. })
581 if source.kind() == std::io::ErrorKind::NotFound =>
582 {
583 Ok(None)
584 }
585 Err(e) => Err(e),
586 }
587 },
588 );
589 let (docs, syms) = docs_and_syms?;
590 Ok((docs, syms, refs?))
591 },
592 );
593 let (fst, post, post_len) = fst_and_post?;
594 let (docs, syms, refs) = tables?;
595
596 if syms.doc_count() != docs.len()
600 || refs.as_ref().is_some_and(|r| r.doc_count() != docs.len())
601 {
602 return Err(Error::Corrupt(format!(
603 "segment {seg_id}: side-table doc count does not match docs table"
604 )));
605 }
606
607 let live = read_bitmap(&paths.live_file(seg_id))?;
608
609 Ok(Segment {
610 id: seg_id,
611 data: Arc::new(SegmentData {
612 fst,
613 post,
614 post_len,
615 docs,
616 syms,
617 refs,
618 }),
619 live,
620 })
621 }
622
623 pub fn reopen(&self, paths: &Paths) -> Result<Segment> {
627 let live = read_bitmap(&paths.live_file(self.id))?;
628 Ok(Segment {
629 id: self.id,
630 data: self.data.clone(),
631 live,
632 })
633 }
634
635 pub fn is_live(&self, doc_id: u32) -> bool {
636 self.live.contains(doc_id)
637 }
638
639 pub fn subtract_live(&mut self, doc_ids: &[u32]) {
644 for &id in doc_ids {
645 self.live.remove(id);
646 }
647 }
648
649 pub fn live_count(&self) -> u64 {
650 self.live.len()
651 }
652
653 pub fn all_live(&self) -> RoaringBitmap {
655 self.live.clone()
656 }
657
658 pub fn candidates(&self, query: &TrigramQuery) -> Result<RoaringBitmap> {
661 let mut filtering = query
662 .dnfs
663 .iter()
664 .filter(|d| trigram::dnf_filters(d))
665 .peekable();
666 if filtering.peek().is_none() {
667 return Ok(self.all_live());
668 }
669 let mut result: Option<RoaringBitmap> = None;
670 for dnf in filtering {
671 let bm = self.data.dnf_bitmap(dnf)?;
672 result = Some(match result.take() {
673 None => bm,
674 Some(a) => a & bm,
675 });
676 if result.as_ref().is_some_and(|b| b.is_empty()) {
677 break;
678 }
679 }
680 let mut out = result.unwrap_or_default();
681 out &= &self.live;
682 Ok(out)
683 }
684}
685
686impl SegmentData {
687 pub fn doc(&self, doc_id: u32) -> Option<&DocMeta> {
688 self.docs.get(doc_id as usize)
689 }
690
691 pub fn doc_path_score(&self, doc_id: u32) -> f32 {
698 self.doc(doc_id)
699 .map(|d| crate::search::path_score(&d.path))
700 .unwrap_or(0.0)
701 }
702
703 pub fn sym_count(&self) -> usize {
705 self.syms.len()
706 }
707
708 pub fn sym(&self, i: u32) -> Option<SymbolEntry> {
710 self.syms.get(i)
711 }
712
713 pub(crate) fn sym_view(&self, i: u32) -> Option<crate::table::SymView<'_>> {
717 self.syms.view(i)
718 }
719
720 pub(crate) fn doc_sym_rows(&self, doc_id: u32) -> std::ops::Range<u32> {
722 self.syms.doc_range(doc_id)
723 }
724
725 pub(crate) fn doc_sym_views(
727 &self,
728 doc_id: u32,
729 ) -> impl Iterator<Item = crate::table::SymView<'_>> {
730 self.syms
731 .doc_range(doc_id)
732 .filter_map(move |i| self.syms.view(i))
733 }
734
735 pub(crate) fn doc_ref_views(
737 &self,
738 doc_id: u32,
739 ) -> impl Iterator<Item = crate::table::RefView<'_>> {
740 self.refs
741 .iter()
742 .flat_map(move |t| t.doc_range(doc_id).filter_map(move |i| t.view(i)))
743 }
744
745 pub(crate) fn ref_views_named<'s>(
747 &'s self,
748 name: &'s str,
749 ) -> impl Iterator<Item = crate::table::RefView<'s>> + 's {
750 self.refs
751 .iter()
752 .flat_map(move |t| t.rows_named(name).filter_map(move |i| t.view(i)))
753 }
754
755 pub fn doc_syms(&self, doc_id: u32) -> impl Iterator<Item = SymbolEntry> + '_ {
757 self.syms
758 .doc_range(doc_id)
759 .filter_map(move |i| self.syms.get(i))
760 }
761
762 pub fn doc_sym_count(&self, doc_id: u32) -> u32 {
764 let r = self.syms.doc_range(doc_id);
765 r.end - r.start
766 }
767
768 pub fn doc_refs(&self, doc_id: u32) -> impl Iterator<Item = RefEntry> + '_ {
770 self.refs
771 .iter()
772 .flat_map(move |t| t.doc_range(doc_id).filter_map(move |i| t.get(i)))
773 }
774
775 pub fn refs_named<'s>(&'s self, name: &'s str) -> impl Iterator<Item = RefEntry> + 's {
779 self.refs
780 .iter()
781 .flat_map(move |t| t.rows_named(name).filter_map(move |i| t.get(i)))
782 }
783
784 pub fn calls_to<'s>(&'s self, name: &'s str) -> impl Iterator<Item = RefEntry> + 's {
786 self.refs_named(name).filter(|r| r.kind == RefKind::Call)
787 }
788
789 pub fn syms_by_lower<'s>(&'s self, lower: &'s str) -> impl Iterator<Item = u32> + 's {
792 self.syms.rows_named(lower)
793 }
794
795 pub fn sym_name(&self, i: u32) -> &str {
797 self.syms.name(i)
798 }
799
800 pub fn sym_name_lower(&self, i: u32) -> &str {
802 self.syms.name_lower(i)
803 }
804
805 fn posting_entry(&self, tri: Trigram) -> Option<u64> {
807 self.fst.get(tri)
808 }
809
810 fn posting_at(&self, offset: u64) -> Result<RoaringBitmap> {
817 let start = offset as usize;
818 let slice = self.post.get(start..self.post_len).ok_or_else(|| {
819 Error::Corrupt(format!(
820 "posting offset {start} out of range for postings blob of length {}",
821 self.post_len
822 ))
823 })?;
824 RoaringBitmap::deserialize_from(slice)
825 .map_err(|e| Error::Corrupt(format!("posting list: {e}")))
826 }
827
828 fn dnf_bitmap(&self, dnf: &TrigramDnf) -> Result<RoaringBitmap> {
830 let mut acc = RoaringBitmap::new();
831 for group in dnf {
832 acc |= self.group_bitmap(group)?;
833 }
834 Ok(acc)
835 }
836
837 fn group_bitmap(&self, group: &[Trigram]) -> Result<RoaringBitmap> {
844 let mut entries: Vec<u64> = Vec::with_capacity(group.len());
845 for tri in group {
846 match self.posting_entry(*tri) {
847 Some(v) => entries.push(v),
848 None => return Ok(RoaringBitmap::new()),
850 }
851 }
852 entries.sort_unstable_by_key(|&v| unpack_card(v));
853 entries.truncate(MAX_GROUP_TRIGRAMS);
854
855 let mut acc: Option<RoaringBitmap> = None;
856 for v in entries {
857 let bm = self.posting_at(unpack_offset(v))?;
858 acc = Some(match acc.take() {
859 None => bm,
860 Some(a) => a & bm,
861 });
862 if acc.as_ref().is_some_and(|b| b.is_empty()) {
863 break;
864 }
865 }
866 Ok(acc.unwrap_or_default())
867 }
868}
869
870pub(crate) fn write_bitmap(path: &std::path::Path, bm: &RoaringBitmap) -> Result<()> {
871 let mut buf = Vec::with_capacity(bm.serialized_size() + 8);
872 bm.serialize_into(&mut buf)
873 .map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
874 append_checksum(&mut buf);
875 write_atomic(path, &buf)
876}
877
878pub(crate) fn read_bitmap(path: &std::path::Path) -> Result<RoaringBitmap> {
879 let bytes = std::fs::read(path).map_err(|e| Error::io(path, e))?;
880 RoaringBitmap::deserialize_from(verify_checksum(&bytes, "live bitmap")?)
881 .map_err(|e| Error::Corrupt(format!("live bitmap: {e}")))
882}