1use byteorder::{LittleEndian, WriteBytesExt};
14use std::io::{self, Read, Write};
15
16use super::opt_p4d::{find_optimal_bit_width, pack_with_exceptions, unpack_with_exceptions};
17use super::posting_common::{read_vint, write_vint};
18use crate::DocId;
19use crate::directories::OwnedBytes;
20use crate::structures::simd;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum PostingCodec {
26 #[default]
29 Rounded = 0,
30 Packed = 1,
33 Pfor = 2,
36}
37
38impl PostingCodec {
39 const HEADER_SHIFT: u32 = 6;
41 const WIDTH_MASK: u8 = 0x3F;
42
43 fn from_header_byte(doc_bits: u8) -> io::Result<(Self, u8)> {
44 let width = doc_bits & Self::WIDTH_MASK;
45 let codec = match doc_bits >> Self::HEADER_SHIFT {
46 0 => PostingCodec::Rounded,
47 1 => PostingCodec::Packed,
48 2 => PostingCodec::Pfor,
49 other => {
50 return Err(io::Error::new(
51 io::ErrorKind::InvalidData,
52 format!(
53 "posting block uses unknown codec id {other}; the index was written by a \
54 newer Hermes"
55 ),
56 ));
57 }
58 };
59 if width > 32 {
60 return Err(io::Error::new(
61 io::ErrorKind::InvalidData,
62 format!("posting block doc-id width {width} exceeds 32 bits"),
63 ));
64 }
65 Ok((codec, width))
66 }
67
68 fn header_byte(self, width: u8) -> u8 {
69 ((self as u8) << Self::HEADER_SHIFT) | width
70 }
71
72 pub fn parse(s: &str) -> Option<Self> {
73 match s.to_ascii_lowercase().as_str() {
74 "rounded" | "default" => Some(PostingCodec::Rounded),
75 "packed" | "bp128" | "exact" => Some(PostingCodec::Packed),
76 "pfor" | "optp4d" | "patched" => Some(PostingCodec::Pfor),
77 _ => None,
78 }
79 }
80}
81
82impl std::fmt::Display for PostingCodec {
83 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84 f.write_str(match self {
85 PostingCodec::Rounded => "rounded",
86 PostingCodec::Packed => "packed",
87 PostingCodec::Pfor => "pfor",
88 })
89 }
90}
91
92#[inline]
96fn packed_bytes(count: usize, width: u8) -> usize {
97 (count * width as usize).div_ceil(8)
98}
99
100fn pack_bits(values: &[u32], width: u8, out: &mut Vec<u8>) {
102 if width == 0 || values.is_empty() {
103 return;
104 }
105 if width == 32 {
106 for &v in values {
107 out.extend_from_slice(&v.to_le_bytes());
108 }
109 return;
110 }
111 let start = out.len();
112 out.resize(start + packed_bytes(values.len(), width), 0);
113 let dst = &mut out[start..];
114 let mut bit_pos = 0usize;
115 for &v in values {
116 let mut acc = (v as u64) << (bit_pos & 7);
117 let mut byte = bit_pos >> 3;
118 let mut remaining = (bit_pos & 7) + width as usize;
119 while remaining > 0 {
120 dst[byte] |= acc as u8;
121 acc >>= 8;
122 byte += 1;
123 remaining = remaining.saturating_sub(8);
124 }
125 bit_pos += width as usize;
126 }
127}
128
129fn unpack_bits(input: &[u8], width: u8, out: &mut [u32], count: usize) {
134 match width {
135 0 => out[..count].fill(0),
136 8 => simd::unpack_8bit(input, out, count),
137 16 => simd::unpack_16bit(input, out, count),
138 32 => simd::unpack_32bit(input, out, count),
139 _ => {
140 let mask = (1u64 << width) - 1;
141 let mut bit_pos = 0usize;
142 for slot in out[..count].iter_mut() {
143 let byte = bit_pos >> 3;
144 let word = if byte + 8 <= input.len() {
145 u64::from_le_bytes(input[byte..byte + 8].try_into().unwrap())
146 } else {
147 let mut word = 0u64;
148 for (i, &b) in input[byte..].iter().enumerate() {
149 word |= (b as u64) << (i * 8);
150 }
151 word
152 };
153 *slot = ((word >> (bit_pos & 7)) & mask) as u32;
154 bit_pos += width as usize;
155 }
156 }
157 }
158}
159
160fn pack_pfor(values: &[u32], out: &mut Vec<u8>) -> u8 {
164 let (width, _, _) = find_optimal_bit_width(values);
165 let (packed, exceptions) = pack_with_exceptions(values, width);
166 out.push(exceptions.len() as u8);
167 out.extend_from_slice(&packed);
168 for (pos, high) in exceptions {
169 out.push(pos);
170 out.extend_from_slice(&high.to_le_bytes());
171 }
172 width
173}
174
175fn pfor_payload_len(input: &[u8], count: usize, width: u8) -> io::Result<usize> {
177 let n_exceptions = *input
178 .first()
179 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "posting block truncated"))?
180 as usize;
181 Ok(1 + packed_bytes(count, width) + n_exceptions * 5)
182}
183
184fn unpack_pfor(input: &[u8], width: u8, out: &mut [u32], count: usize) -> io::Result<()> {
185 let n_exceptions = *input
186 .first()
187 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "posting block truncated"))?
188 as usize;
189 let packed_len = packed_bytes(count, width);
190 let table_at = 1 + packed_len;
191 if input.len() < table_at + n_exceptions * 5 {
192 return Err(io::Error::new(
193 io::ErrorKind::InvalidData,
194 "posting block exception table truncated",
195 ));
196 }
197 let packed = &input[1..table_at];
198 let mut exceptions: [(u8, u32); 128] = [(0, 0); 128];
199 for (i, entry) in input[table_at..table_at + n_exceptions * 5]
200 .chunks_exact(5)
201 .enumerate()
202 .take(128)
203 {
204 exceptions[i] = (
205 entry[0],
206 u32::from_le_bytes([entry[1], entry[2], entry[3], entry[4]]),
207 );
208 }
209 if width == 0 {
210 out[..count].fill(0);
212 for &(pos, value) in &exceptions[..n_exceptions.min(128)] {
213 if (pos as usize) < count {
214 out[pos as usize] = value;
215 }
216 }
217 return Ok(());
218 }
219 unpack_with_exceptions(
220 packed,
221 width,
222 &exceptions[..n_exceptions.min(128)],
223 count,
224 out,
225 );
226 Ok(())
227}
228
229#[derive(Debug, Clone, Copy, PartialEq, Eq)]
231pub struct Posting {
232 pub doc_id: DocId,
233 pub term_freq: u32,
234}
235
236#[derive(Debug, Clone, Default)]
238pub struct PostingList {
239 postings: Vec<Posting>,
240}
241
242impl PostingList {
243 pub fn new() -> Self {
244 Self::default()
245 }
246
247 pub fn with_capacity(capacity: usize) -> Self {
248 Self {
249 postings: Vec::with_capacity(capacity),
250 }
251 }
252
253 pub fn push(&mut self, doc_id: DocId, term_freq: u32) {
255 debug_assert!(
256 self.postings.is_empty() || self.postings.last().unwrap().doc_id < doc_id,
257 "Postings must be added in sorted order"
258 );
259 self.postings.push(Posting { doc_id, term_freq });
260 }
261
262 pub fn add(&mut self, doc_id: DocId, term_freq: u32) {
264 if let Some(last) = self.postings.last_mut()
265 && last.doc_id == doc_id
266 {
267 last.term_freq += term_freq;
268 return;
269 }
270 self.postings.push(Posting { doc_id, term_freq });
271 }
272
273 pub fn doc_count(&self) -> u32 {
275 self.postings.len() as u32
276 }
277
278 pub fn len(&self) -> usize {
279 self.postings.len()
280 }
281
282 pub fn is_empty(&self) -> bool {
283 self.postings.is_empty()
284 }
285
286 pub fn iter(&self) -> impl Iterator<Item = &Posting> {
287 self.postings.iter()
288 }
289
290 pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
292 write_vint(writer, self.postings.len() as u64)?;
294
295 let mut prev_doc_id = 0u32;
296 for posting in &self.postings {
297 let delta = posting.doc_id - prev_doc_id;
299 write_vint(writer, delta as u64)?;
300 write_vint(writer, posting.term_freq as u64)?;
301 prev_doc_id = posting.doc_id;
302 }
303
304 Ok(())
305 }
306
307 pub fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
309 let count = read_vint(reader)? as usize;
310 let mut postings = Vec::with_capacity(count);
311
312 let mut prev_doc_id = 0u32;
313 for _ in 0..count {
314 let delta = read_vint(reader)? as u32;
315 let term_freq = read_vint(reader)? as u32;
316 let doc_id = prev_doc_id + delta;
317 postings.push(Posting { doc_id, term_freq });
318 prev_doc_id = doc_id;
319 }
320
321 Ok(Self { postings })
322 }
323}
324
325pub struct PostingListIterator<'a> {
327 postings: &'a [Posting],
328 position: usize,
329}
330
331impl<'a> PostingListIterator<'a> {
332 pub fn new(posting_list: &'a PostingList) -> Self {
333 Self {
334 postings: &posting_list.postings,
335 position: 0,
336 }
337 }
338
339 pub fn doc(&self) -> DocId {
341 if self.position < self.postings.len() {
342 self.postings[self.position].doc_id
343 } else {
344 TERMINATED
345 }
346 }
347
348 pub fn term_freq(&self) -> u32 {
350 if self.position < self.postings.len() {
351 self.postings[self.position].term_freq
352 } else {
353 0
354 }
355 }
356
357 pub fn advance(&mut self) -> DocId {
359 self.position += 1;
360 self.doc()
361 }
362
363 pub fn seek(&mut self, target: DocId) -> DocId {
365 let remaining = &self.postings[self.position..];
366 let offset = remaining.partition_point(|p| p.doc_id < target);
367 self.position += offset;
368 self.doc()
369 }
370
371 pub fn size_hint(&self) -> usize {
373 self.postings.len().saturating_sub(self.position)
374 }
375}
376
377pub const TERMINATED: DocId = DocId::MAX;
379
380pub const BLOCK_SIZE: usize = 128;
389
390const L1_INTERVAL: usize = 8;
392
393const L0_SIZE: usize = 16;
396
397const L1_SIZE: usize = 4;
399
400const FOOTER_SIZE: usize = 24;
402
403const FOOTER_V2_SIZE: usize = FOOTER_SIZE + 20;
409
410const FOOTER_MAGIC: u32 = 0x324C_5042;
412
413const FLAG_POS_CURSORS: u32 = 1;
415
416const FLAG_LEN_BOUNDS: u32 = 2;
420
421const FLAG_L1_BOUNDS: u32 = 4;
425
426fn group_bounds_from_l0(l0: &[u8], l0_count: usize) -> Vec<u32> {
429 let mut groups = Vec::with_capacity(l0_count.div_ceil(L1_INTERVAL));
430 let mut idx = 0;
431 while idx < l0_count {
432 let end = (idx + L1_INTERVAL).min(l0_count);
433 let mut max_tf = 0u32;
434 let mut min_len = u32::MAX;
435 for block in idx..end {
436 let (_, _, _, word) = read_l0(l0, block);
437 let (tf, len) = unpack_bounds(word, true);
438 max_tf = max_tf.max(tf);
439 min_len = min_len.min(len.unwrap_or(1));
440 }
441 groups.push(pack_bounds(max_tf, min_len));
442 idx = end;
443 }
444 groups
445}
446
447#[inline]
449fn pack_bounds(max_tf: u32, min_len: u32) -> u32 {
450 max_tf.min(u16::MAX as u32) | (min_len.min(u16::MAX as u32) << 16)
451}
452
453#[inline]
456fn unpack_bounds(word: u32, packed: bool) -> (u32, Option<u32>) {
457 if packed {
458 (word & 0xFFFF, Some(word >> 16))
459 } else {
460 (f32::from_bits(word) as u32, None)
461 }
462}
463
464const CURSOR_SIZE: usize = 8;
467
468struct Footer {
470 stream_len: usize,
471 l0_count: usize,
472 l1_count: usize,
473 doc_count: u32,
474 max_tf: u32,
475 total_positions: u64,
476 has_cursors: bool,
477 len_bounds: bool,
478 l1_bounds: bool,
479 min_len: u32,
480}
481
482impl Footer {
483 fn parse(raw: &[u8]) -> io::Result<Self> {
484 if raw.len() < FOOTER_SIZE {
485 return Err(io::Error::new(
486 io::ErrorKind::InvalidData,
487 "posting data too short",
488 ));
489 }
490 let extended = raw.len() >= FOOTER_V2_SIZE
491 && u32::from_le_bytes(raw[raw.len() - 4..].try_into().unwrap()) == FOOTER_MAGIC;
492 let f = raw.len()
493 - if extended {
494 FOOTER_V2_SIZE
495 } else {
496 FOOTER_SIZE
497 };
498 let stream_len = u64::from_le_bytes(raw[f..f + 8].try_into().unwrap()) as usize;
499 let l0_count = u32::from_le_bytes(raw[f + 8..f + 12].try_into().unwrap()) as usize;
500 let l1_count = u32::from_le_bytes(raw[f + 12..f + 16].try_into().unwrap()) as usize;
501 let doc_count = u32::from_le_bytes(raw[f + 16..f + 20].try_into().unwrap());
502 let max_tf = u32::from_le_bytes(raw[f + 20..f + 24].try_into().unwrap());
503 let (total_positions, flags, min_len) = if extended {
504 let total = u64::from_le_bytes(raw[f + 24..f + 32].try_into().unwrap());
505 let flags = u32::from_le_bytes(raw[f + 32..f + 36].try_into().unwrap());
506 let min_len = u32::from_le_bytes(raw[f + 36..f + 40].try_into().unwrap());
507 (total, flags, min_len)
508 } else {
509 (0, 0, 0)
510 };
511 let footer = Self {
512 stream_len,
513 l0_count,
514 l1_count,
515 doc_count,
516 max_tf,
517 total_positions,
518 has_cursors: flags & FLAG_POS_CURSORS != 0,
519 len_bounds: flags & FLAG_LEN_BOUNDS != 0,
520 l1_bounds: flags & FLAG_L1_BOUNDS != 0,
521 min_len,
522 };
523 if footer.cursors_end() > f {
524 return Err(io::Error::new(
525 io::ErrorKind::InvalidData,
526 "posting list sections exceed the footer offset",
527 ));
528 }
529 Ok(footer)
530 }
531
532 fn l0_start(&self) -> usize {
533 self.stream_len
534 }
535 fn l0_end(&self) -> usize {
536 self.l0_start() + self.l0_count * L0_SIZE
537 }
538 fn l1_end(&self) -> usize {
539 self.l0_end() + self.l1_count * L1_SIZE
540 }
541 fn l1_bounds_end(&self) -> usize {
542 self.l1_end() + if self.l1_bounds { self.l1_count * 4 } else { 0 }
543 }
544 fn cursors_end(&self) -> usize {
545 self.l1_bounds_end()
546 + if self.has_cursors {
547 self.l0_count * CURSOR_SIZE
548 } else {
549 0
550 }
551 }
552}
553
554#[inline]
561fn read_l0(bytes: &[u8], idx: usize) -> (u32, u32, u32, u32) {
562 let b = &bytes[idx * L0_SIZE..][..L0_SIZE];
563 let first_doc = u32::from_le_bytes([b[0], b[1], b[2], b[3]]);
564 let last_doc = u32::from_le_bytes([b[4], b[5], b[6], b[7]]);
565 let offset = u32::from_le_bytes([b[8], b[9], b[10], b[11]]);
566 let bounds = u32::from_le_bytes([b[12], b[13], b[14], b[15]]);
567 (first_doc, last_doc, offset, bounds)
568}
569
570#[inline]
572fn write_l0(buf: &mut Vec<u8>, first_doc: u32, last_doc: u32, offset: u32, bounds: u32) {
573 buf.extend_from_slice(&first_doc.to_le_bytes());
574 buf.extend_from_slice(&last_doc.to_le_bytes());
575 buf.extend_from_slice(&offset.to_le_bytes());
576 buf.extend_from_slice(&bounds.to_le_bytes());
577}
578
579#[inline]
583fn block_len_from_l0(l0_bytes: &[u8], l0_count: usize, stream_len: usize, idx: usize) -> usize {
584 let (_, _, offset, _) = read_l0(l0_bytes, idx);
585 let end = if idx + 1 < l0_count {
586 read_l0(l0_bytes, idx + 1).2 as usize
587 } else {
588 stream_len
589 };
590 end.saturating_sub(offset as usize)
591}
592
593struct EncodedBlock {
596 doc_bits: u8,
597 tf_bits: u8,
598}
599
600fn encode_block_arrays(
602 codec: PostingCodec,
603 deltas: &[u32],
604 tfs: &[u32],
605 stream: &mut Vec<u8>,
606) -> EncodedBlock {
607 match codec {
608 PostingCodec::Rounded => {
609 let max_delta = deltas.iter().copied().max().unwrap_or(0);
610 let doc_bits = simd::round_bit_width(simd::bits_needed(max_delta));
611 let max_tf = tfs.iter().copied().max().unwrap_or(0);
612 let tf_bits = simd::round_bit_width(simd::bits_needed(max_tf));
613 if !deltas.is_empty() {
614 let rounded = simd::RoundedBitWidth::from_u8(doc_bits);
615 let start = stream.len();
616 stream.resize(start + deltas.len() * rounded.bytes_per_value(), 0);
617 simd::pack_rounded(deltas, rounded, &mut stream[start..]);
618 }
619 {
620 let rounded = simd::RoundedBitWidth::from_u8(tf_bits);
621 let start = stream.len();
622 stream.resize(start + tfs.len() * rounded.bytes_per_value(), 0);
623 simd::pack_rounded(tfs, rounded, &mut stream[start..]);
624 }
625 EncodedBlock {
626 doc_bits: codec.header_byte(doc_bits),
627 tf_bits,
628 }
629 }
630 PostingCodec::Packed => {
631 let max_delta = deltas.iter().copied().max().unwrap_or(0);
632 let doc_bits = simd::bits_needed(max_delta);
633 let max_tf = tfs.iter().copied().max().unwrap_or(0);
634 let tf_bits = simd::bits_needed(max_tf);
635 pack_bits(deltas, doc_bits, stream);
636 pack_bits(tfs, tf_bits, stream);
637 EncodedBlock {
638 doc_bits: codec.header_byte(doc_bits),
639 tf_bits,
640 }
641 }
642 PostingCodec::Pfor => {
643 let doc_bits = if deltas.is_empty() {
644 0
645 } else {
646 pack_pfor(deltas, stream)
647 };
648 let tf_bits = pack_pfor(tfs, stream);
649 EncodedBlock {
650 doc_bits: codec.header_byte(doc_bits),
651 tf_bits,
652 }
653 }
654 }
655}
656
657#[derive(Debug, Clone)]
658pub struct BlockPostingList {
659 stream: OwnedBytes,
661 l0_bytes: OwnedBytes,
664 l0_count: usize,
666 l1_docs: Vec<u32>,
669 l1_bounds: Vec<u32>,
672 doc_count: u32,
674 max_tf: u32,
676 pos_cursors: Option<OwnedBytes>,
680 total_positions: u64,
683 len_bounds: bool,
685 min_len: u32,
687}
688
689impl BlockPostingList {
690 #[inline]
692 fn read_l0_entry(&self, idx: usize) -> (u32, u32, u32, u32) {
693 read_l0(&self.l0_bytes, idx)
694 }
695
696 pub fn from_posting_list(list: &PostingList) -> io::Result<Self> {
705 Self::build(list, false, None, PostingCodec::Rounded)
706 }
707
708 pub fn from_posting_list_with_codec(
710 list: &PostingList,
711 codec: PostingCodec,
712 ) -> io::Result<Self> {
713 Self::build(list, false, None, codec)
714 }
715
716 pub fn from_posting_list_with_positions(list: &PostingList) -> io::Result<Self> {
721 Self::build(list, true, None, PostingCodec::Rounded)
722 }
723
724 pub fn from_posting_list_with(
729 list: &PostingList,
730 with_positions: bool,
731 length_of: Option<&dyn Fn(DocId) -> u32>,
732 ) -> io::Result<Self> {
733 Self::build(list, with_positions, length_of, PostingCodec::Rounded)
734 }
735
736 pub fn from_posting_list_with_options(
738 list: &PostingList,
739 with_positions: bool,
740 length_of: Option<&dyn Fn(DocId) -> u32>,
741 codec: PostingCodec,
742 ) -> io::Result<Self> {
743 Self::build(list, with_positions, length_of, codec)
744 }
745
746 fn build(
747 list: &PostingList,
748 with_positions: bool,
749 length_of: Option<&dyn Fn(DocId) -> u32>,
750 codec: PostingCodec,
751 ) -> io::Result<Self> {
752 let mut stream: Vec<u8> = Vec::new();
753 let mut l0_buf: Vec<u8> = Vec::new();
754 let mut l1_docs: Vec<u32> = Vec::new();
755 let mut cursors: Vec<u8> = Vec::new();
756 let mut positions_so_far = 0u64;
757 let mut l0_count = 0usize;
758 let mut max_tf = 0u32;
759 let mut list_min_len = u32::MAX;
760
761 let postings = &list.postings;
762 let mut i = 0;
763
764 let mut deltas = Vec::with_capacity(BLOCK_SIZE);
766 let mut tf_buf = Vec::with_capacity(BLOCK_SIZE);
767
768 while i < postings.len() {
769 if stream.len() > u32::MAX as usize {
770 return Err(io::Error::new(
771 io::ErrorKind::InvalidData,
772 "posting list stream exceeds u32::MAX bytes",
773 ));
774 }
775 let block_start = stream.len() as u32;
776 let block_end = (i + BLOCK_SIZE).min(postings.len());
777 let block = &postings[i..block_end];
778 let count = block.len();
779
780 let block_max_tf = block.iter().map(|p| p.term_freq).max().unwrap_or(0);
782 max_tf = max_tf.max(block_max_tf);
783
784 let base_doc_id = block.first().unwrap().doc_id;
785 let last_doc_id = block.last().unwrap().doc_id;
786
787 deltas.clear();
789 let mut prev = base_doc_id;
790 for posting in block.iter().skip(1) {
791 deltas.push(posting.doc_id - prev);
792 prev = posting.doc_id;
793 }
794
795 tf_buf.clear();
797 tf_buf.extend(block.iter().map(|p| p.term_freq));
798
799 stream.write_u16::<LittleEndian>(count as u16)?;
803 stream.write_u32::<LittleEndian>(base_doc_id)?;
804 let header_at = stream.len();
805 stream.push(0);
806 stream.push(0);
807 let encoded = encode_block_arrays(codec, &deltas, &tf_buf, &mut stream);
808 stream[header_at] = encoded.doc_bits;
809 stream[header_at + 1] = encoded.tf_bits;
810
811 let block_min_len = length_of.map_or(1, |length_of| {
813 block
814 .iter()
815 .map(|p| length_of(p.doc_id).max(1))
816 .min()
817 .unwrap_or(1)
818 });
819 list_min_len = list_min_len.min(block_min_len);
820 write_l0(
821 &mut l0_buf,
822 base_doc_id,
823 last_doc_id,
824 block_start,
825 pack_bounds(block_max_tf, block_min_len),
826 );
827 l0_count += 1;
828 if with_positions {
829 cursors.extend_from_slice(&positions_so_far.to_le_bytes());
830 positions_so_far += block.iter().map(|p| p.term_freq as u64).sum::<u64>();
831 }
832
833 if l0_count.is_multiple_of(L1_INTERVAL) {
835 l1_docs.push(last_doc_id);
836 }
837
838 i = block_end;
839 }
840
841 if !l0_count.is_multiple_of(L1_INTERVAL) && l0_count > 0 {
843 let (_, last_doc, _, _) = read_l0(&l0_buf, l0_count - 1);
844 l1_docs.push(last_doc);
845 }
846 let l1_bounds = group_bounds_from_l0(&l0_buf, l0_count);
847
848 Ok(Self {
849 stream: OwnedBytes::new(stream),
850 l0_bytes: OwnedBytes::new(l0_buf),
851 l0_count,
852 l1_docs,
853 l1_bounds,
854 doc_count: postings.len() as u32,
855 max_tf,
856 pos_cursors: with_positions.then(|| OwnedBytes::new(cursors)),
857 total_positions: positions_so_far,
858 len_bounds: true,
859 min_len: if list_min_len == u32::MAX {
860 1
861 } else {
862 list_min_len
863 },
864 })
865 }
866
867 pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
880 writer.write_all(&self.stream)?;
881 writer.write_all(&self.l0_bytes)?;
882 for &doc in &self.l1_docs {
883 writer.write_u32::<LittleEndian>(doc)?;
884 }
885 for &bounds in &self.l1_bounds {
886 writer.write_u32::<LittleEndian>(bounds)?;
887 }
888 if let Some(cursors) = &self.pos_cursors {
889 writer.write_all(cursors)?;
890 }
891 Self::write_footer(
892 writer,
893 self.stream.len() as u64,
894 self.l0_count,
895 self.l1_docs.len(),
896 self.doc_count,
897 self.max_tf,
898 self.total_positions,
899 self.pos_cursors.is_some(),
900 self.len_bounds.then_some(self.min_len),
901 !self.l1_bounds.is_empty(),
902 )
903 }
904
905 #[allow(clippy::too_many_arguments)]
906 fn write_footer<W: Write>(
907 writer: &mut W,
908 stream_len: u64,
909 l0_count: usize,
910 l1_count: usize,
911 doc_count: u32,
912 max_tf: u32,
913 total_positions: u64,
914 has_cursors: bool,
915 min_len: Option<u32>,
916 l1_bounds: bool,
917 ) -> io::Result<()> {
918 writer.write_u64::<LittleEndian>(stream_len)?;
919 writer.write_u32::<LittleEndian>(l0_count as u32)?;
920 writer.write_u32::<LittleEndian>(l1_count as u32)?;
921 writer.write_u32::<LittleEndian>(doc_count)?;
922 writer.write_u32::<LittleEndian>(max_tf)?;
923 writer.write_u64::<LittleEndian>(total_positions)?;
924 let mut flags = 0u32;
925 if has_cursors {
926 flags |= FLAG_POS_CURSORS;
927 }
928 if min_len.is_some() {
929 flags |= FLAG_LEN_BOUNDS;
930 }
931 if l1_bounds {
932 flags |= FLAG_L1_BOUNDS;
933 }
934 writer.write_u32::<LittleEndian>(flags)?;
935 writer.write_u32::<LittleEndian>(min_len.unwrap_or(0))?;
936 writer.write_u32::<LittleEndian>(FOOTER_MAGIC)?;
937 Ok(())
938 }
939
940 pub fn deserialize(raw: &[u8]) -> io::Result<Self> {
942 Self::deserialize_zero_copy(OwnedBytes::new(raw.to_vec()))
943 }
944
945 pub fn deserialize_zero_copy(raw: OwnedBytes) -> io::Result<Self> {
949 let footer = Footer::parse(raw.as_slice())?;
950 let l1_docs = Self::extract_l1_docs(&raw[footer.l0_end()..], footer.l1_count);
951 let l1_bounds = if footer.l1_bounds {
952 Self::extract_l1_docs(&raw[footer.l1_end()..], footer.l1_count)
953 } else {
954 Vec::new()
955 };
956 let pos_cursors = footer
957 .has_cursors
958 .then(|| raw.slice(footer.l1_bounds_end()..footer.cursors_end()));
959
960 Ok(Self {
961 stream: raw.slice(0..footer.stream_len),
962 l0_bytes: raw.slice(footer.l0_start()..footer.l0_end()),
963 l0_count: footer.l0_count,
964 l1_docs,
965 l1_bounds,
966 doc_count: footer.doc_count,
967 max_tf: footer.max_tf,
968 pos_cursors,
969 total_positions: footer.total_positions,
970 len_bounds: footer.len_bounds,
971 min_len: footer.min_len,
972 })
973 }
974
975 pub fn min_len(&self) -> Option<u32> {
978 self.len_bounds.then_some(self.min_len)
979 }
980
981 #[inline]
983 pub fn block_bounds(&self, block_idx: usize) -> Option<(u32, Option<u32>)> {
984 if block_idx >= self.l0_count {
985 return None;
986 }
987 let (_, _, _, word) = self.read_l0_entry(block_idx);
988 Some(unpack_bounds(word, self.len_bounds))
989 }
990
991 #[inline]
994 pub fn group_bounds(&self, block_idx: usize) -> Option<(u32, u32)> {
995 if block_idx >= self.l0_count {
996 return None;
997 }
998 let word = *self.l1_bounds.get(block_idx / L1_INTERVAL)?;
999 let (max_tf, min_len) = unpack_bounds(word, true);
1000 Some((max_tf, min_len.unwrap_or(1)))
1001 }
1002
1003 #[inline]
1005 pub fn group_last_doc(&self, block_idx: usize) -> Option<DocId> {
1006 self.l1_docs.get(block_idx / L1_INTERVAL).copied()
1007 }
1008
1009 #[inline]
1011 pub fn is_group_start(&self, block_idx: usize) -> bool {
1012 block_idx.is_multiple_of(L1_INTERVAL)
1013 }
1014
1015 #[inline]
1018 pub fn next_group_block(&self, block_idx: usize) -> usize {
1019 ((block_idx / L1_INTERVAL + 1) * L1_INTERVAL).min(self.l0_count)
1020 }
1021
1022 pub fn has_cursors_bytes(raw: &[u8]) -> bool {
1024 Footer::parse(raw).is_ok_and(|footer| footer.has_cursors)
1025 }
1026
1027 pub fn has_position_cursors(&self) -> bool {
1029 self.pos_cursors.is_some()
1030 }
1031
1032 pub fn total_positions(&self) -> u64 {
1034 self.total_positions
1035 }
1036
1037 #[inline]
1039 pub fn pos_cursor(&self, block_idx: usize) -> Option<u64> {
1040 let cursors = self.pos_cursors.as_ref()?;
1041 let p = block_idx * CURSOR_SIZE;
1042 cursors
1043 .get(p..p + CURSOR_SIZE)
1044 .map(|b| u64::from_le_bytes(b.try_into().unwrap()))
1045 }
1046
1047 fn extract_l1_docs(bytes: &[u8], count: usize) -> Vec<u32> {
1049 let mut docs = Vec::with_capacity(count);
1050 for i in 0..count {
1051 let p = i * L1_SIZE;
1052 docs.push(u32::from_le_bytes(bytes[p..p + 4].try_into().unwrap()));
1053 }
1054 docs
1055 }
1056
1057 pub fn doc_count(&self) -> u32 {
1058 self.doc_count
1059 }
1060
1061 pub fn max_tf(&self) -> u32 {
1063 self.max_tf
1064 }
1065
1066 pub fn num_blocks(&self) -> usize {
1068 self.l0_count
1069 }
1070
1071 pub fn block_max_tf(&self, block_idx: usize) -> Option<u32> {
1073 self.block_bounds(block_idx).map(|(max_tf, _)| max_tf)
1074 }
1075
1076 pub fn concatenate_blocks(sources: &[(BlockPostingList, u32)]) -> io::Result<Self> {
1079 let mut stream: Vec<u8> = Vec::new();
1080 let mut l0_buf: Vec<u8> = Vec::new();
1081 let mut l1_docs: Vec<u32> = Vec::new();
1082 let mut l0_count = 0usize;
1083 let mut total_docs = 0u32;
1084 let mut max_tf = 0u32;
1085 let all_cursors = sources.iter().all(|(s, _)| s.has_position_cursors());
1086 if !all_cursors && sources.iter().any(|(s, _)| s.has_position_cursors()) {
1087 return Err(io::Error::new(
1088 io::ErrorKind::InvalidData,
1089 "cannot concatenate posting lists with and without position cursors",
1090 ));
1091 }
1092 let mut cursors: Vec<u8> = Vec::new();
1093 let mut positions_before = 0u64;
1094 let mut min_len = u32::MAX;
1095
1096 for (source, doc_offset) in sources {
1097 max_tf = max_tf.max(source.max_tf);
1098 min_len = min_len.min(source.min_len().unwrap_or(1));
1099 for block_idx in 0..source.num_blocks() {
1100 if all_cursors {
1101 let cursor = source.pos_cursor(block_idx).unwrap_or(0) + positions_before;
1102 cursors.extend_from_slice(&cursor.to_le_bytes());
1103 }
1104 let (first_doc, last_doc, offset, word) = source.read_l0_entry(block_idx);
1105 let (block_max_tf, block_min_len) = unpack_bounds(word, source.len_bounds);
1106 let bounds = pack_bounds(block_max_tf, block_min_len.unwrap_or(1));
1107 let blk_size = source.block_len(block_idx);
1108 let block_bytes = &source.stream[offset as usize..offset as usize + blk_size];
1109
1110 let count = u16::from_le_bytes(block_bytes[0..2].try_into().unwrap());
1111 if stream.len() > u32::MAX as usize {
1112 return Err(io::Error::new(
1113 io::ErrorKind::InvalidData,
1114 "posting list stream exceeds u32::MAX bytes during concatenation",
1115 ));
1116 }
1117 let new_offset = stream.len() as u32;
1118
1119 stream.write_u16::<LittleEndian>(count)?;
1121 stream.write_u32::<LittleEndian>(first_doc + doc_offset)?;
1122 stream.extend_from_slice(&block_bytes[6..]);
1123
1124 let new_last = last_doc + doc_offset;
1125 write_l0(
1126 &mut l0_buf,
1127 first_doc + doc_offset,
1128 new_last,
1129 new_offset,
1130 bounds,
1131 );
1132 l0_count += 1;
1133 total_docs += count as u32;
1134
1135 if l0_count.is_multiple_of(L1_INTERVAL) {
1136 l1_docs.push(new_last);
1137 }
1138 }
1139 positions_before += source.total_positions;
1140 }
1141
1142 if !l0_count.is_multiple_of(L1_INTERVAL) && l0_count > 0 {
1144 let (_, last_doc, _, _) = read_l0(&l0_buf, l0_count - 1);
1145 l1_docs.push(last_doc);
1146 }
1147 let l1_bounds = group_bounds_from_l0(&l0_buf, l0_count);
1148
1149 Ok(Self {
1150 stream: OwnedBytes::new(stream),
1151 l0_bytes: OwnedBytes::new(l0_buf),
1152 l0_count,
1153 l1_docs,
1154 l1_bounds,
1155 doc_count: total_docs,
1156 max_tf,
1157 pos_cursors: all_cursors.then(|| OwnedBytes::new(cursors)),
1158 total_positions: if all_cursors { positions_before } else { 0 },
1159 len_bounds: true,
1160 min_len: if min_len == u32::MAX { 1 } else { min_len },
1161 })
1162 }
1163
1164 pub fn concatenate_streaming<W: Write>(
1179 sources: &[(&[u8], u32)], writer: &mut W,
1181 ) -> crate::Result<(u32, usize)> {
1182 let mut metas: Vec<Footer> = Vec::with_capacity(sources.len());
1183 let mut total_docs = 0u32;
1184 let mut merged_max_tf = 0u32;
1185 let mut merged_min_len = u32::MAX;
1186
1187 for (source_index, (raw, _)) in sources.iter().enumerate() {
1188 let footer = Footer::parse(raw).map_err(|e| {
1189 crate::Error::Corruption(format!(
1190 "posting list source {source_index} has an invalid footer: {e}"
1191 ))
1192 })?;
1193 total_docs += footer.doc_count;
1194 merged_max_tf = merged_max_tf.max(footer.max_tf);
1195 merged_min_len = merged_min_len.min(if footer.len_bounds { footer.min_len } else { 1 });
1196 metas.push(footer);
1197 }
1198
1199 if sources.len() == 1 && sources[0].1 == 0 {
1202 writer.write_all(sources[0].0)?;
1203 return Ok((metas[0].doc_count, sources[0].0.len()));
1204 }
1205
1206 let all_cursors = metas.iter().all(|m| m.has_cursors);
1207 if !all_cursors && metas.iter().any(|m| m.has_cursors) {
1208 return Err(crate::Error::Corruption(
1209 "cannot concatenate posting lists with and without position cursors".into(),
1210 ));
1211 }
1212
1213 let mut out_l0: Vec<u8> = Vec::new();
1216 let mut out_l1_docs: Vec<u32> = Vec::new();
1217 let mut out_cursors: Vec<u8> = Vec::new();
1218 let mut positions_before = 0u64;
1219 let mut out_l0_count = 0usize;
1220 let mut stream_written = 0u64;
1221 let mut patch_buf = [0u8; 8];
1222
1223 for (src_idx, meta) in metas.iter().enumerate() {
1224 let (raw, doc_offset) = &sources[src_idx];
1225 let l0_base = meta.l0_start(); let src_stream = &raw[..meta.stream_len];
1227 let cursors_base = meta.l1_bounds_end();
1228
1229 for i in 0..meta.l0_count {
1230 let (first_doc, last_doc, offset, word) = read_l0(&raw[l0_base..], i);
1232 let (block_max_tf, block_min_len) = unpack_bounds(word, meta.len_bounds);
1233 let bounds = pack_bounds(block_max_tf, block_min_len.unwrap_or(1));
1234 if all_cursors {
1235 let p = cursors_base + i * CURSOR_SIZE;
1236 let cursor = u64::from_le_bytes(raw[p..p + CURSOR_SIZE].try_into().unwrap());
1237 out_cursors.extend_from_slice(&(cursor + positions_before).to_le_bytes());
1238 }
1239
1240 let blk_size =
1242 block_len_from_l0(&raw[l0_base..], meta.l0_count, meta.stream_len, i);
1243 let block = &src_stream[offset as usize..offset as usize + blk_size];
1244
1245 let new_last = last_doc + doc_offset;
1247 if stream_written > u32::MAX as u64 {
1248 return Err(io::Error::new(
1249 io::ErrorKind::InvalidData,
1250 "posting list stream exceeds u32::MAX bytes during streaming merge",
1251 )
1252 .into());
1253 }
1254 write_l0(
1255 &mut out_l0,
1256 first_doc + doc_offset,
1257 new_last,
1258 stream_written as u32,
1259 bounds,
1260 );
1261 out_l0_count += 1;
1262
1263 if out_l0_count.is_multiple_of(L1_INTERVAL) {
1265 out_l1_docs.push(new_last);
1266 }
1267
1268 patch_buf.copy_from_slice(&block[0..8]);
1270 let blk_first = u32::from_le_bytes(patch_buf[2..6].try_into().unwrap());
1271 patch_buf[2..6].copy_from_slice(&(blk_first + doc_offset).to_le_bytes());
1272 writer.write_all(&patch_buf)?;
1273 writer.write_all(&block[8..])?;
1274
1275 stream_written += blk_size as u64;
1276 }
1277 positions_before += meta.total_positions;
1278 }
1279
1280 if !out_l0_count.is_multiple_of(L1_INTERVAL) && out_l0_count > 0 {
1282 let (_, last_doc, _, _) = read_l0(&out_l0, out_l0_count - 1);
1283 out_l1_docs.push(last_doc);
1284 }
1285
1286 let out_l1_bounds = group_bounds_from_l0(&out_l0, out_l0_count);
1288 writer.write_all(&out_l0)?;
1289 for &doc in &out_l1_docs {
1290 writer.write_u32::<LittleEndian>(doc)?;
1291 }
1292 for &bounds in &out_l1_bounds {
1293 writer.write_u32::<LittleEndian>(bounds)?;
1294 }
1295 writer.write_all(&out_cursors)?;
1296 Self::write_footer(
1297 writer,
1298 stream_written,
1299 out_l0_count,
1300 out_l1_docs.len(),
1301 total_docs,
1302 merged_max_tf,
1303 if all_cursors { positions_before } else { 0 },
1304 all_cursors,
1305 Some(if merged_min_len == u32::MAX {
1306 1
1307 } else {
1308 merged_min_len
1309 }),
1310 true,
1311 )?;
1312
1313 let l1_bytes_len = out_l1_docs.len() * L1_SIZE + out_l1_bounds.len() * 4;
1314 let total_bytes = stream_written as usize
1315 + out_l0.len()
1316 + l1_bytes_len
1317 + out_cursors.len()
1318 + FOOTER_V2_SIZE;
1319 Ok((total_docs, total_bytes))
1320 }
1321
1322 pub fn decode_block_into(
1329 &self,
1330 block_idx: usize,
1331 doc_ids: &mut Vec<u32>,
1332 tfs: &mut Vec<u32>,
1333 ) -> bool {
1334 if let Some((offset, tf_start, count)) = self.decode_block_doc_ids_only(block_idx, doc_ids)
1335 {
1336 self.decode_block_tfs_deferred(offset, tf_start, count, tfs);
1337 true
1338 } else {
1339 false
1340 }
1341 }
1342
1343 pub fn decode_block_doc_ids_only(
1348 &self,
1349 block_idx: usize,
1350 doc_ids: &mut Vec<u32>,
1351 ) -> Option<(usize, usize, usize)> {
1352 if block_idx >= self.l0_count {
1353 return None;
1354 }
1355
1356 let (_, _, offset, _) = self.read_l0_entry(block_idx);
1357 let pos = offset as usize;
1358 let blk_size = self.block_len(block_idx);
1359 let block_data = &self.stream[pos..pos + blk_size];
1360
1361 let count = u16::from_le_bytes(block_data[0..2].try_into().unwrap()) as usize;
1363 let first_doc = u32::from_le_bytes(block_data[2..6].try_into().unwrap());
1364 let (codec, doc_width) = PostingCodec::from_header_byte(block_data[6]).ok()?;
1365
1366 doc_ids.clear();
1367 doc_ids.resize(count, 0);
1368 doc_ids[0] = first_doc;
1369
1370 let payload = &block_data[8..];
1371 let deltas_bytes = if count > 1 {
1372 match codec {
1373 PostingCodec::Rounded => {
1374 let rounded = simd::RoundedBitWidth::from_u8(doc_width);
1375 let bytes = (count - 1) * rounded.bytes_per_value();
1376 simd::unpack_rounded(&payload[..bytes], rounded, &mut doc_ids[1..], count - 1);
1377 bytes
1378 }
1379 PostingCodec::Packed => {
1380 let bytes = packed_bytes(count - 1, doc_width);
1381 unpack_bits(&payload[..bytes], doc_width, &mut doc_ids[1..], count - 1);
1382 bytes
1383 }
1384 PostingCodec::Pfor => {
1385 let bytes = pfor_payload_len(payload, count - 1, doc_width).ok()?;
1386 unpack_pfor(&payload[..bytes], doc_width, &mut doc_ids[1..], count - 1).ok()?;
1387 bytes
1388 }
1389 }
1390 } else {
1391 0
1392 };
1393 for i in 1..count {
1394 doc_ids[i] = doc_ids[i].wrapping_add(doc_ids[i - 1]);
1395 }
1396
1397 let tfs_start = 8 + deltas_bytes;
1398 Some((pos, tfs_start, count))
1399 }
1400
1401 pub fn decode_block_tfs_deferred(
1405 &self,
1406 block_offset: usize,
1407 tf_start: usize,
1408 count: usize,
1409 tfs: &mut Vec<u32>,
1410 ) {
1411 let block_data = &self.stream[block_offset..];
1412 let codec = PostingCodec::from_header_byte(block_data[6])
1413 .map(|(codec, _)| codec)
1414 .unwrap_or_default();
1415 let tf_bits = block_data[7];
1416
1417 tfs.clear();
1418 tfs.resize(count, 0);
1419 let payload = &block_data[tf_start..];
1420 match codec {
1421 PostingCodec::Rounded => {
1422 let rounded = simd::RoundedBitWidth::from_u8(tf_bits);
1423 simd::unpack_rounded(
1424 &payload[..count * rounded.bytes_per_value()],
1425 rounded,
1426 tfs,
1427 count,
1428 );
1429 }
1430 PostingCodec::Packed => {
1431 unpack_bits(
1432 &payload[..packed_bytes(count, tf_bits)],
1433 tf_bits,
1434 tfs,
1435 count,
1436 );
1437 }
1438 PostingCodec::Pfor => {
1439 if let Ok(len) = pfor_payload_len(payload, count, tf_bits) {
1440 let _ = unpack_pfor(&payload[..len], tf_bits, tfs, count);
1441 }
1442 }
1443 }
1444 }
1445
1446 #[inline]
1448 fn block_len(&self, block_idx: usize) -> usize {
1449 block_len_from_l0(&self.l0_bytes, self.l0_count, self.stream.len(), block_idx)
1450 }
1451
1452 pub fn block_codec(&self, block_idx: usize) -> Option<PostingCodec> {
1454 if block_idx >= self.l0_count {
1455 return None;
1456 }
1457 let (_, _, offset, _) = self.read_l0_entry(block_idx);
1458 PostingCodec::from_header_byte(self.stream[offset as usize + 6])
1459 .ok()
1460 .map(|(codec, _)| codec)
1461 }
1462
1463 #[inline]
1465 pub fn block_first_doc(&self, block_idx: usize) -> Option<DocId> {
1466 if block_idx >= self.l0_count {
1467 return None;
1468 }
1469 let (first_doc, _, _, _) = self.read_l0_entry(block_idx);
1470 Some(first_doc)
1471 }
1472
1473 #[inline]
1475 pub fn block_last_doc(&self, block_idx: usize) -> Option<DocId> {
1476 if block_idx >= self.l0_count {
1477 return None;
1478 }
1479 let (_, last_doc, _, _) = self.read_l0_entry(block_idx);
1480 Some(last_doc)
1481 }
1482
1483 pub fn seek_block(&self, target: DocId, from_block: usize) -> Option<usize> {
1491 if from_block >= self.l0_count {
1492 return None;
1493 }
1494
1495 let from_l1 = from_block / L1_INTERVAL;
1496
1497 let l1_idx = if !self.l1_docs.is_empty() {
1499 let idx = from_l1 + simd::find_first_ge_u32(&self.l1_docs[from_l1..], target);
1500 if idx >= self.l1_docs.len() {
1501 return None;
1502 }
1503 idx
1504 } else {
1505 return None;
1506 };
1507
1508 let start = (l1_idx * L1_INTERVAL).max(from_block);
1510 let end = ((l1_idx + 1) * L1_INTERVAL).min(self.l0_count);
1511 let count = end - start;
1512
1513 let mut last_docs = [u32::MAX; L1_INTERVAL];
1514 for (j, idx) in (start..end).enumerate() {
1515 let (_, ld, _, _) = read_l0(&self.l0_bytes, idx);
1516 last_docs[j] = ld;
1517 }
1518 let within = simd::find_first_ge_u32(&last_docs[..count], target);
1519 let block_idx = start + within;
1520
1521 if block_idx < self.l0_count {
1522 Some(block_idx)
1523 } else {
1524 None
1525 }
1526 }
1527
1528 pub fn iterator(&self) -> BlockPostingIterator<'_> {
1530 BlockPostingIterator::new(self)
1531 }
1532
1533 pub fn into_iterator(self) -> BlockPostingIterator<'static> {
1535 BlockPostingIterator::owned(self)
1536 }
1537}
1538
1539pub struct BlockPostingIterator<'a> {
1546 block_list: std::borrow::Cow<'a, BlockPostingList>,
1547 current_block: usize,
1548 block_doc_ids: Vec<u32>,
1549 block_tfs: Vec<u32>,
1550 position_in_block: usize,
1551 tf_prefix: u64,
1555 exhausted: bool,
1556}
1557
1558impl<'a> BlockPostingIterator<'a> {
1559 fn new(block_list: &'a BlockPostingList) -> Self {
1560 let exhausted = block_list.l0_count == 0;
1561 let mut iter = Self {
1562 block_list: std::borrow::Cow::Borrowed(block_list),
1563 current_block: 0,
1564 block_doc_ids: Vec::with_capacity(BLOCK_SIZE),
1565 block_tfs: Vec::with_capacity(BLOCK_SIZE),
1566 position_in_block: 0,
1567 tf_prefix: 0,
1568 exhausted,
1569 };
1570 if !iter.exhausted {
1571 iter.load_block(0);
1572 }
1573 iter
1574 }
1575
1576 fn owned(block_list: BlockPostingList) -> BlockPostingIterator<'static> {
1577 let exhausted = block_list.l0_count == 0;
1578 let mut iter = BlockPostingIterator {
1579 block_list: std::borrow::Cow::Owned(block_list),
1580 current_block: 0,
1581 block_doc_ids: Vec::with_capacity(BLOCK_SIZE),
1582 block_tfs: Vec::with_capacity(BLOCK_SIZE),
1583 position_in_block: 0,
1584 tf_prefix: 0,
1585 exhausted,
1586 };
1587 if !iter.exhausted {
1588 iter.load_block(0);
1589 }
1590 iter
1591 }
1592
1593 fn load_block(&mut self, block_idx: usize) {
1594 if block_idx >= self.block_list.l0_count {
1595 self.exhausted = true;
1596 return;
1597 }
1598
1599 self.current_block = block_idx;
1600 self.position_in_block = 0;
1601 self.tf_prefix = 0;
1602
1603 self.block_list
1604 .decode_block_into(block_idx, &mut self.block_doc_ids, &mut self.block_tfs);
1605 }
1606
1607 #[inline]
1612 pub fn position_cursor(&self) -> u64 {
1613 self.block_list.pos_cursor(self.current_block).unwrap_or(0) + self.tf_prefix
1614 }
1615
1616 pub fn doc(&self) -> DocId {
1617 if self.exhausted {
1618 TERMINATED
1619 } else if self.position_in_block < self.block_doc_ids.len() {
1620 self.block_doc_ids[self.position_in_block]
1621 } else {
1622 TERMINATED
1623 }
1624 }
1625
1626 pub fn term_freq(&self) -> u32 {
1627 if self.exhausted || self.position_in_block >= self.block_tfs.len() {
1628 0
1629 } else {
1630 self.block_tfs[self.position_in_block]
1631 }
1632 }
1633
1634 pub fn advance(&mut self) -> DocId {
1635 if self.exhausted {
1636 return TERMINATED;
1637 }
1638
1639 if let Some(&tf) = self.block_tfs.get(self.position_in_block) {
1640 self.tf_prefix += tf as u64;
1641 }
1642 self.position_in_block += 1;
1643 if self.position_in_block >= self.block_doc_ids.len() {
1644 self.load_block(self.current_block + 1);
1645 }
1646 self.doc()
1647 }
1648
1649 pub fn seek(&mut self, target: DocId) -> DocId {
1650 if self.exhausted {
1651 return TERMINATED;
1652 }
1653
1654 let block_idx = match self.block_list.seek_block(target, self.current_block) {
1656 Some(idx) => idx,
1657 None => {
1658 self.exhausted = true;
1659 return TERMINATED;
1660 }
1661 };
1662
1663 if block_idx != self.current_block {
1664 self.load_block(block_idx);
1665 }
1666
1667 let remaining = &self.block_doc_ids[self.position_in_block..];
1669 let pos = crate::structures::simd::find_first_ge_u32(remaining, target);
1670 self.tf_prefix += self.block_tfs[self.position_in_block..self.position_in_block + pos]
1671 .iter()
1672 .map(|&tf| tf as u64)
1673 .sum::<u64>();
1674 self.position_in_block += pos;
1675
1676 if self.position_in_block >= self.block_doc_ids.len() {
1677 self.load_block(self.current_block + 1);
1678 }
1679 self.doc()
1680 }
1681
1682 pub fn skip_to_next_block(&mut self) -> DocId {
1686 if self.exhausted {
1687 return TERMINATED;
1688 }
1689 self.load_block(self.current_block + 1);
1690 self.doc()
1691 }
1692
1693 #[inline]
1695 pub fn current_block_idx(&self) -> usize {
1696 self.current_block
1697 }
1698
1699 #[inline]
1701 pub fn num_blocks(&self) -> usize {
1702 self.block_list.l0_count
1703 }
1704
1705 #[inline]
1707 pub fn current_block_max_tf(&self) -> u32 {
1708 if self.exhausted || self.current_block >= self.block_list.l0_count {
1709 0
1710 } else {
1711 self.block_list
1712 .block_max_tf(self.current_block)
1713 .unwrap_or(0)
1714 }
1715 }
1716}
1717
1718#[cfg(test)]
1719mod tests {
1720 use super::*;
1721
1722 #[test]
1723 fn test_posting_list_basic() {
1724 let mut list = PostingList::new();
1725 list.push(1, 2);
1726 list.push(5, 1);
1727 list.push(10, 3);
1728
1729 assert_eq!(list.len(), 3);
1730
1731 let mut iter = PostingListIterator::new(&list);
1732 assert_eq!(iter.doc(), 1);
1733 assert_eq!(iter.term_freq(), 2);
1734
1735 assert_eq!(iter.advance(), 5);
1736 assert_eq!(iter.term_freq(), 1);
1737
1738 assert_eq!(iter.advance(), 10);
1739 assert_eq!(iter.term_freq(), 3);
1740
1741 assert_eq!(iter.advance(), TERMINATED);
1742 }
1743
1744 #[test]
1745 fn test_posting_list_serialization() {
1746 let mut list = PostingList::new();
1747 for i in 0..100 {
1748 list.push(i * 3, (i % 5) + 1);
1749 }
1750
1751 let mut buffer = Vec::new();
1752 list.serialize(&mut buffer).unwrap();
1753
1754 let deserialized = PostingList::deserialize(&mut &buffer[..]).unwrap();
1755 assert_eq!(deserialized.len(), list.len());
1756
1757 for (a, b) in list.iter().zip(deserialized.iter()) {
1758 assert_eq!(a, b);
1759 }
1760 }
1761
1762 #[test]
1763 fn test_posting_list_seek() {
1764 let mut list = PostingList::new();
1765 for i in 0..100 {
1766 list.push(i * 2, 1);
1767 }
1768
1769 let mut iter = PostingListIterator::new(&list);
1770
1771 assert_eq!(iter.seek(50), 50);
1772 assert_eq!(iter.seek(51), 52);
1773 assert_eq!(iter.seek(200), TERMINATED);
1774 }
1775
1776 #[test]
1777 fn test_block_posting_list() {
1778 let mut list = PostingList::new();
1779 for i in 0..500 {
1780 list.push(i * 2, (i % 10) + 1);
1781 }
1782
1783 let block_list = BlockPostingList::from_posting_list(&list).unwrap();
1784 assert_eq!(block_list.doc_count(), 500);
1785
1786 let mut iter = block_list.iterator();
1787 assert_eq!(iter.doc(), 0);
1788 assert_eq!(iter.term_freq(), 1);
1789
1790 assert_eq!(iter.seek(500), 500);
1792 assert_eq!(iter.seek(998), 998);
1793 assert_eq!(iter.seek(1000), TERMINATED);
1794 }
1795
1796 #[test]
1797 fn test_block_posting_list_serialization() {
1798 let mut list = PostingList::new();
1799 for i in 0..300 {
1800 list.push(i * 3, i + 1);
1801 }
1802
1803 let block_list = BlockPostingList::from_posting_list(&list).unwrap();
1804
1805 let mut buffer = Vec::new();
1806 block_list.serialize(&mut buffer).unwrap();
1807
1808 let deserialized = BlockPostingList::deserialize(&buffer[..]).unwrap();
1809 assert_eq!(deserialized.doc_count(), block_list.doc_count());
1810
1811 let mut iter1 = block_list.iterator();
1813 let mut iter2 = deserialized.iterator();
1814
1815 while iter1.doc() != TERMINATED {
1816 assert_eq!(iter1.doc(), iter2.doc());
1817 assert_eq!(iter1.term_freq(), iter2.term_freq());
1818 iter1.advance();
1819 iter2.advance();
1820 }
1821 assert_eq!(iter2.doc(), TERMINATED);
1822 }
1823
1824 fn collect_postings(bpl: &BlockPostingList) -> Vec<(u32, u32)> {
1826 let mut result = Vec::new();
1827 let mut it = bpl.iterator();
1828 while it.doc() != TERMINATED {
1829 result.push((it.doc(), it.term_freq()));
1830 it.advance();
1831 }
1832 result
1833 }
1834
1835 fn build_bpl(postings: &[(u32, u32)]) -> BlockPostingList {
1837 let mut pl = PostingList::new();
1838 for &(doc_id, tf) in postings {
1839 pl.push(doc_id, tf);
1840 }
1841 BlockPostingList::from_posting_list(&pl).unwrap()
1842 }
1843
1844 fn serialize_bpl(bpl: &BlockPostingList) -> Vec<u8> {
1846 let mut buf = Vec::new();
1847 bpl.serialize(&mut buf).unwrap();
1848 buf
1849 }
1850
1851 #[test]
1852 fn test_concatenate_blocks_two_segments() {
1853 let a: Vec<(u32, u32)> = (0..100).map(|i| (i * 2, i + 1)).collect();
1855 let bpl_a = build_bpl(&a);
1856
1857 let b: Vec<(u32, u32)> = (0..100).map(|i| (i * 3, i + 2)).collect();
1859 let bpl_b = build_bpl(&b);
1860
1861 let merged =
1863 BlockPostingList::concatenate_blocks(&[(bpl_a.clone(), 0), (bpl_b.clone(), 200)])
1864 .unwrap();
1865
1866 assert_eq!(merged.doc_count(), 200);
1867
1868 let postings = collect_postings(&merged);
1869 assert_eq!(postings.len(), 200);
1870
1871 for (i, p) in postings.iter().enumerate().take(100) {
1873 assert_eq!(*p, (i as u32 * 2, i as u32 + 1));
1874 }
1875 for i in 0..100 {
1877 assert_eq!(postings[100 + i], (i as u32 * 3 + 200, i as u32 + 2));
1878 }
1879 }
1880
1881 #[test]
1882 fn test_concatenate_streaming_matches_blocks() {
1883 let seg_a: Vec<(u32, u32)> = (0..250).map(|i| (i * 2, (i % 7) + 1)).collect();
1885 let seg_b: Vec<(u32, u32)> = (0..180).map(|i| (i * 5, (i % 3) + 1)).collect();
1886 let seg_c: Vec<(u32, u32)> = (0..90).map(|i| (i * 10, (i % 11) + 1)).collect();
1887
1888 let bpl_a = build_bpl(&seg_a);
1889 let bpl_b = build_bpl(&seg_b);
1890 let bpl_c = build_bpl(&seg_c);
1891
1892 let offset_b = 1000u32;
1893 let offset_c = 2000u32;
1894
1895 let ref_merged = BlockPostingList::concatenate_blocks(&[
1897 (bpl_a.clone(), 0),
1898 (bpl_b.clone(), offset_b),
1899 (bpl_c.clone(), offset_c),
1900 ])
1901 .unwrap();
1902 let mut ref_buf = Vec::new();
1903 ref_merged.serialize(&mut ref_buf).unwrap();
1904
1905 let bytes_a = serialize_bpl(&bpl_a);
1907 let bytes_b = serialize_bpl(&bpl_b);
1908 let bytes_c = serialize_bpl(&bpl_c);
1909
1910 let sources: Vec<(&[u8], u32)> =
1911 vec![(&bytes_a, 0), (&bytes_b, offset_b), (&bytes_c, offset_c)];
1912 let mut stream_buf = Vec::new();
1913 let (doc_count, bytes_written) =
1914 BlockPostingList::concatenate_streaming(&sources, &mut stream_buf).unwrap();
1915
1916 assert_eq!(doc_count, 520); assert_eq!(bytes_written, stream_buf.len());
1918
1919 let ref_postings = collect_postings(&BlockPostingList::deserialize(&ref_buf).unwrap());
1921 let stream_postings =
1922 collect_postings(&BlockPostingList::deserialize(&stream_buf).unwrap());
1923
1924 assert_eq!(ref_postings.len(), stream_postings.len());
1925 for (i, (r, s)) in ref_postings.iter().zip(stream_postings.iter()).enumerate() {
1926 assert_eq!(r, s, "mismatch at posting {}", i);
1927 }
1928 }
1929
1930 #[test]
1931 fn test_concatenate_streaming_short_source_returns_corruption() {
1932 let seg_a: Vec<(u32, u32)> = (0..250).map(|i| (i * 2, (i % 7) + 1)).collect();
1937 let seg_c: Vec<(u32, u32)> = (0..90).map(|i| (i * 10, (i % 11) + 1)).collect();
1938 let bytes_a = serialize_bpl(&build_bpl(&seg_a));
1939 let bytes_c = serialize_bpl(&build_bpl(&seg_c));
1940 let short = vec![0u8; FOOTER_SIZE - 1]; let sources: Vec<(&[u8], u32)> = vec![(&bytes_a, 0), (&short, 1000), (&bytes_c, 2000)];
1943 let mut out = Vec::new();
1944 let result = BlockPostingList::concatenate_streaming(&sources, &mut out);
1945 assert!(
1946 matches!(result, Err(crate::Error::Corruption(_))),
1947 "short/corrupt source must be a Corruption error, not silently skipped: {:?}",
1948 result.map(|r| r.0)
1949 );
1950 }
1951
1952 #[test]
1953 fn test_multi_round_merge() {
1954 let segments: Vec<Vec<(u32, u32)>> = (0..4)
1961 .map(|seg| (0..200).map(|i| (i * 3, (i + seg * 7) % 10 + 1)).collect())
1962 .collect();
1963
1964 let bpls: Vec<BlockPostingList> = segments.iter().map(|s| build_bpl(s)).collect();
1965 let serialized: Vec<Vec<u8>> = bpls.iter().map(serialize_bpl).collect();
1966
1967 let mut merged_01 = Vec::new();
1969 let sources_01: Vec<(&[u8], u32)> = vec![(&serialized[0], 0), (&serialized[1], 600)];
1970 let (dc_01, _) =
1971 BlockPostingList::concatenate_streaming(&sources_01, &mut merged_01).unwrap();
1972 assert_eq!(dc_01, 400);
1973
1974 let mut merged_23 = Vec::new();
1975 let sources_23: Vec<(&[u8], u32)> = vec![(&serialized[2], 0), (&serialized[3], 600)];
1976 let (dc_23, _) =
1977 BlockPostingList::concatenate_streaming(&sources_23, &mut merged_23).unwrap();
1978 assert_eq!(dc_23, 400);
1979
1980 let mut final_merged = Vec::new();
1982 let sources_final: Vec<(&[u8], u32)> = vec![(&merged_01, 0), (&merged_23, 1200)];
1983 let (dc_final, _) =
1984 BlockPostingList::concatenate_streaming(&sources_final, &mut final_merged).unwrap();
1985 assert_eq!(dc_final, 800);
1986
1987 let final_bpl = BlockPostingList::deserialize(&final_merged).unwrap();
1989 let postings = collect_postings(&final_bpl);
1990 assert_eq!(postings.len(), 800);
1991
1992 assert_eq!(postings[0].0, 0); assert_eq!(postings[199].0, 597); assert_eq!(postings[200].0, 600); assert_eq!(postings[399].0, 1197); assert_eq!(postings[400].0, 1200); assert_eq!(postings[799].0, 2397); for seg in 0u32..4 {
2005 for i in 0u32..200 {
2006 let idx = (seg * 200 + i) as usize;
2007 assert_eq!(
2008 postings[idx].1,
2009 (i + seg * 7) % 10 + 1,
2010 "seg{} tf[{}]",
2011 seg,
2012 i
2013 );
2014 }
2015 }
2016
2017 let mut it = final_bpl.iterator();
2019 assert_eq!(it.seek(600), 600);
2020 assert_eq!(it.seek(1200), 1200);
2021 assert_eq!(it.seek(2397), 2397);
2022 assert_eq!(it.seek(2398), TERMINATED);
2023 }
2024
2025 #[test]
2026 fn test_large_scale_merge() {
2027 let num_segments = 5;
2030 let docs_per_segment = 2000;
2031 let docs_gap = 3; let segments: Vec<Vec<(u32, u32)>> = (0..num_segments)
2034 .map(|seg| {
2035 (0..docs_per_segment)
2036 .map(|i| (i as u32 * docs_gap, (i as u32 + seg as u32) % 20 + 1))
2037 .collect()
2038 })
2039 .collect();
2040
2041 let bpls: Vec<BlockPostingList> = segments.iter().map(|s| build_bpl(s)).collect();
2042
2043 for bpl in &bpls {
2045 assert!(
2046 bpl.num_blocks() >= 15,
2047 "expected >=15 blocks, got {}",
2048 bpl.num_blocks()
2049 );
2050 }
2051
2052 let serialized: Vec<Vec<u8>> = bpls.iter().map(serialize_bpl).collect();
2053
2054 let max_doc_per_seg = (docs_per_segment as u32 - 1) * docs_gap;
2056 let offsets: Vec<u32> = (0..num_segments)
2057 .map(|i| i as u32 * (max_doc_per_seg + 1))
2058 .collect();
2059
2060 let sources: Vec<(&[u8], u32)> = serialized
2061 .iter()
2062 .zip(offsets.iter())
2063 .map(|(b, o)| (b.as_slice(), *o))
2064 .collect();
2065
2066 let mut merged = Vec::new();
2067 let (doc_count, _) =
2068 BlockPostingList::concatenate_streaming(&sources, &mut merged).unwrap();
2069 assert_eq!(doc_count, (num_segments * docs_per_segment) as u32);
2070
2071 let merged_bpl = BlockPostingList::deserialize(&merged).unwrap();
2073 let postings = collect_postings(&merged_bpl);
2074 assert_eq!(postings.len(), num_segments * docs_per_segment);
2075
2076 for i in 1..postings.len() {
2078 assert!(
2079 postings[i].0 > postings[i - 1].0 || (i % docs_per_segment == 0), "doc_id not increasing at {}: {} vs {}",
2081 i,
2082 postings[i - 1].0,
2083 postings[i].0,
2084 );
2085 }
2086
2087 let mut it = merged_bpl.iterator();
2089 for (seg, &expected_first) in offsets.iter().enumerate() {
2090 assert_eq!(
2091 it.seek(expected_first),
2092 expected_first,
2093 "seek to segment {} start",
2094 seg
2095 );
2096 }
2097 }
2098
2099 #[test]
2100 fn test_merge_edge_cases() {
2101 let bpl_a = build_bpl(&[(0, 5)]);
2103 let bpl_b = build_bpl(&[(0, 3)]);
2104
2105 let merged =
2106 BlockPostingList::concatenate_blocks(&[(bpl_a.clone(), 0), (bpl_b.clone(), 1)])
2107 .unwrap();
2108 assert_eq!(merged.doc_count(), 2);
2109 let p = collect_postings(&merged);
2110 assert_eq!(p, vec![(0, 5), (1, 3)]);
2111
2112 let exact_block: Vec<(u32, u32)> = (0..BLOCK_SIZE as u32).map(|i| (i, i % 5 + 1)).collect();
2114 let bpl_exact = build_bpl(&exact_block);
2115 assert_eq!(bpl_exact.num_blocks(), 1);
2116
2117 let bytes = serialize_bpl(&bpl_exact);
2118 let mut out = Vec::new();
2119 let sources: Vec<(&[u8], u32)> = vec![(&bytes, 0), (&bytes, BLOCK_SIZE as u32)];
2120 let (dc, _) = BlockPostingList::concatenate_streaming(&sources, &mut out).unwrap();
2121 assert_eq!(dc, BLOCK_SIZE as u32 * 2);
2122
2123 let merged = BlockPostingList::deserialize(&out).unwrap();
2124 let postings = collect_postings(&merged);
2125 assert_eq!(postings.len(), BLOCK_SIZE * 2);
2126 assert_eq!(postings[BLOCK_SIZE].0, BLOCK_SIZE as u32);
2128
2129 let over_block: Vec<(u32, u32)> = (0..BLOCK_SIZE as u32 + 1).map(|i| (i * 2, 1)).collect();
2131 let bpl_over = build_bpl(&over_block);
2132 assert_eq!(bpl_over.num_blocks(), 2);
2133 }
2134
2135 #[test]
2136 fn test_streaming_roundtrip_single_source() {
2137 let docs: Vec<(u32, u32)> = (0..500).map(|i| (i * 7, i % 15 + 1)).collect();
2139 let bpl = build_bpl(&docs);
2140 let direct = serialize_bpl(&bpl);
2141
2142 let sources: Vec<(&[u8], u32)> = vec![(&direct, 0)];
2143 let mut streamed = Vec::new();
2144 BlockPostingList::concatenate_streaming(&sources, &mut streamed).unwrap();
2145
2146 let p1 = collect_postings(&BlockPostingList::deserialize(&direct).unwrap());
2148 let p2 = collect_postings(&BlockPostingList::deserialize(&streamed).unwrap());
2149 assert_eq!(p1, p2);
2150 }
2151
2152 #[test]
2153 fn test_max_tf_preserved_through_merge() {
2154 let mut a = Vec::new();
2156 for i in 0..200 {
2157 a.push((i * 2, if i == 100 { 50 } else { 1 }));
2158 }
2159 let bpl_a = build_bpl(&a);
2160 assert_eq!(bpl_a.max_tf(), 50);
2161
2162 let mut b = Vec::new();
2164 for i in 0..200 {
2165 b.push((i * 2, if i == 50 { 30 } else { 2 }));
2166 }
2167 let bpl_b = build_bpl(&b);
2168 assert_eq!(bpl_b.max_tf(), 30);
2169
2170 let bytes_a = serialize_bpl(&bpl_a);
2172 let bytes_b = serialize_bpl(&bpl_b);
2173 let sources: Vec<(&[u8], u32)> = vec![(&bytes_a, 0), (&bytes_b, 1000)];
2174 let mut out = Vec::new();
2175 BlockPostingList::concatenate_streaming(&sources, &mut out).unwrap();
2176
2177 let merged = BlockPostingList::deserialize(&out).unwrap();
2178 assert_eq!(merged.max_tf(), 50);
2179 assert_eq!(merged.doc_count(), 400);
2180 }
2181
2182 #[test]
2185 fn test_l0_l1_counts() {
2186 let bpl = build_bpl(&(0..50u32).map(|i| (i, 1)).collect::<Vec<_>>());
2188 assert_eq!(bpl.num_blocks(), 1);
2189 assert_eq!(bpl.l1_docs.len(), 1);
2190
2191 let n = BLOCK_SIZE * L1_INTERVAL;
2193 let bpl = build_bpl(&(0..n as u32).map(|i| (i * 2, 1)).collect::<Vec<_>>());
2194 assert_eq!(bpl.num_blocks(), L1_INTERVAL);
2195 assert_eq!(bpl.l1_docs.len(), 1);
2196
2197 let n = BLOCK_SIZE * L1_INTERVAL + 1;
2199 let bpl = build_bpl(&(0..n as u32).map(|i| (i * 2, 1)).collect::<Vec<_>>());
2200 assert_eq!(bpl.num_blocks(), L1_INTERVAL + 1);
2201 assert_eq!(bpl.l1_docs.len(), 2);
2202
2203 let n = BLOCK_SIZE * L1_INTERVAL * 3;
2205 let bpl = build_bpl(&(0..n as u32).map(|i| (i, 1)).collect::<Vec<_>>());
2206 assert_eq!(bpl.num_blocks(), L1_INTERVAL * 3);
2207 assert_eq!(bpl.l1_docs.len(), 3);
2208 }
2209
2210 #[test]
2211 fn test_l1_last_doc_values() {
2212 let n = BLOCK_SIZE * 20;
2214 let docs: Vec<(u32, u32)> = (0..n as u32).map(|i| (i * 3, 1)).collect();
2215 let bpl = build_bpl(&docs);
2216 assert_eq!(bpl.num_blocks(), 20);
2217 assert_eq!(bpl.l1_docs.len(), 3); let expected_l1_0 = bpl.block_last_doc(7).unwrap();
2221 assert_eq!(bpl.l1_docs[0], expected_l1_0);
2222
2223 let expected_l1_1 = bpl.block_last_doc(15).unwrap();
2225 assert_eq!(bpl.l1_docs[1], expected_l1_1);
2226
2227 let expected_l1_2 = bpl.block_last_doc(19).unwrap();
2229 assert_eq!(bpl.l1_docs[2], expected_l1_2);
2230 }
2231
2232 #[test]
2233 fn test_seek_block_basic() {
2234 let n = BLOCK_SIZE * 20;
2236 let docs: Vec<(u32, u32)> = (0..n as u32).map(|i| (i * 10, 1)).collect();
2237 let bpl = build_bpl(&docs);
2238
2239 assert_eq!(bpl.seek_block(0, 0), Some(0));
2241
2242 for blk in 0..20 {
2244 let first = bpl.block_first_doc(blk).unwrap();
2245 assert_eq!(
2246 bpl.seek_block(first, 0),
2247 Some(blk),
2248 "seek to block {} first_doc",
2249 blk
2250 );
2251 }
2252
2253 for blk in 0..20 {
2255 let last = bpl.block_last_doc(blk).unwrap();
2256 assert_eq!(
2257 bpl.seek_block(last, 0),
2258 Some(blk),
2259 "seek to block {} last_doc",
2260 blk
2261 );
2262 }
2263
2264 let max_doc = bpl.block_last_doc(19).unwrap();
2266 assert_eq!(bpl.seek_block(max_doc + 1, 0), None);
2267
2268 let mid_doc = bpl.block_first_doc(10).unwrap();
2270 assert_eq!(bpl.seek_block(mid_doc, 10), Some(10));
2271 assert_eq!(
2272 bpl.seek_block(mid_doc, 11),
2273 Some(11).or(bpl.seek_block(mid_doc, 11))
2274 );
2275 }
2276
2277 #[test]
2278 fn test_seek_block_across_l1_boundaries() {
2279 let n = BLOCK_SIZE * 24;
2281 let docs: Vec<(u32, u32)> = (0..n as u32).map(|i| (i * 5, 1)).collect();
2282 let bpl = build_bpl(&docs);
2283 assert_eq!(bpl.l1_docs.len(), 3);
2284
2285 for group in 0..3 {
2287 let blk = group * L1_INTERVAL;
2288 let target = bpl.block_first_doc(blk).unwrap();
2289 assert_eq!(
2290 bpl.seek_block(target, 0),
2291 Some(blk),
2292 "seek to group {} block {}",
2293 group,
2294 blk
2295 );
2296 }
2297
2298 let target = bpl.block_first_doc(20).unwrap() + 1;
2300 assert_eq!(bpl.seek_block(target, 0), Some(20));
2301 }
2302
2303 #[test]
2304 fn block_len_matches_l0_offsets() {
2305 let bpl = build_bpl(&(0..1000).map(|i| (i * 3, 1 + i % 4)).collect::<Vec<_>>());
2307 let mut total = 0usize;
2308 for b in 0..bpl.num_blocks() {
2309 let (_, _, offset, _) = bpl.read_l0_entry(b);
2310 assert_eq!(offset as usize, total, "block {b} offset");
2311 total += bpl.block_len(b);
2312 }
2313 assert_eq!(total, bpl.stream.len());
2314 }
2315
2316 #[test]
2320 fn every_codec_round_trips_and_seeks() {
2321 let mut postings: Vec<(u32, u32)> = Vec::new();
2322 let mut doc = 0u32;
2323 for i in 0..5000u32 {
2324 doc += if i % 97 == 0 { 100_000 } else { 1 + i % 7 };
2327 let tf = if i % 131 == 0 { 5000 } else { 1 + i % 3 };
2328 postings.push((doc, tf));
2329 }
2330 let mut list = PostingList::new();
2331 for &(d, tf) in &postings {
2332 list.push(d, tf);
2333 }
2334 let rounded = BlockPostingList::from_posting_list(&list).unwrap();
2335 let mut sizes = Vec::new();
2336 for codec in [
2337 PostingCodec::Rounded,
2338 PostingCodec::Packed,
2339 PostingCodec::Pfor,
2340 ] {
2341 let bpl = BlockPostingList::from_posting_list_with_codec(&list, codec).unwrap();
2342 assert_eq!(collect_postings(&bpl), postings, "{codec}");
2343 for b in 0..bpl.num_blocks() {
2344 assert_eq!(bpl.block_codec(b), Some(codec));
2345 assert_eq!(bpl.block_max_tf(b), rounded.block_max_tf(b));
2346 }
2347 let bytes = serialize_bpl(&bpl);
2349 let back = BlockPostingList::deserialize(&bytes).unwrap();
2350 assert_eq!(collect_postings(&back), postings, "{codec} deserialize");
2351 let back =
2352 BlockPostingList::deserialize_zero_copy(OwnedBytes::new(bytes.clone())).unwrap();
2353 assert_eq!(collect_postings(&back), postings, "{codec} zero-copy");
2354 let mut a = rounded.iterator();
2356 let mut b = back.iterator();
2357 for target in (0..postings.last().unwrap().0 + 10).step_by(2_003) {
2358 assert_eq!(a.seek(target), b.seek(target), "{codec} seek {target}");
2359 assert_eq!(a.term_freq(), b.term_freq());
2360 }
2361 sizes.push((codec, bytes.len()));
2362 }
2363 let rounded_bytes = serialize_bpl(&rounded);
2364 assert_eq!(
2365 serialize_bpl(
2366 &BlockPostingList::from_posting_list_with_codec(&list, PostingCodec::Rounded)
2367 .unwrap()
2368 ),
2369 rounded_bytes,
2370 "Rounded must stay byte-identical"
2371 );
2372 assert!(matches!(rounded.stream[6], 0 | 8 | 16 | 32));
2374 let size = |c: PostingCodec| sizes.iter().find(|(k, _)| *k == c).unwrap().1;
2375 assert!(size(PostingCodec::Packed) < size(PostingCodec::Rounded));
2376 assert!(size(PostingCodec::Pfor) < size(PostingCodec::Packed));
2377 }
2378
2379 #[test]
2381 fn mixed_codec_sources_concatenate() {
2382 let a: Vec<(u32, u32)> = (0..300u32).map(|i| (i * 5, 1 + i % 9)).collect();
2383 let b: Vec<(u32, u32)> = (0..300u32).map(|i| (i * 11 + 3, 2 + i % 5)).collect();
2384 let list_a = {
2385 let mut l = PostingList::new();
2386 a.iter().for_each(|&(d, t)| l.push(d, t));
2387 BlockPostingList::from_posting_list_with_codec(&l, PostingCodec::Pfor).unwrap()
2388 };
2389 let list_b = {
2390 let mut l = PostingList::new();
2391 b.iter().for_each(|&(d, t)| l.push(d, t));
2392 BlockPostingList::from_posting_list_with_codec(&l, PostingCodec::Packed).unwrap()
2393 };
2394 let offset_b = a.last().unwrap().0 + 1;
2395 let expected: Vec<(u32, u32)> = a
2396 .iter()
2397 .copied()
2398 .chain(b.iter().map(|&(d, t)| (d + offset_b, t)))
2399 .collect();
2400
2401 let merged = BlockPostingList::concatenate_blocks(&[
2402 (list_a.clone(), 0),
2403 (list_b.clone(), offset_b),
2404 ])
2405 .unwrap();
2406 assert_eq!(collect_postings(&merged), expected);
2407
2408 let bytes_a = serialize_bpl(&list_a);
2409 let bytes_b = serialize_bpl(&list_b);
2410 let mut out = Vec::new();
2411 let (docs, written) = BlockPostingList::concatenate_streaming(
2412 &[(bytes_a.as_slice(), 0), (bytes_b.as_slice(), offset_b)],
2413 &mut out,
2414 )
2415 .unwrap();
2416 assert_eq!(docs, 600);
2417 assert_eq!(written, out.len());
2418 let streamed = BlockPostingList::deserialize(&out).unwrap();
2419 assert_eq!(collect_postings(&streamed), expected);
2420 assert_eq!(streamed.block_codec(0), Some(PostingCodec::Pfor));
2421 assert_eq!(
2422 streamed.block_codec(streamed.num_blocks() - 1),
2423 Some(PostingCodec::Packed)
2424 );
2425 }
2426
2427 #[test]
2428 fn test_l0_entry_roundtrip() {
2429 let docs: Vec<(u32, u32)> = (0..1000u32).map(|i| (i * 3, (i % 10) + 1)).collect();
2431 let bpl = build_bpl(&docs);
2432
2433 let bytes = serialize_bpl(&bpl);
2434 let bpl2 = BlockPostingList::deserialize(&bytes).unwrap();
2435
2436 assert_eq!(bpl.num_blocks(), bpl2.num_blocks());
2437 for blk in 0..bpl.num_blocks() {
2438 assert_eq!(
2439 bpl.read_l0_entry(blk),
2440 bpl2.read_l0_entry(blk),
2441 "L0 entry mismatch at block {}",
2442 blk
2443 );
2444 }
2445
2446 assert_eq!(bpl.l1_docs, bpl2.l1_docs);
2448 }
2449
2450 #[test]
2451 fn test_zero_copy_deserialize_matches() {
2452 let docs: Vec<(u32, u32)> = (0..2000u32).map(|i| (i * 2, (i % 5) + 1)).collect();
2453 let bpl = build_bpl(&docs);
2454 let bytes = serialize_bpl(&bpl);
2455
2456 let copied = BlockPostingList::deserialize(&bytes).unwrap();
2457 let zero_copy =
2458 BlockPostingList::deserialize_zero_copy(OwnedBytes::new(bytes.clone())).unwrap();
2459
2460 assert_eq!(copied.l0_count, zero_copy.l0_count);
2462 assert_eq!(copied.l1_docs, zero_copy.l1_docs);
2463 assert_eq!(copied.doc_count, zero_copy.doc_count);
2464 assert_eq!(copied.max_tf, zero_copy.max_tf);
2465
2466 let p1 = collect_postings(&copied);
2468 let p2 = collect_postings(&zero_copy);
2469 assert_eq!(p1, p2);
2470 }
2471
2472 #[test]
2473 fn test_l1_preserved_through_streaming_merge() {
2474 let seg_a = build_bpl(&(0..1000u32).map(|i| (i * 2, 1)).collect::<Vec<_>>());
2476 let seg_b = build_bpl(&(0..800u32).map(|i| (i * 3, 2)).collect::<Vec<_>>());
2477 let seg_c = build_bpl(&(0..500u32).map(|i| (i * 5, 3)).collect::<Vec<_>>());
2478
2479 let bytes_a = serialize_bpl(&seg_a);
2480 let bytes_b = serialize_bpl(&seg_b);
2481 let bytes_c = serialize_bpl(&seg_c);
2482
2483 let sources: Vec<(&[u8], u32)> = vec![(&bytes_a, 0), (&bytes_b, 10000), (&bytes_c, 20000)];
2484 let mut out = Vec::new();
2485 BlockPostingList::concatenate_streaming(&sources, &mut out).unwrap();
2486
2487 let merged = BlockPostingList::deserialize(&out).unwrap();
2488 let expected_l1_count = merged.num_blocks().div_ceil(L1_INTERVAL);
2489 assert_eq!(merged.l1_docs.len(), expected_l1_count);
2490
2491 for (i, &l1_doc) in merged.l1_docs.iter().enumerate() {
2493 let last_block_in_group = ((i + 1) * L1_INTERVAL - 1).min(merged.num_blocks() - 1);
2494 let expected = merged.block_last_doc(last_block_in_group).unwrap();
2495 assert_eq!(l1_doc, expected, "L1[{}] mismatch", i);
2496 }
2497
2498 for blk in 0..merged.num_blocks() {
2500 let first = merged.block_first_doc(blk).unwrap();
2501 assert_eq!(merged.seek_block(first, 0), Some(blk));
2502 }
2503 }
2504
2505 #[test]
2506 fn test_seek_block_single_block() {
2507 let bpl = build_bpl(&[(0, 1), (10, 2), (20, 3)]);
2509 assert_eq!(bpl.num_blocks(), 1);
2510 assert_eq!(bpl.l1_docs.len(), 1);
2511
2512 assert_eq!(bpl.seek_block(0, 0), Some(0));
2513 assert_eq!(bpl.seek_block(10, 0), Some(0));
2514 assert_eq!(bpl.seek_block(20, 0), Some(0));
2515 assert_eq!(bpl.seek_block(21, 0), None);
2516 }
2517
2518 #[test]
2519 fn test_footer_size() {
2520 let docs: Vec<(u32, u32)> = (0..500u32).map(|i| (i * 2, 1)).collect();
2522 let bpl = build_bpl(&docs);
2523 let bytes = serialize_bpl(&bpl);
2524
2525 let expected = bpl.stream.len()
2526 + bpl.l0_count * L0_SIZE
2527 + bpl.l1_docs.len() * (L1_SIZE + 4)
2528 + FOOTER_V2_SIZE;
2529 assert_eq!(bytes.len(), expected);
2530 }
2531
2532 fn build_bpl_with_positions(postings: &[(u32, u32)]) -> BlockPostingList {
2533 let mut list = PostingList::new();
2534 for &(doc, tf) in postings {
2535 list.push(doc, tf);
2536 }
2537 BlockPostingList::from_posting_list_with_positions(&list).unwrap()
2538 }
2539
2540 fn expected_cursors(postings: &[(u32, u32)]) -> Vec<u64> {
2542 let mut acc = 0u64;
2543 postings
2544 .iter()
2545 .map(|&(_, tf)| {
2546 let c = acc;
2547 acc += tf as u64;
2548 c
2549 })
2550 .collect()
2551 }
2552
2553 fn iterator_cursors(bpl: &BlockPostingList) -> Vec<u64> {
2554 let mut it = bpl.iterator();
2555 let mut out = Vec::new();
2556 while it.doc() != TERMINATED {
2557 out.push(it.position_cursor());
2558 it.advance();
2559 }
2560 out
2561 }
2562
2563 #[test]
2564 fn position_cursors_survive_serialization_and_seeks() {
2565 let docs: Vec<(u32, u32)> = (0..700u32).map(|i| (i * 3, i % 5 + 1)).collect();
2566 let bpl = build_bpl_with_positions(&docs);
2567 assert!(bpl.has_position_cursors());
2568 assert_eq!(
2569 bpl.total_positions(),
2570 docs.iter().map(|&(_, tf)| tf as u64).sum::<u64>()
2571 );
2572 assert_eq!(bpl.pos_cursor(0), Some(0));
2573 assert_eq!(
2574 bpl.pos_cursor(1),
2575 Some(docs[..128].iter().map(|&(_, tf)| tf as u64).sum::<u64>())
2576 );
2577 assert_eq!(iterator_cursors(&bpl), expected_cursors(&docs));
2578
2579 let bytes = serialize_bpl(&bpl);
2580 assert_eq!(
2581 bytes.len(),
2582 bpl.stream.len()
2583 + bpl.l0_count * (L0_SIZE + CURSOR_SIZE)
2584 + bpl.l1_docs.len() * (L1_SIZE + 4)
2585 + FOOTER_V2_SIZE
2586 );
2587 assert!(BlockPostingList::has_cursors_bytes(&bytes));
2588 let decoded =
2589 BlockPostingList::deserialize_zero_copy(OwnedBytes::new(bytes.clone())).unwrap();
2590 assert_eq!(iterator_cursors(&decoded), expected_cursors(&docs));
2591 assert_eq!(decoded.total_positions(), bpl.total_positions());
2592
2593 let mut it = decoded.iterator();
2595 let expected = expected_cursors(&docs);
2596 for (i, &(doc, _)) in docs.iter().enumerate().step_by(37) {
2597 assert_eq!(it.seek(doc), doc);
2598 assert_eq!(it.position_cursor(), expected[i], "cursor at doc {doc}");
2599 }
2600 let mut it = decoded.iterator();
2601 assert_eq!(it.seek(docs[600].0 + 1), docs[601].0);
2602 assert_eq!(it.position_cursor(), expected[601]);
2603
2604 let plain = build_bpl(&docs);
2607 assert!(!plain.has_position_cursors());
2608 assert_eq!(plain.pos_cursor(0), None);
2609 assert_eq!(plain.total_positions(), 0);
2610 }
2611
2612 #[test]
2613 fn length_bounds_are_packed_per_block_and_survive_merges() {
2614 let docs: Vec<(u32, u32)> = (0..300u32).map(|i| (i, i % 3 + 1)).collect();
2615 let length_of = |doc: u32| 10 + (doc % 50) * 7;
2616 let mut list = PostingList::new();
2617 for &(doc, tf) in &docs {
2618 list.push(doc, tf);
2619 }
2620 let bpl = BlockPostingList::from_posting_list_with(&list, true, Some(&length_of)).unwrap();
2621 assert_eq!(bpl.min_len(), Some(10));
2622 assert_eq!(bpl.block_bounds(0), Some((3, Some(10))));
2623 assert_eq!(bpl.block_bounds(2), Some((3, Some(52))));
2625 assert_eq!(bpl.block_max_tf(2), Some(3));
2626
2627 let bytes = serialize_bpl(&bpl);
2628 let decoded = BlockPostingList::deserialize(&bytes).unwrap();
2629 assert_eq!(decoded.min_len(), Some(10));
2630 assert_eq!(decoded.block_bounds(2), Some((3, Some(52))));
2631 assert_eq!(decoded.group_bounds(0), Some((3, 10)));
2634 assert_eq!(decoded.group_bounds(2), Some((3, 10)));
2635 assert_eq!(decoded.group_bounds(3), None);
2636 assert_eq!(decoded.group_last_doc(1), Some(299));
2637 assert_eq!(decoded.next_group_block(1), 3);
2638
2639 let plain = build_bpl(&docs);
2641 assert_eq!(plain.min_len(), Some(1));
2642 assert_eq!(plain.block_bounds(0), Some((3, Some(1))));
2643
2644 let mut out = Vec::new();
2646 BlockPostingList::concatenate_streaming(&[(&bytes, 0), (&bytes, 1000)], &mut out).unwrap();
2647 let merged = BlockPostingList::deserialize(&out).unwrap();
2648 assert_eq!(merged.min_len(), Some(10));
2649 assert_eq!(merged.block_bounds(2), Some((3, Some(52))));
2650 assert_eq!(merged.block_bounds(3), Some((3, Some(10))));
2651 assert_eq!(merged.block_max_tf(5), Some(3));
2652 assert_eq!(merged.group_bounds(5), Some((3, 10)));
2655 assert_eq!(merged.group_last_doc(5), Some(1299));
2656 assert_eq!(merged.next_group_block(5), 6);
2657 }
2658
2659 #[test]
2660 fn legacy_footer_without_magic_still_deserializes() {
2661 let docs: Vec<(u32, u32)> = (0..300u32).map(|i| (i * 2, 1 + i % 3)).collect();
2662 let bpl = build_bpl(&docs);
2663 let bytes = serialize_bpl(&bpl);
2664 let legacy = bytes[..bytes.len() - (FOOTER_V2_SIZE - FOOTER_SIZE)].to_vec();
2666 assert!(!BlockPostingList::has_cursors_bytes(&legacy));
2667 let decoded = BlockPostingList::deserialize(&legacy).unwrap();
2669 assert_eq!(collect_postings(&decoded), docs);
2670 assert_eq!(decoded.max_tf(), 3);
2671 assert!(!decoded.has_position_cursors());
2672 assert_eq!(decoded.min_len(), None);
2673 assert_eq!(decoded.group_bounds(0), None);
2674 let mut out = Vec::new();
2676 let (count, written) =
2677 BlockPostingList::concatenate_streaming(&[(&legacy, 0), (&legacy, 1000)], &mut out)
2678 .unwrap();
2679 assert_eq!(count, 600);
2680 assert_eq!(written, out.len());
2681 let merged = BlockPostingList::deserialize(&out).unwrap();
2682 assert_eq!(merged.doc_count(), 600);
2683 assert!(!merged.has_position_cursors());
2684 }
2685
2686 #[test]
2687 fn streaming_merge_rebases_position_cursors() {
2688 let a: Vec<(u32, u32)> = (0..200u32).map(|i| (i, i % 4 + 1)).collect();
2689 let b: Vec<(u32, u32)> = (0..150u32).map(|i| (i * 2, 2)).collect();
2690 let bytes_a = serialize_bpl(&build_bpl_with_positions(&a));
2691 let bytes_b = serialize_bpl(&build_bpl_with_positions(&b));
2692 let mut out = Vec::new();
2693 let (count, written) =
2694 BlockPostingList::concatenate_streaming(&[(&bytes_a, 0), (&bytes_b, 1000)], &mut out)
2695 .unwrap();
2696 assert_eq!(count, 350);
2697 assert_eq!(written, out.len());
2698 let merged = BlockPostingList::deserialize(&out).unwrap();
2699 assert!(merged.has_position_cursors());
2700 let all: Vec<(u32, u32)> = a
2701 .iter()
2702 .copied()
2703 .chain(b.iter().map(|&(d, tf)| (d + 1000, tf)))
2704 .collect();
2705 assert_eq!(collect_postings(&merged), all);
2706 assert_eq!(iterator_cursors(&merged), expected_cursors(&all));
2707 assert_eq!(
2708 merged.total_positions(),
2709 all.iter().map(|&(_, tf)| tf as u64).sum::<u64>()
2710 );
2711 let reference = BlockPostingList::concatenate_blocks(&[
2713 (build_bpl_with_positions(&a), 0),
2714 (build_bpl_with_positions(&b), 1000),
2715 ])
2716 .unwrap();
2717 assert_eq!(iterator_cursors(&reference), expected_cursors(&all));
2718 let plain = serialize_bpl(&build_bpl(&b));
2720 assert!(
2721 BlockPostingList::concatenate_streaming(
2722 &[(&bytes_a, 0), (&plain, 1000)],
2723 &mut Vec::new()
2724 )
2725 .is_err()
2726 );
2727 }
2728
2729 #[test]
2730 fn test_seek_block_from_block_skips_earlier() {
2731 let n = BLOCK_SIZE * 16;
2733 let docs: Vec<(u32, u32)> = (0..n as u32).map(|i| (i * 3, 1)).collect();
2734 let bpl = build_bpl(&docs);
2735
2736 let target_in_5 = bpl.block_first_doc(5).unwrap() + 1;
2738 let result = bpl.seek_block(target_in_5, 8);
2741 assert!(result.is_some());
2742 assert!(result.unwrap() >= 8);
2743 }
2744}