use std::io::BufWriter;
use std::ops::Deref;
use std::sync::Arc;
use memmap2::Mmap;
use roaring::RoaringBitmap;
use serde::{Deserialize, Serialize};
use crate::error::{Error, Result};
use crate::fsutil::{write_atomic, AtomicFile};
use crate::paths::Paths;
use crate::trigram::{self, Trigram, TrigramDnf, TrigramQuery};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DocMeta {
pub path: String,
pub lang: String,
pub size: u64,
pub hash: u64,
pub lines: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolEntry {
pub doc_id: u32,
pub name: String,
pub kind: String,
pub line_start: u32,
pub line_end: u32,
pub container: Option<String>,
pub signature: Option<String>,
}
#[derive(Debug, Clone)]
pub struct RawSymbol {
pub name: String,
pub kind: String,
pub line_start: u32,
pub line_end: u32,
pub container: Option<String>,
pub signature: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RefKind {
Call,
Import,
}
impl RefKind {
pub fn as_str(self) -> &'static str {
match self {
RefKind::Call => "call",
RefKind::Import => "import",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RefEntry {
pub doc_id: u32,
pub name: String,
pub kind: RefKind,
pub line: u32,
pub column: u32,
}
#[derive(Debug, Clone)]
pub struct RawRef {
pub name: String,
pub kind: RefKind,
pub line: u32,
pub column: u32,
}
const OFFSET_BITS: u32 = 40;
const OFFSET_MASK: u64 = (1 << OFFSET_BITS) - 1;
const CARD_CAP: u64 = (1 << (64 - OFFSET_BITS)) - 1;
const MAX_GROUP_TRIGRAMS: usize = 4;
fn pack_entry(offset: u64, cardinality: u64) -> Result<u64> {
if offset > OFFSET_MASK {
return Err(Error::other(format!(
"postings blob offset {offset} exceeds the packable maximum"
)));
}
Ok((cardinality.min(CARD_CAP) << OFFSET_BITS) | offset)
}
fn unpack_offset(value: u64) -> u64 {
value & OFFSET_MASK
}
fn unpack_card(value: u64) -> u64 {
value >> OFFSET_BITS
}
pub struct SegmentWriter {
docs: Vec<DocMeta>,
syms: crate::table::SymTableBuilder,
refs: crate::table::RefTableBuilder,
pairs: Vec<u64>,
}
impl Default for SegmentWriter {
fn default() -> Self {
Self::new()
}
}
impl SegmentWriter {
pub fn new() -> Self {
SegmentWriter {
docs: Vec::new(),
syms: crate::table::SymTableBuilder::new(),
refs: crate::table::RefTableBuilder::new(),
pairs: Vec::new(),
}
}
pub fn is_empty(&self) -> bool {
self.docs.is_empty()
}
pub fn doc_count(&self) -> usize {
self.docs.len()
}
pub fn symbol_count(&self) -> usize {
self.syms.len()
}
pub fn add_doc(
&mut self,
meta: DocMeta,
trigram_keys: &[u32],
symbols: Vec<RawSymbol>,
refs: Vec<RawRef>,
) -> u32 {
let doc_id = self.docs.len() as u32;
self.docs.push(meta);
self.pairs.extend(
trigram_keys
.iter()
.map(|&k| (u64::from(k) << 32) | u64::from(doc_id)),
);
for s in &symbols {
self.syms
.push(
doc_id,
&s.name,
&s.kind,
s.line_start,
s.line_end,
s.container.as_deref(),
s.signature.as_deref(),
)
.expect("writer doc ids are ascending");
}
for r in &refs {
self.refs
.push(doc_id, &r.name, r.kind, r.line, r.column)
.expect("writer doc ids are ascending");
}
doc_id
}
pub fn write(self, paths: &Paths, seg_id: u64) -> Result<()> {
let SegmentWriter {
docs,
syms,
refs,
pairs,
} = self;
let (postings, tables) = rayon::join(
|| build_postings_blob(pairs),
|| rayon::join(|| syms.finish(docs.len()), || refs.finish(docs.len())),
);
let (post_blob, fst_entries) = postings?;
let (syms_enc, refs_enc) = tables;
write_segment_files(
paths,
seg_id,
&docs,
syms_enc?,
refs_enc?,
&fst_entries,
post_blob,
)
}
}
fn append_checksum(buf: &mut Vec<u8>) {
let h = xxhash_rust::xxh3::xxh3_64(buf);
buf.extend_from_slice(&h.to_le_bytes());
}
fn verify_checksum<'a>(bytes: &'a [u8], what: &str) -> Result<&'a [u8]> {
if bytes.len() < 8 {
return Err(Error::Corrupt(format!(
"{what}: too short for checksum footer"
)));
}
let (payload, footer) = bytes.split_at(bytes.len() - 8);
let want = u64::from_le_bytes(footer.try_into().expect("8-byte footer"));
if xxhash_rust::xxh3::xxh3_64(payload) != want {
return Err(Error::Corrupt(format!("{what}: checksum mismatch")));
}
Ok(payload)
}
fn encode_table<T: Serialize>(rows: &[T]) -> Result<Vec<u8>> {
let mut buf = postcard::to_allocvec(rows)?;
append_checksum(&mut buf);
Ok(buf)
}
fn read_table<T: serde::de::DeserializeOwned>(
path: &std::path::Path,
what: &str,
) -> Result<Vec<T>> {
let bytes = std::fs::read(path).map_err(|e| Error::io(path, e))?;
Ok(postcard::from_bytes(verify_checksum(&bytes, what)?)?)
}
pub(crate) fn write_segment_files(
paths: &Paths,
seg_id: u64,
docs: &[DocMeta],
syms: crate::table::EncodedTable,
refs: crate::table::EncodedTable,
fst_entries: &[(Trigram, u64)],
mut post_blob: Vec<u8>,
) -> Result<()> {
std::fs::create_dir_all(paths.segments_dir())
.map_err(|e| Error::io(paths.segments_dir(), e))?;
let fst_path = paths.fst_file(seg_id);
let mut fst_out = AtomicFile::create(&fst_path)?;
let mut builder = fst::MapBuilder::new(BufWriter::new(fst_out.file()))?;
for (tri, value) in fst_entries {
builder.insert(tri, *value)?;
}
builder.finish()?;
fst_out.commit()?;
append_checksum(&mut post_blob);
write_atomic(&paths.post_file(seg_id), &post_blob)?;
write_atomic(&paths.docs_file(seg_id), &encode_table(docs)?)?;
syms.write_atomic(&paths.syms_file(seg_id))?;
refs.write_atomic(&paths.refs_file(seg_id))?;
let mut live = RoaringBitmap::new();
live.insert_range(0..docs.len() as u32);
write_bitmap(&paths.live_file(seg_id), &live)?;
Ok(())
}
pub(crate) type PostingsBlob = (Vec<u8>, Vec<(Trigram, u64)>);
fn build_postings_blob(mut pairs: Vec<u64>) -> Result<PostingsBlob> {
use rayon::prelude::*;
pairs.par_sort_unstable();
if pairs.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let n = pairs.len();
let parts = rayon::current_num_threads().clamp(1, 64);
let mut bounds: Vec<usize> = vec![0];
for p in 1..parts {
let mut at = (n * p / parts).max(1);
while at < n && (pairs[at - 1] >> 32) == (pairs[at] >> 32) {
at += 1;
}
if at > *bounds.last().expect("non-empty") && at < n {
bounds.push(at);
}
}
bounds.push(n);
type Chunk = (Vec<u8>, Vec<(u32, u64, u64)>);
let chunks: Vec<Chunk> = bounds
.par_windows(2)
.map(|w| {
let span = &pairs[w[0]..w[1]];
let mut blob: Vec<u8> = Vec::new();
let mut entries: Vec<(u32, u64, u64)> = Vec::new();
let mut i = 0usize;
while i < span.len() {
let key = (span[i] >> 32) as u32;
let start = i;
while i < span.len() && (span[i] >> 32) as u32 == key {
i += 1;
}
let mut bm =
RoaringBitmap::from_sorted_iter(span[start..i].iter().map(|&p| p as u32))
.map_err(|e| Error::other(format!("postings pairs not sorted: {e}")))?;
bm.optimize();
let offset = blob.len() as u64;
bm.serialize_into(&mut blob)
.map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
entries.push((key, offset, bm.len()));
}
Ok((blob, entries))
})
.collect::<Result<_>>()?;
drop(pairs);
let total: usize = chunks.iter().map(|(b, _)| b.len()).sum();
let mut post_blob: Vec<u8> = Vec::with_capacity(total);
let mut fst_entries: Vec<(Trigram, u64)> =
Vec::with_capacity(chunks.iter().map(|(_, e)| e.len()).sum());
for (blob, entries) in chunks {
let base = post_blob.len() as u64;
post_blob.extend_from_slice(&blob);
for (key, offset, card) in entries {
fst_entries.push((trigram::tri_of(key), pack_entry(base + offset, card)?));
}
}
Ok((post_blob, fst_entries))
}
pub(crate) fn merge_postings(segments: &[Segment], remaps: &[Vec<u32>]) -> Result<PostingsBlob> {
use fst::Streamer;
let mut op = fst::map::OpBuilder::new();
for seg in segments {
op.push(seg.data.fst.stream());
}
let mut union = op.union();
let mut post_blob: Vec<u8> = Vec::new();
let mut fst_entries: Vec<(Trigram, u64)> = Vec::new();
while let Some((key, vals)) = union.next() {
if key.len() != 3 {
continue;
}
let tri: Trigram = [key[0], key[1], key[2]];
let mut out = RoaringBitmap::new();
for iv in vals {
let seg = &segments[iv.index];
let remap = &remaps[iv.index];
let bm = seg.data.posting_at(unpack_offset(iv.value))?;
for old in bm {
if let Some(&new_id) = remap.get(old as usize) {
if new_id != u32::MAX {
out.insert(new_id);
}
}
}
}
if out.is_empty() {
continue;
}
out.optimize();
let offset = post_blob.len() as u64;
out.serialize_into(&mut post_blob)
.map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
fst_entries.push((tri, pack_entry(offset, out.len())?));
}
Ok((post_blob, fst_entries))
}
pub struct SegmentData {
fst: fst::Map<Mmap>,
post: Mmap,
post_len: usize,
pub docs: Vec<DocMeta>,
syms: crate::table::SymTable,
refs: Option<crate::table::RefTable>,
}
pub struct Segment {
pub id: u64,
data: Arc<SegmentData>,
live: RoaringBitmap,
}
impl Deref for Segment {
type Target = SegmentData;
fn deref(&self) -> &SegmentData {
&self.data
}
}
impl Segment {
pub fn open(paths: &Paths, seg_id: u64) -> Result<Segment> {
let (fst_and_post, tables) = rayon::join(
|| -> Result<(fst::Map<Mmap>, Mmap, usize)> {
let fst_path = paths.fst_file(seg_id);
let (fst, post_and_len) = rayon::join(
|| -> Result<fst::Map<Mmap>> {
let fst_file =
std::fs::File::open(&fst_path).map_err(|e| Error::io(&fst_path, e))?;
let fst_mmap =
unsafe { Mmap::map(&fst_file).map_err(|e| Error::io(&fst_path, e))? };
let fst = fst::Map::new(fst_mmap)?;
fst.as_fst()
.verify()
.map_err(|e| Error::Corrupt(format!("fst checksum: {e}")))?;
Ok(fst)
},
|| -> Result<(Mmap, usize)> {
let post_path = paths.post_file(seg_id);
let post_file =
std::fs::File::open(&post_path).map_err(|e| Error::io(&post_path, e))?;
let post = unsafe {
Mmap::map(&post_file).map_err(|e| Error::io(&post_path, e))?
};
let post_len = verify_checksum(&post, "postings blob")?.len();
Ok((post, post_len))
},
);
let (post, post_len) = post_and_len?;
Ok((fst?, post, post_len))
},
|| -> Result<(
Vec<DocMeta>,
crate::table::SymTable,
Option<crate::table::RefTable>,
)> {
let (docs_and_syms, refs) = rayon::join(
|| -> Result<(Vec<DocMeta>, crate::table::SymTable)> {
let docs: Vec<DocMeta> =
read_table(&paths.docs_file(seg_id), "docs table")?;
let syms = crate::table::SymTable::open(&paths.syms_file(seg_id))?;
Ok((docs, syms))
},
|| -> Result<Option<crate::table::RefTable>> {
let refs_path = paths.refs_file(seg_id);
match crate::table::RefTable::open(&refs_path) {
Ok(t) => Ok(Some(t)),
Err(Error::Io { source, .. })
if source.kind() == std::io::ErrorKind::NotFound =>
{
Ok(None)
}
Err(e) => Err(e),
}
},
);
let (docs, syms) = docs_and_syms?;
Ok((docs, syms, refs?))
},
);
let (fst, post, post_len) = fst_and_post?;
let (docs, syms, refs) = tables?;
if syms.doc_count() != docs.len()
|| refs.as_ref().is_some_and(|r| r.doc_count() != docs.len())
{
return Err(Error::Corrupt(format!(
"segment {seg_id}: side-table doc count does not match docs table"
)));
}
let live = read_bitmap(&paths.live_file(seg_id))?;
Ok(Segment {
id: seg_id,
data: Arc::new(SegmentData {
fst,
post,
post_len,
docs,
syms,
refs,
}),
live,
})
}
pub fn reopen(&self, paths: &Paths) -> Result<Segment> {
let live = read_bitmap(&paths.live_file(self.id))?;
Ok(Segment {
id: self.id,
data: self.data.clone(),
live,
})
}
pub fn is_live(&self, doc_id: u32) -> bool {
self.live.contains(doc_id)
}
pub fn subtract_live(&mut self, doc_ids: &[u32]) {
for &id in doc_ids {
self.live.remove(id);
}
}
pub fn live_count(&self) -> u64 {
self.live.len()
}
pub fn all_live(&self) -> RoaringBitmap {
self.live.clone()
}
pub fn candidates(&self, query: &TrigramQuery) -> Result<RoaringBitmap> {
let mut filtering = query
.dnfs
.iter()
.filter(|d| trigram::dnf_filters(d))
.peekable();
if filtering.peek().is_none() {
return Ok(self.all_live());
}
let mut result: Option<RoaringBitmap> = None;
for dnf in filtering {
let bm = self.data.dnf_bitmap(dnf)?;
result = Some(match result.take() {
None => bm,
Some(a) => a & bm,
});
if result.as_ref().is_some_and(|b| b.is_empty()) {
break;
}
}
let mut out = result.unwrap_or_default();
out &= &self.live;
Ok(out)
}
}
impl SegmentData {
pub fn doc(&self, doc_id: u32) -> Option<&DocMeta> {
self.docs.get(doc_id as usize)
}
pub fn doc_path_score(&self, doc_id: u32) -> f32 {
self.doc(doc_id)
.map(|d| crate::search::path_score(&d.path))
.unwrap_or(0.0)
}
pub fn sym_count(&self) -> usize {
self.syms.len()
}
pub fn sym(&self, i: u32) -> Option<SymbolEntry> {
self.syms.get(i)
}
pub(crate) fn sym_view(&self, i: u32) -> Option<crate::table::SymView<'_>> {
self.syms.view(i)
}
pub(crate) fn doc_sym_rows(&self, doc_id: u32) -> std::ops::Range<u32> {
self.syms.doc_range(doc_id)
}
pub(crate) fn doc_sym_views(
&self,
doc_id: u32,
) -> impl Iterator<Item = crate::table::SymView<'_>> {
self.syms
.doc_range(doc_id)
.filter_map(move |i| self.syms.view(i))
}
pub(crate) fn doc_ref_views(
&self,
doc_id: u32,
) -> impl Iterator<Item = crate::table::RefView<'_>> {
self.refs
.iter()
.flat_map(move |t| t.doc_range(doc_id).filter_map(move |i| t.view(i)))
}
pub(crate) fn ref_views_named<'s>(
&'s self,
name: &'s str,
) -> impl Iterator<Item = crate::table::RefView<'s>> + 's {
self.refs
.iter()
.flat_map(move |t| t.rows_named(name).filter_map(move |i| t.view(i)))
}
pub fn doc_syms(&self, doc_id: u32) -> impl Iterator<Item = SymbolEntry> + '_ {
self.syms
.doc_range(doc_id)
.filter_map(move |i| self.syms.get(i))
}
pub fn doc_sym_count(&self, doc_id: u32) -> u32 {
let r = self.syms.doc_range(doc_id);
r.end - r.start
}
pub fn doc_refs(&self, doc_id: u32) -> impl Iterator<Item = RefEntry> + '_ {
self.refs
.iter()
.flat_map(move |t| t.doc_range(doc_id).filter_map(move |i| t.get(i)))
}
pub fn refs_named<'s>(&'s self, name: &'s str) -> impl Iterator<Item = RefEntry> + 's {
self.refs
.iter()
.flat_map(move |t| t.rows_named(name).filter_map(move |i| t.get(i)))
}
pub fn calls_to<'s>(&'s self, name: &'s str) -> impl Iterator<Item = RefEntry> + 's {
self.refs_named(name).filter(|r| r.kind == RefKind::Call)
}
pub fn syms_by_lower<'s>(&'s self, lower: &'s str) -> impl Iterator<Item = u32> + 's {
self.syms.rows_named(lower)
}
pub fn sym_name(&self, i: u32) -> &str {
self.syms.name(i)
}
pub fn sym_name_lower(&self, i: u32) -> &str {
self.syms.name_lower(i)
}
fn posting_entry(&self, tri: Trigram) -> Option<u64> {
self.fst.get(tri)
}
fn posting_at(&self, offset: u64) -> Result<RoaringBitmap> {
let start = offset as usize;
let slice = self.post.get(start..self.post_len).ok_or_else(|| {
Error::Corrupt(format!(
"posting offset {start} out of range for postings blob of length {}",
self.post_len
))
})?;
RoaringBitmap::deserialize_from(slice)
.map_err(|e| Error::Corrupt(format!("posting list: {e}")))
}
fn dnf_bitmap(&self, dnf: &TrigramDnf) -> Result<RoaringBitmap> {
let mut acc = RoaringBitmap::new();
for group in dnf {
acc |= self.group_bitmap(group)?;
}
Ok(acc)
}
fn group_bitmap(&self, group: &[Trigram]) -> Result<RoaringBitmap> {
let mut entries: Vec<u64> = Vec::with_capacity(group.len());
for tri in group {
match self.posting_entry(*tri) {
Some(v) => entries.push(v),
None => return Ok(RoaringBitmap::new()),
}
}
entries.sort_unstable_by_key(|&v| unpack_card(v));
entries.truncate(MAX_GROUP_TRIGRAMS);
let mut acc: Option<RoaringBitmap> = None;
for v in entries {
let bm = self.posting_at(unpack_offset(v))?;
acc = Some(match acc.take() {
None => bm,
Some(a) => a & bm,
});
if acc.as_ref().is_some_and(|b| b.is_empty()) {
break;
}
}
Ok(acc.unwrap_or_default())
}
}
pub(crate) fn write_bitmap(path: &std::path::Path, bm: &RoaringBitmap) -> Result<()> {
let mut buf = Vec::with_capacity(bm.serialized_size() + 8);
bm.serialize_into(&mut buf)
.map_err(|e| Error::other(format!("roaring serialize: {e}")))?;
append_checksum(&mut buf);
write_atomic(path, &buf)
}
pub(crate) fn read_bitmap(path: &std::path::Path) -> Result<RoaringBitmap> {
let bytes = std::fs::read(path).map_err(|e| Error::io(path, e))?;
RoaringBitmap::deserialize_from(verify_checksum(&bytes, "live bitmap")?)
.map_err(|e| Error::Corrupt(format!("live bitmap: {e}")))
}