1#[cfg(feature = "native")]
43mod compact;
44mod directory;
45#[cfg(feature = "native")]
46pub(crate) use compact::PositionRangeSource;
47
48use std::io::{self, Write};
49
50use byteorder::{LittleEndian, WriteBytesExt};
51
52use super::{PostingCodec, bitpacking4x};
53use crate::directories::OwnedBytes;
54use crate::structures::simd;
55
56pub const POSITION_STREAM_BLOCK: usize = 128;
58
59const BLOCK_HEADER: usize = 4;
60const INDEX_ENTRY: usize = 12;
61const FOOTER: usize = 16;
62const MAGIC: u32 = 0x3553_4F50;
64const UNIQUE_POSITIONS: u32 = 1 << 31;
66
67fn has_unique_positions(raw: &[u8]) -> bool {
68 let at = raw.len() - FOOTER;
69 u32::from_le_bytes(raw[at..at + 4].try_into().unwrap()) & UNIQUE_POSITIONS != 0
70}
71
72fn block_count_and_flags(count: usize, unique: bool) -> io::Result<u32> {
73 let count = u32::try_from(count)
74 .ok()
75 .filter(|&count| count < UNIQUE_POSITIONS)
76 .ok_or_else(|| io::Error::other("too many position blocks"))?;
77 Ok(count | if unique { UNIQUE_POSITIONS } else { 0 })
78}
79
80fn footer_magic(compact: bool) -> u32 {
81 if compact {
82 directory::COMPACT_MAGIC
83 } else {
84 MAGIC
85 }
86}
87
88pub struct PositionStreamEncoder<W: Write> {
90 writer: W,
91 pending: Vec<u32>,
92 index: Vec<(u32, u64)>,
93 index_limit: Option<usize>,
94 written: u64,
95 total: u64,
96 scratch: Vec<u8>,
97 codec: PostingCodec,
98 compact: bool,
99 descriptors: Vec<u16>,
100 unique_positions: bool,
101}
102
103impl<W: Write> PositionStreamEncoder<W> {
104 pub fn new(writer: W) -> Self {
105 Self::with_posting_codec(writer, PostingCodec::Rounded)
106 }
107
108 pub fn with_posting_codec(writer: W, codec: PostingCodec) -> Self {
111 Self {
112 writer,
113 pending: Vec::with_capacity(POSITION_STREAM_BLOCK),
114 index: Vec::new(),
115 index_limit: None,
116 written: 0,
117 total: 0,
118 scratch: Vec::with_capacity(BLOCK_HEADER + POSITION_STREAM_BLOCK * 4),
119 codec,
120 compact: false,
121 descriptors: Vec::new(),
122 unique_positions: true,
123 }
124 }
125
126 pub fn with_compact_directory(mut self) -> Self {
128 assert!(
129 self.index.is_empty() && self.pending.is_empty(),
130 "select the position format before appending values"
131 );
132 self.compact = true;
133 self
134 }
135
136 pub fn push_doc(&mut self, positions: &mut [u32]) -> io::Result<()> {
140 positions.sort_unstable();
141 let mut prev = 0u32;
142 for (index, &position) in positions.iter().enumerate() {
143 self.unique_positions &= index == 0 || position != prev;
144 self.push_value(position - prev)?;
145 prev = position;
146 }
147 Ok(())
148 }
149
150 pub fn push_values(&mut self, values: &[u32]) -> io::Result<()> {
152 self.unique_positions = false;
153 for &value in values {
154 self.push_value(value)?;
155 }
156 Ok(())
157 }
158
159 #[inline]
160 fn push_value(&mut self, value: u32) -> io::Result<()> {
161 self.pending.push(value);
162 self.total += 1;
163 if self.pending.len() == POSITION_STREAM_BLOCK {
164 self.flush_block()?;
165 }
166 Ok(())
167 }
168
169 fn reserve_index_entry(&mut self) -> io::Result<()> {
170 if let Some(limit) = self.index_limit {
171 if self.index.len() >= limit {
172 return Err(io::Error::other(
173 "position output directory exceeds compaction scratch budget",
174 ));
175 }
176 if self.index.len() == self.index.capacity() {
177 let capacity = self.index.capacity().saturating_mul(2).max(16).min(limit);
178 self.index.reserve_exact(capacity - self.index.len());
179 }
180 }
181 if self.compact && self.descriptors.len() == self.descriptors.capacity() {
182 let limit = self.index_limit.unwrap_or(usize::MAX);
183 let capacity = self
184 .descriptors
185 .capacity()
186 .saturating_mul(2)
187 .max(16)
188 .min(limit);
189 self.descriptors
190 .reserve_exact(capacity - self.descriptors.len());
191 }
192 Ok(())
193 }
194
195 fn flush_block(&mut self) -> io::Result<()> {
196 if self.pending.is_empty() {
197 return Ok(());
198 }
199 if self.written > u32::MAX as u64 {
200 return Err(io::Error::new(
201 io::ErrorKind::InvalidData,
202 "position stream exceeds u32::MAX bytes",
203 ));
204 }
205 self.reserve_index_entry()?;
206 self.index
207 .push((self.written as u32, self.total - self.pending.len() as u64));
208 let count = self.pending.len();
209 self.scratch.clear();
210 self.scratch.resize(BLOCK_HEADER, 0);
211 self.scratch[0..2].copy_from_slice(&(count as u16).to_le_bytes());
212 if self.codec.for_count(count) == PostingCodec::Simd4x {
213 let width = bitpacking4x::encode(&self.pending, &mut self.scratch);
214 self.scratch[2] = width;
215 self.scratch[3] = 1;
216 } else {
217 let max = self.pending.iter().copied().max().unwrap_or(0);
218 let width = simd::RoundedBitWidth::from_exact(simd::bits_needed(max));
219 self.scratch
220 .resize(BLOCK_HEADER + count * width.bytes_per_value(), 0);
221 self.scratch[2] = width.as_u8();
222 simd::pack_rounded(&self.pending, width, &mut self.scratch[BLOCK_HEADER..]);
223 }
224 if self.compact {
225 self.descriptors.push(directory::descriptor(&self.scratch)?);
226 self.writer.write_all(&self.scratch[BLOCK_HEADER..])?;
227 self.written += (self.scratch.len() - BLOCK_HEADER) as u64;
228 } else {
229 self.writer.write_all(&self.scratch)?;
230 self.written += self.scratch.len() as u64;
231 }
232 self.pending.clear();
233 Ok(())
234 }
235
236 fn append_encoded_block(&mut self, bytes: &[u8]) -> io::Result<()> {
237 self.unique_positions = false;
238 let count = PositionStream::block_count(bytes).ok_or_else(|| {
239 io::Error::new(io::ErrorKind::InvalidData, "invalid copied position block")
240 })?;
241 self.flush_block()?;
242 self.reserve_index_entry()?;
243 let offset = u32::try_from(self.written)
244 .map_err(|_| io::Error::other("position output exceeds u32 offsets"))?;
245 self.index.push((offset, self.total));
246 let payload = if self.compact {
247 self.descriptors.push(directory::descriptor(bytes)?);
248 &bytes[BLOCK_HEADER..]
249 } else {
250 bytes
251 };
252 self.writer.write_all(payload)?;
253 self.written += payload.len() as u64;
254 self.total += count as u64;
255 Ok(())
256 }
257
258 pub fn finish(self) -> io::Result<(u64, u64)> {
261 self.finish_checked(|| Ok(()))
262 }
263
264 fn finish_checked(
265 mut self,
266 mut check: impl FnMut() -> io::Result<()>,
267 ) -> io::Result<(u64, u64)> {
268 check()?;
269 self.flush_block()?;
270 if self.compact {
271 directory::write(&mut self.writer, &self.index, &self.descriptors, &mut check)?;
272 } else {
273 for (i, &(offset, value_start)) in self.index.iter().enumerate() {
274 if i.is_multiple_of(4096) {
275 check()?;
276 }
277 self.writer.write_u32::<LittleEndian>(offset)?;
278 self.writer.write_u64::<LittleEndian>(value_start)?;
279 }
280 }
281 self.writer
282 .write_u32::<LittleEndian>(block_count_and_flags(
283 self.index.len(),
284 self.unique_positions,
285 )?)?;
286 self.writer.write_u64::<LittleEndian>(self.total)?;
287 self.writer
288 .write_u32::<LittleEndian>(footer_magic(self.compact))?;
289 let index_len = if self.compact {
290 directory::directory_len(self.index.len()).unwrap()
291 } else {
292 self.index.len() * INDEX_ENTRY
293 };
294 let bytes = self.written + index_len as u64 + FOOTER as u64;
295 Ok((self.total, bytes))
296 }
297}
298
299#[derive(Debug, Clone)]
301pub struct PositionStream {
302 bytes: OwnedBytes,
303 num_blocks: usize,
304 index_start: usize,
305 total: u64,
306 canonical_blocks: bool,
307 compact: bool,
308}
309
310#[derive(Default)]
311struct PositionBlockCache {
312 index: Option<usize>,
313 value_start: u64,
314 values: Vec<u32>,
315 #[cfg(test)]
316 decodes: usize,
317 #[cfg(test)]
318 lookups: usize,
319}
320
321impl PositionBlockCache {
322 #[inline]
323 fn deltas(&self, cursor: u64, tf: u32) -> Option<&[u32]> {
324 self.index?;
325 let start = usize::try_from(cursor.checked_sub(self.value_start)?).ok()?;
326 self.values.get(start..start.checked_add(tf as usize)?)
327 }
328}
329
330impl PositionStream {
331 pub fn is_stream(raw: &[u8]) -> bool {
333 directory::is_compact(raw)
334 || (raw.len() >= FOOTER
335 && u32::from_le_bytes(raw[raw.len() - 4..].try_into().unwrap()) == MAGIC)
336 }
337
338 pub fn open(bytes: OwnedBytes) -> io::Result<Self> {
339 let (num_blocks, index_start, total) = Self::parse_layout(&bytes)?;
340 Self::validate_blocks(&bytes, total)?;
341 Ok(Self::from_layout(bytes, num_blocks, index_start, total))
342 }
343
344 pub(super) fn open_for_query(bytes: OwnedBytes) -> io::Result<Self> {
346 let (num_blocks, index_start, total) = Self::parse_layout(&bytes)?;
347 Ok(Self::from_layout(bytes, num_blocks, index_start, total))
348 }
349
350 fn from_layout(bytes: OwnedBytes, num_blocks: usize, index_start: usize, total: u64) -> Self {
351 let canonical_blocks = num_blocks == 0
355 || Self::entry_for(&bytes, index_start, num_blocks, num_blocks - 1).1
356 == (num_blocks as u64 - 1) * POSITION_STREAM_BLOCK as u64;
357 Self {
358 compact: directory::is_compact(&bytes),
359 bytes,
360 num_blocks,
361 index_start,
362 total,
363 canonical_blocks,
364 }
365 }
366
367 fn parse_layout(raw: &[u8]) -> io::Result<(usize, usize, u64)> {
368 Self::parse_layout_tail(raw, raw.len())
369 }
370
371 fn parse_layout_tail(raw: &[u8], total_len: usize) -> io::Result<(usize, usize, u64)> {
372 if !Self::is_stream(raw) {
373 return Err(io::Error::new(
374 io::ErrorKind::InvalidData,
375 "unsupported position stream format; rebuild the index",
376 ));
377 }
378 let footer_start = raw.len() - FOOTER;
379 let num_blocks =
380 (u32::from_le_bytes(raw[footer_start..footer_start + 4].try_into().unwrap())
381 & !UNIQUE_POSITIONS) as usize;
382 let total =
383 u64::from_le_bytes(raw[footer_start + 4..footer_start + 12].try_into().unwrap());
384 let index_len = (if directory::is_compact(raw) {
385 directory::directory_len(num_blocks)
386 } else {
387 num_blocks.checked_mul(INDEX_ENTRY)
388 })
389 .ok_or_else(|| {
390 io::Error::new(io::ErrorKind::InvalidData, "position block index overflows")
391 })?;
392 let index_start = total_len
393 .saturating_sub(FOOTER)
394 .checked_sub(index_len)
395 .ok_or_else(|| {
396 io::Error::new(
397 io::ErrorKind::InvalidData,
398 "position block index longer than the stream",
399 )
400 })?;
401 Ok((num_blocks, index_start, total))
402 }
403
404 pub fn total_positions(&self) -> u64 {
405 self.total
406 }
407
408 pub fn num_blocks(&self) -> usize {
409 self.num_blocks
410 }
411
412 #[inline]
413 fn index_entry(raw: &[u8], index_start: usize, idx: usize) -> (usize, u64) {
414 let p = index_start + idx * INDEX_ENTRY;
415 (
416 u32::from_le_bytes(raw[p..p + 4].try_into().unwrap()) as usize,
417 u64::from_le_bytes(raw[p + 4..p + 12].try_into().unwrap()),
418 )
419 }
420
421 fn entry_for(raw: &[u8], index_start: usize, blocks: usize, idx: usize) -> (usize, u64) {
422 if directory::is_compact(raw) {
423 directory::entry(&raw[index_start..raw.len() - FOOTER], blocks, idx)
424 } else {
425 Self::index_entry(raw, index_start, idx)
426 }
427 }
428
429 #[inline]
430 fn entry(&self, idx: usize) -> (usize, u64) {
431 if self.compact {
432 directory::entry(
433 &self.bytes[self.index_start..self.bytes.len() - FOOTER],
434 self.num_blocks,
435 idx,
436 )
437 } else {
438 Self::index_entry(self.bytes.as_slice(), self.index_start, idx)
439 }
440 }
441
442 fn block_range(&self, idx: usize) -> Option<(usize, usize, u64)> {
443 if idx >= self.num_blocks {
444 return None;
445 }
446 let (start, value_start) = self.entry(idx);
447 let end = if self.compact {
448 let tag = directory::tag(&self.bytes[self.index_start..], self.num_blocks, idx);
449 start + directory::parts(tag).ok()?.3
450 } else if idx + 1 < self.num_blocks {
451 self.entry(idx + 1).0
452 } else {
453 self.index_start
454 };
455 (start <= end && end <= self.index_start).then_some((start, end, value_start))
456 }
457
458 fn block_count(raw: &[u8]) -> Option<usize> {
459 if raw.len() < BLOCK_HEADER {
460 return None;
461 }
462 let count = u16::from_le_bytes(raw[0..2].try_into().unwrap()) as usize;
463 if count == 0 || count > POSITION_STREAM_BLOCK {
464 return None;
465 }
466 let payload_len = match (raw[3], raw[2]) {
469 (0, 0 | 8 | 16 | 32) => count * (raw[2] as usize / 8),
470 (1, 0..=32) if count == POSITION_STREAM_BLOCK => {
471 bitpacking4x::encoded_len(count, raw[2])
472 }
473 _ => return None,
474 };
475 (raw.len() == BLOCK_HEADER + payload_len).then_some(count)
476 }
477
478 fn locate_value(&self, cursor: u64, forward_from: Option<usize>) -> Option<(usize, usize)> {
479 if cursor >= self.total || self.num_blocks == 0 {
480 return None;
481 }
482 if self.canonical_blocks {
483 let idx = usize::try_from(cursor / POSITION_STREAM_BLOCK as u64).ok()?;
484 return Some((idx, (cursor % POSITION_STREAM_BLOCK as u64) as usize));
485 }
486 if self.compact {
487 return directory::locate(
488 &self.bytes[self.index_start..],
489 self.num_blocks,
490 cursor,
491 forward_from,
492 );
493 }
494 let mut low = 0usize;
495 let mut high = self.num_blocks;
496 if let Some(first) = forward_from {
497 low = first.min(self.num_blocks);
498 high = low.saturating_add(1).min(self.num_blocks);
499 let mut step = 1usize;
500 while high < self.num_blocks && self.entry(high).1 <= cursor {
501 low = high;
502 step = step.saturating_mul(2);
503 high = high.saturating_add(step).min(self.num_blocks);
504 }
505 }
506 while low < high {
507 let mid = low + (high - low) / 2;
508 let value_start = self.entry(mid).1;
509 if value_start <= cursor {
510 low = mid + 1;
511 } else {
512 high = mid;
513 }
514 }
515 let idx = low.checked_sub(1)?;
516 let value_start = self.entry(idx).1;
519 let in_block = usize::try_from(cursor.checked_sub(value_start)?).ok()?;
520 Some((idx, in_block))
521 }
522
523 pub fn decode_block(&self, idx: usize, out: &mut Vec<u32>) -> bool {
525 let Some((start, end, _)) = self.block_range(idx) else {
526 return false;
527 };
528 self.decode_payload(idx, &self.bytes.as_slice()[start..end], out)
529 }
530
531 fn decode_payload(&self, idx: usize, raw: &[u8], out: &mut Vec<u32>) -> bool {
532 if !self.compact {
533 return Self::decode_block_bytes(raw, out);
534 }
535 let tag = directory::tag(&self.bytes[self.index_start..], self.num_blocks, idx);
536 let Ok((count, width, codec, bytes)) = directory::parts(tag) else {
537 return false;
538 };
539 if bytes != raw.len() {
540 return false;
541 }
542 out.resize(count, 0);
543 crate::observe::search_work!(
544 position_blocks += 1,
545 position_values += count,
546 position_payload_bytes += bytes
547 );
548 if codec == 1 {
549 bitpacking4x::decode(raw, width, out);
550 } else {
551 simd::unpack_rounded(
552 raw,
553 simd::RoundedBitWidth::try_from_u8(width).unwrap(),
554 out,
555 count,
556 );
557 }
558 true
559 }
560
561 fn decode_block_bytes(raw: &[u8], out: &mut Vec<u32>) -> bool {
562 let Some(count) = Self::block_count(raw) else {
563 return false;
564 };
565 out.resize(count, 0);
566 crate::observe::search_work!(
567 position_blocks += 1,
568 position_values += count,
569 position_payload_bytes += raw.len() - BLOCK_HEADER
570 );
571 if raw[3] == 1 {
572 bitpacking4x::decode(&raw[BLOCK_HEADER..], raw[2], out);
573 } else {
574 let Some(width) = simd::RoundedBitWidth::try_from_u8(raw[2]) else {
577 out.clear();
578 return false;
579 };
580 simd::unpack_rounded(&raw[BLOCK_HEADER..], width, out, count);
581 }
582 true
583 }
584
585 pub fn read_into(
593 &self,
594 cursor: u64,
595 tf: u32,
596 scratch: &mut Vec<u32>,
597 out: &mut Vec<u32>,
598 ) -> bool {
599 let mut cache = PositionBlockCache {
600 values: std::mem::take(scratch),
601 ..Default::default()
602 };
603 let found = self.read_cached(cursor, tf, &mut cache, out);
604 *scratch = cache.values;
605 found
606 }
607
608 #[inline]
609 fn read_cached(
610 &self,
611 cursor: u64,
612 tf: u32,
613 cache: &mut PositionBlockCache,
614 out: &mut Vec<u32>,
615 ) -> bool {
616 out.clear();
617 if tf == 0 {
618 return true;
619 }
620 if cursor
621 .checked_add(tf as u64)
622 .is_none_or(|end| end > self.total)
623 {
624 return false;
625 }
626 if let Some(deltas) = cache.deltas(cursor, tf) {
629 let mut position = 0u32;
630 out.extend(deltas.iter().map(|&delta| {
631 position = position.wrapping_add(delta);
632 position
633 }));
634 return true;
635 }
636 self.read_uncached(cursor, tf, cache, out)
637 }
638
639 fn read_uncached(
640 &self,
641 cursor: u64,
642 tf: u32,
643 cache: &mut PositionBlockCache,
644 out: &mut Vec<u32>,
645 ) -> bool {
646 let located = if self.canonical_blocks {
647 self.locate_value(cursor, None)
648 } else {
649 let cached = cache.index.and_then(|idx| {
650 let offset = cursor.checked_sub(cache.value_start)?;
651 (offset < cache.values.len() as u64).then_some((idx, offset as usize))
652 });
653 if cached.is_some() {
654 cached
655 } else {
656 #[cfg(test)]
657 {
658 cache.lookups += 1;
659 }
660 let forward = cache
661 .index
662 .filter(|_| cursor >= cache.value_start)
663 .map(|idx| idx + 1);
664 self.locate_value(cursor, forward)
665 }
666 };
667 let Some((mut idx, mut in_block)) = located else {
668 return false;
669 };
670 let mut remaining = tf as usize;
671 let mut next_cursor = cursor;
672 let mut prev = 0u32;
673 out.reserve(remaining);
674 while remaining > 0 {
675 if cache.index != Some(idx) {
676 cache.index = None;
677 let Some((start, end, value_start)) = self.block_range(idx) else {
678 return false;
679 };
680 if value_start.checked_add(in_block as u64) != Some(next_cursor)
683 || !self.decode_payload(
684 idx,
685 &self.bytes.as_slice()[start..end],
686 &mut cache.values,
687 )
688 {
689 return false;
690 }
691 cache.value_start = value_start;
692 cache.index = Some(idx);
693 #[cfg(test)]
694 {
695 cache.decodes += 1;
696 }
697 }
698 if cache.values.len() <= in_block {
699 return false;
700 }
701 let take = remaining.min(cache.values.len() - in_block);
702 for &delta in &cache.values[in_block..in_block + take] {
703 prev = prev.wrapping_add(delta);
704 out.push(prev);
705 }
706 remaining -= take;
707 next_cursor += take as u64;
708 in_block = 0;
709 idx += 1;
710 }
711 true
712 }
713
714 pub fn concatenate_streaming<W: Write>(
717 sources: &[&[u8]],
718 writer: &mut W,
719 ) -> crate::Result<(u64, u64)> {
720 let layouts: Vec<_> = sources
721 .iter()
722 .map(|raw| Self::parse_layout(raw))
723 .collect::<io::Result<_>>()?;
724
725 if sources.len() == 1 {
728 let raw = sources[0];
729 let total = layouts[0].2;
730 Self::validate_blocks(raw, total)?;
731 writer.write_all(raw)?;
732 return Ok((total, raw.len() as u64));
733 }
734
735 if sources.iter().any(|raw| directory::is_compact(raw)) {
736 let mut encoder = PositionStreamEncoder::new(writer).with_compact_directory();
737 let mut block = Vec::with_capacity(BLOCK_HEADER + POSITION_STREAM_BLOCK * 4);
738 for (raw, &(blocks, index_start, total)) in sources.iter().zip(&layouts) {
739 Self::validate_blocks(raw, total)?;
740 for idx in 0..blocks {
741 let (start, _) = Self::entry_for(raw, index_start, blocks, idx);
742 let end = if idx + 1 == blocks {
743 index_start
744 } else {
745 Self::entry_for(raw, index_start, blocks, idx + 1).0
746 };
747 if directory::is_compact(raw) {
748 let tag = directory::tag(&raw[index_start..], blocks, idx);
749 let (count, width, codec, _) = directory::parts(tag)?;
750 block.clear();
751 block.extend_from_slice(&(count as u16).to_le_bytes());
752 block.extend_from_slice(&[width, codec]);
753 block.extend_from_slice(&raw[start..end]);
754 encoder.append_encoded_block(&block)?;
755 } else {
756 encoder.append_encoded_block(&raw[start..end])?;
757 }
758 }
759 }
760 encoder.unique_positions = sources.iter().all(|raw| has_unique_positions(raw));
761 return Ok(encoder.finish()?);
762 }
763
764 let total_blocks: usize = layouts.iter().map(|(blocks, _, _)| *blocks).sum();
765 let count_and_flags = block_count_and_flags(
766 total_blocks,
767 sources.iter().all(|raw| has_unique_positions(raw)),
768 )?;
769 let mut out_index = Vec::with_capacity(total_blocks * INDEX_ENTRY);
770 let mut data_written = 0u64;
771 let mut total_positions = 0u64;
772
773 for (raw, &(num_blocks, index_start, source_total)) in sources.iter().zip(&layouts) {
774 let mut expected_start = 0u64;
775 for idx in 0..num_blocks {
776 let (start, value_start) = Self::index_entry(raw, index_start, idx);
777 let end = if idx + 1 < num_blocks {
778 Self::index_entry(raw, index_start, idx + 1).0
779 } else {
780 index_start
781 };
782 if start > end || end > index_start || value_start != expected_start {
783 return Err(crate::Error::Corruption(
784 "invalid position block index during merge".into(),
785 ));
786 }
787 let block = &raw[start..end];
788 let count = Self::block_count(block).ok_or_else(|| {
789 crate::Error::Corruption("invalid position block during merge".into())
790 })?;
791 if data_written > u32::MAX as u64 {
792 return Err(io::Error::new(
793 io::ErrorKind::InvalidData,
794 "position stream exceeds u32::MAX bytes during merge",
795 )
796 .into());
797 }
798 out_index.write_u32::<LittleEndian>(data_written as u32)?;
799 out_index.write_u64::<LittleEndian>(
800 total_positions.checked_add(value_start).ok_or_else(|| {
801 io::Error::new(
802 io::ErrorKind::InvalidData,
803 "position count overflows u64 during merge",
804 )
805 })?,
806 )?;
807 writer.write_all(block)?;
808 data_written += block.len() as u64;
809 expected_start += count as u64;
810 }
811 if expected_start != source_total {
812 return Err(crate::Error::Corruption(
813 "position stream total does not match its blocks".into(),
814 ));
815 }
816 total_positions = total_positions.checked_add(source_total).ok_or_else(|| {
817 io::Error::new(
818 io::ErrorKind::InvalidData,
819 "position count overflows u64 during merge",
820 )
821 })?;
822 }
823
824 writer.write_all(&out_index)?;
825 writer.write_u32::<LittleEndian>(count_and_flags)?;
826 writer.write_u64::<LittleEndian>(total_positions)?;
827 writer.write_u32::<LittleEndian>(footer_magic(false))?;
828 let bytes_written = data_written + out_index.len() as u64 + FOOTER as u64;
829 Ok((total_positions, bytes_written))
830 }
831
832 fn validate_blocks(raw: &[u8], expected_total: u64) -> io::Result<()> {
833 let (num_blocks, index_start, _) = Self::parse_layout(raw)?;
834 if directory::is_compact(raw) {
835 return directory::validate(
836 &raw[index_start..raw.len() - FOOTER],
837 num_blocks,
838 index_start,
839 expected_total,
840 );
841 }
842 if (num_blocks == 0 && index_start != 0)
843 || (num_blocks != 0 && Self::index_entry(raw, index_start, 0).0 != 0)
844 {
845 return Err(io::Error::new(
846 io::ErrorKind::InvalidData,
847 "invalid position block index: unaddressed payload prefix",
848 ));
849 }
850 let mut total = 0u64;
851 for idx in 0..num_blocks {
852 let (start, value_start) = Self::index_entry(raw, index_start, idx);
853 let end = if idx + 1 < num_blocks {
854 Self::index_entry(raw, index_start, idx + 1).0
855 } else {
856 index_start
857 };
858 if start > end || end > index_start || value_start != total {
859 return Err(io::Error::new(
860 io::ErrorKind::InvalidData,
861 "invalid position block index",
862 ));
863 }
864 total += Self::block_count(&raw[start..end]).ok_or_else(|| {
865 io::Error::new(io::ErrorKind::InvalidData, "invalid position block")
866 })? as u64;
867 }
868 if total != expected_total {
869 return Err(io::Error::new(
870 io::ErrorKind::InvalidData,
871 "position stream total does not match its blocks",
872 ));
873 }
874 Ok(())
875 }
876}
877
878#[derive(Debug, Clone)]
880pub struct TermPositions(pub(super) PositionStream);
881
882pub(crate) struct TermPositionCursor {
885 positions: TermPositions,
886 cache: PositionBlockCache,
887}
888
889impl TermPositionCursor {
890 #[inline]
894 pub(crate) fn read_one(&mut self, cursor: u64, scratch: &mut Vec<u32>) -> Option<u32> {
895 if let Some(&position) = self
896 .cache
897 .deltas(cursor, 1)
898 .and_then(|values| values.first())
899 {
900 crate::observe::search_work!(position_reads += 1, positions_requested += 1);
901 return Some(position);
902 }
903 self.read_into(cursor, 1, scratch).then(|| scratch[0])
904 }
905
906 #[inline]
909 pub(crate) fn contains(
910 &mut self,
911 cursor: u64,
912 tf: u32,
913 target: u32,
914 scratch: &mut Vec<u32>,
915 ) -> bool {
916 if tf == 1 {
917 return self.read_one(cursor, scratch) == Some(target);
918 }
919 if let Some(deltas) = self.cache.deltas(cursor, tf) {
920 crate::observe::search_work!(position_reads += 1, positions_requested += tf);
921 let mut position = 0u32;
922 for &delta in deltas {
923 position = position.wrapping_add(delta);
924 if position >= target {
925 return position == target;
926 }
927 }
928 return false;
929 }
930 self.read_into(cursor, tf, scratch) && scratch.binary_search(&target).is_ok()
931 }
932
933 pub(crate) fn read_into(&mut self, cursor: u64, tf: u32, out: &mut Vec<u32>) -> bool {
934 crate::observe::search_work!(position_reads += 1, positions_requested += tf);
935 self.positions
936 .0
937 .read_cached(cursor, tf, &mut self.cache, out)
938 }
939}
940
941impl TermPositions {
942 pub(crate) fn has_unique_positions(&self) -> bool {
943 has_unique_positions(&self.0.bytes)
944 }
945
946 #[cfg(all(feature = "native", test))]
948 pub(crate) fn has_compact_directory(&self) -> bool {
949 self.0.compact
950 }
951
952 pub(crate) fn into_cursor(self) -> TermPositionCursor {
953 TermPositionCursor {
954 positions: self,
955 cache: PositionBlockCache::default(),
956 }
957 }
958 pub fn open(bytes: OwnedBytes) -> io::Result<Self> {
959 Ok(Self(PositionStream::open(bytes)?))
960 }
961
962 pub fn positions_into(
965 &self,
966 cursor: u64,
967 tf: u32,
968 scratch: &mut Vec<u32>,
969 out: &mut Vec<u32>,
970 ) -> bool {
971 self.0.read_into(cursor, tf, scratch, out)
972 }
973
974 pub fn positions(&self, cursor: u64, tf: u32) -> Option<Vec<u32>> {
980 let mut out = Vec::new();
981 let mut scratch = Vec::new();
982 self.positions_into(cursor, tf, &mut scratch, &mut out)
983 .then_some(out)
984 }
985}
986
987#[cfg(test)]
988mod compact_directory_tests {
989 use super::*;
990
991 fn encode(values: &[u32], compact: bool, codec: PostingCodec) -> Vec<u8> {
992 let mut bytes = Vec::new();
993 let mut encoder = PositionStreamEncoder::with_posting_codec(&mut bytes, codec);
994 if compact {
995 encoder = encoder.with_compact_directory();
996 }
997 encoder.push_values(values).unwrap();
998 encoder.finish().unwrap();
999 bytes
1000 }
1001
1002 #[test]
1003 fn singleton_reads_preserve_cached_values_across_short_blocks_and_reverse_probes() {
1004 for codec in [PostingCodec::Rounded, PostingCodec::Simd4x] {
1005 for compact in [false, true] {
1006 let values: Vec<_> = (0..267u32).map(|i| i.wrapping_mul(1234567)).collect();
1007 let first = encode(&values[..13], compact, codec);
1008 let second = encode(&values[13..], compact, codec);
1009 let mut bytes = Vec::new();
1010 PositionStream::concatenate_streaming(&[&first, &second], &mut bytes).unwrap();
1011 let original = bytes.clone();
1012 let mut cursor = TermPositions::open(OwnedBytes::new(bytes))
1013 .unwrap()
1014 .into_cursor();
1015 let mut scratch = Vec::new();
1016 for index in (0..values.len()).chain((0..values.len()).rev()) {
1017 assert_eq!(
1018 cursor.read_one(index as u64, &mut scratch),
1019 Some(values[index])
1020 );
1021 let decodes = cursor.cache.decodes;
1022 assert_eq!(
1023 cursor.read_one(index as u64, &mut scratch),
1024 Some(values[index])
1025 );
1026 assert_eq!(cursor.cache.decodes, decodes);
1027 }
1028 assert_eq!(cursor.read_one(values.len() as u64, &mut scratch), None);
1029 assert_eq!(cursor.read_one(u64::MAX, &mut scratch), None);
1030 assert_eq!(cursor.positions.0.bytes.as_slice(), original);
1031 }
1032 }
1033 }
1034
1035 #[test]
1036 fn compact_positions_preserve_payloads_and_every_cursor_across_mixed_short_blocks() {
1037 for codec in [PostingCodec::Rounded, PostingCodec::Simd4x] {
1038 for count in [0, 1, 2, 127, 128, 129, 1024, 1025, 4097] {
1039 let values: Vec<u32> = (0..count).map(|i| (i * 31 % 257) as u32).collect();
1040 let old = encode(&values, false, codec);
1041 let new = encode(&values, true, codec);
1042 let interleaved = PositionStream::open(OwnedBytes::new(old.clone())).unwrap();
1043 let compact = PositionStream::open(OwnedBytes::new(new.clone())).unwrap();
1044 assert!(compact.compact);
1045 assert_eq!(
1046 new.len(),
1047 compact.index_start
1048 + directory::directory_len(compact.num_blocks).unwrap()
1049 + FOOTER
1050 );
1051 for i in 0..compact.num_blocks {
1052 let (a, b, _) = interleaved.block_range(i).unwrap();
1053 let (c, d, _) = compact.block_range(i).unwrap();
1054 assert_eq!(&old[a + BLOCK_HEADER..b], &new[c..d]);
1055 }
1056 let tail = encode(&[7, 3, 11], true, codec);
1057 let mut merged = Vec::new();
1058 PositionStream::concatenate_streaming(&[&tail, &old, &new, &tail], &mut merged)
1059 .unwrap();
1060 let stream = PositionStream::open(OwnedBytes::new(merged)).unwrap();
1061 let expected: Vec<_> =
1062 [vec![7, 3, 11], values.clone(), values, vec![7, 3, 11]].concat();
1063 let mut scratch = Vec::new();
1064 let mut out = Vec::new();
1065 for (i, &value) in expected.iter().enumerate() {
1066 assert!(stream.read_into(i as u64, 1, &mut scratch, &mut out));
1067 assert_eq!(out, [value]);
1068 }
1069 }
1070 }
1071 }
1072
1073 #[test]
1074 fn compact_directory_rejects_bad_structure_without_interpreting_payload_values() {
1075 let bytes = encode(&vec![7; 1153], true, PostingCodec::Rounded);
1076 let (_, start, _) = PositionStream::parse_layout(&bytes).unwrap();
1077 let mut changed = bytes.clone();
1079 changed[..start].fill(255);
1080 assert!(PositionStream::open(OwnedBytes::new(changed)).is_ok());
1081 for at in [
1082 start,
1083 start + 4,
1084 start + 12,
1085 start + 16,
1086 bytes.len() - FOOTER - 1,
1087 ] {
1088 let mut corrupt = bytes.clone();
1089 corrupt[at] ^= 128;
1090 assert!(
1091 PositionStream::open(OwnedBytes::new(corrupt)).is_err(),
1092 "byte {at}"
1093 );
1094 }
1095 for n in 0..bytes.len() {
1096 assert!(PositionStream::open(OwnedBytes::new(bytes[..n].to_vec())).is_err());
1097 }
1098 }
1099}
1100
1101#[cfg(test)]
1102mod tests {
1103 use super::*;
1104
1105 #[test]
1106 fn explicit_position_open_rejects_overflowing_checkpoint_before_deriving_layout() {
1107 let mut bytes = Vec::new();
1108 let mut encoder = PositionStreamEncoder::new(&mut bytes).with_compact_directory();
1109 encoder.push_values(&[1; 256]).unwrap();
1110 encoder.finish().unwrap();
1111 let (_, index_start, _) = PositionStream::parse_layout(&bytes).unwrap();
1112 bytes[index_start + 4..index_start + 12].copy_from_slice(&u64::MAX.to_le_bytes());
1113 assert!(PositionStream::open(OwnedBytes::new(bytes)).is_err());
1114 }
1115
1116 #[test]
1117 fn position_decoding_overwrites_stale_values_across_lengths_and_codecs() {
1118 let mut out = vec![u32::MAX; 256];
1119 for codec in [PostingCodec::Rounded, PostingCodec::Simd4x] {
1120 for count in [128, 1, 127, 128, 17] {
1121 for value in [0, 1, 255, 65536] {
1122 for compact in [false, true] {
1123 let mut bytes = Vec::new();
1124 let mut encoder =
1125 PositionStreamEncoder::with_posting_codec(&mut bytes, codec);
1126 if compact {
1127 encoder = encoder.with_compact_directory();
1128 }
1129 encoder.push_values(&vec![value; count]).unwrap();
1130 encoder.finish().unwrap();
1131 let stream = PositionStream::open(OwnedBytes::new(bytes)).unwrap();
1132 out.fill(u32::MAX);
1133 assert!(stream.decode_block(0, &mut out));
1134 assert_eq!(out, vec![value; count]);
1135 }
1136 }
1137 }
1138 }
1139 }
1140
1141 #[test]
1142 fn admitted_mixed_position_directories_locate_every_cursor_across_copied_short_blocks() {
1143 let counts = [1, 127, 128, 3, 65, 128, 7];
1144 let mut sources = Vec::new();
1145 let mut expected = Vec::new();
1146 for (block, &count) in counts.iter().enumerate() {
1147 let mut bytes = Vec::new();
1148 let codec = if block % 2 == 0 {
1149 PostingCodec::Rounded
1150 } else {
1151 PostingCodec::Simd4x
1152 };
1153 let mut encoder = PositionStreamEncoder::with_posting_codec(&mut bytes, codec);
1154 encoder.push_values(&vec![1; count]).unwrap();
1155 encoder.finish().unwrap();
1156 sources.push(bytes);
1157 expected.extend((0..count).map(|offset| (block, offset)));
1158 }
1159 let mut bytes = Vec::new();
1160 PositionStream::concatenate_streaming(
1161 &sources.iter().map(Vec::as_slice).collect::<Vec<_>>(),
1162 &mut bytes,
1163 )
1164 .unwrap();
1165 let original = bytes.clone();
1166 let stream = PositionStream::open(OwnedBytes::new(bytes)).unwrap();
1167 assert!(!stream.canonical_blocks);
1168 for cursor in (0..expected.len()).rev().chain(0..expected.len()) {
1169 let expected = expected[cursor];
1170 assert_eq!(stream.locate_value(cursor as u64, None), Some(expected));
1171 for first in 0..=expected.0 {
1172 assert_eq!(
1173 stream.locate_value(cursor as u64, Some(first)),
1174 Some(expected)
1175 );
1176 }
1177 let mut values = Vec::new();
1178 assert!(stream.read_into(cursor as u64, 1, &mut Vec::new(), &mut values));
1179 assert_eq!(values, [1]);
1180 }
1181 assert_eq!(stream.locate_value(expected.len() as u64, None), None);
1182 assert_eq!(stream.locate_value(u64::MAX, Some(usize::MAX)), None);
1183 assert_eq!(stream.bytes.as_slice(), original);
1184 }
1185
1186 #[test]
1187 fn position_uniqueness_is_certified_per_document_and_conjoined_by_copying_merge() {
1188 for compact in [false, true] {
1189 let encode = |docs: &[Vec<u32>]| {
1190 let mut bytes = Vec::new();
1191 let mut encoder = PositionStreamEncoder::new(&mut bytes);
1192 if compact {
1193 encoder = encoder.with_compact_directory();
1194 }
1195 for doc in docs {
1196 encoder.push_doc(&mut doc.clone()).unwrap();
1197 }
1198 encoder.finish().unwrap();
1199 bytes
1200 };
1201 let unique = encode(&[vec![7, 0, 1], vec![0, 3, u32::MAX]]);
1202 let repeated = encode(&[vec![0, 3, 3]]);
1203 assert!(has_unique_positions(&unique));
1204 assert!(!has_unique_positions(&repeated));
1205 assert_eq!(directory::is_compact(&unique), compact);
1206 let mut uncertified = unique.clone();
1207 let end = uncertified.len();
1208 let count = u32::from_le_bytes(
1209 uncertified[end - FOOTER..end - FOOTER + 4]
1210 .try_into()
1211 .unwrap(),
1212 ) & !UNIQUE_POSITIONS;
1213 uncertified[end - FOOTER..end - FOOTER + 4].copy_from_slice(&count.to_le_bytes());
1214 assert!(TermPositions::open(OwnedBytes::new(uncertified.clone())).is_ok());
1215 for second in [&unique, &repeated, &uncertified] {
1216 let mut merged = Vec::new();
1217 PositionStream::concatenate_streaming(&[&unique, second], &mut merged).unwrap();
1218 let result = PositionStream::open(OwnedBytes::new(merged.clone())).unwrap();
1219 assert_eq!(has_unique_positions(&merged), has_unique_positions(second));
1220 let mut block = 0;
1222 for source in [&unique, second] {
1223 let source = PositionStream::open(OwnedBytes::new(source.clone())).unwrap();
1224 for index in 0..source.num_blocks() {
1225 let (a, b, _) = source.block_range(index).unwrap();
1226 let (c, d, _) = result.block_range(block).unwrap();
1227 assert_eq!(&source.bytes[a..b], &result.bytes[c..d]);
1228 block += 1;
1229 }
1230 }
1231 }
1232 let mut raw = Vec::new();
1233 let mut encoder = PositionStreamEncoder::new(&mut raw);
1234 encoder.push_values(&[0, 1, 2]).unwrap();
1235 encoder.finish().unwrap();
1236 assert!(
1237 !has_unique_positions(&raw),
1238 "raw values cannot prove document boundaries"
1239 );
1240 }
1241 }
1242
1243 #[test]
1244 fn sequential_merged_position_reads_reuse_logical_block_addresses() {
1245 let docs = vec![vec![1, 5]; 13];
1246 let (source, _) = encode(&docs);
1247 let mut bytes = Vec::new();
1248 PositionStream::concatenate_streaming(&vec![source.as_slice(); 30], &mut bytes).unwrap();
1249 let original = bytes.clone();
1250 let mut cursor = TermPositions::open(OwnedBytes::new(bytes))
1251 .unwrap()
1252 .into_cursor();
1253 let mut out = Vec::new();
1254 for doc in 0..390u32 {
1255 assert!(cursor.read_into(u64::from(doc) * 2, 2, &mut out));
1256 assert_eq!(out, [1, 5]);
1257 }
1258 assert_eq!(cursor.cache.decodes, 30);
1259 assert_eq!(
1260 cursor.cache.lookups, 30,
1261 "cached block addresses must serve all covered documents"
1262 );
1263 assert!(cursor.read_into(0, 2, &mut out));
1264 assert_eq!(out, [1, 5]);
1265 assert!(cursor.read_into(778, 2, &mut out));
1266 assert_eq!(out, [1, 5]);
1267 let stream = cursor.positions.0;
1268 assert_eq!(stream.bytes.as_slice(), original);
1269 }
1270
1271 #[test]
1272 fn merged_position_cursor_handles_forward_gaps_backward_reads_and_spanning_documents() {
1273 let mut docs = Vec::new();
1274 let mut sources = Vec::new();
1275 for source in 0..40 {
1276 let part: Vec<Vec<u32>> = (0..17)
1277 .map(|doc| {
1278 (0..1 + (source * 17 + doc) % 301)
1279 .map(|p| p * 3 + doc)
1280 .collect()
1281 })
1282 .collect();
1283 sources.push(encode(&part).0);
1284 docs.extend(part);
1285 }
1286 let mut bytes = Vec::new();
1287 PositionStream::concatenate_streaming(
1288 &sources.iter().map(Vec::as_slice).collect::<Vec<_>>(),
1289 &mut bytes,
1290 )
1291 .unwrap();
1292 let before = bytes.clone();
1293 let mut cursor = TermPositions::open(OwnedBytes::new(bytes))
1294 .unwrap()
1295 .into_cursor();
1296 let mut starts = Vec::new();
1297 let mut next = 0u64;
1298 for doc in &docs {
1299 starts.push(next);
1300 next += doc.len() as u64;
1301 }
1302 let mut out = Vec::new();
1303 for at in (0..docs.len())
1304 .step_by(7)
1305 .chain((0..docs.len()).rev())
1306 .chain(0..docs.len())
1307 {
1308 assert!(cursor.read_into(starts[at], docs[at].len() as u32, &mut out));
1309 assert_eq!(out, docs[at], "document {at}");
1310 }
1311 let stream = cursor.positions.0;
1312 assert_eq!(stream.bytes.as_slice(), before);
1313 }
1314
1315 #[test]
1316 fn failed_cross_block_position_read_does_not_publish_a_stale_cached_range() {
1317 let (a, _) = encode(&[vec![1, 5, 9]]);
1318 let (b, _) = encode(&[(0..300).collect()]);
1319 let mut bytes = Vec::new();
1320 PositionStream::concatenate_streaming(&[&a, &b], &mut bytes).unwrap();
1321 let valid = PositionStream::open(OwnedBytes::new(bytes.clone())).unwrap();
1322 let (start, _, _) = valid.block_range(2).unwrap();
1323 bytes[start + 2] = 7; assert!(TermPositions::open(OwnedBytes::new(bytes.clone())).is_err());
1325 let mut malformed = valid;
1328 malformed.bytes = OwnedBytes::new(bytes);
1329 let mut cursor = TermPositions(malformed).into_cursor();
1330 let mut out = Vec::new();
1331 assert!(cursor.read_into(0, 3, &mut out));
1332 assert!(!cursor.read_into(3, 300, &mut out));
1333 assert_eq!(cursor.cache.index, None);
1334 assert!(cursor.read_into(0, 3, &mut out));
1335 assert_eq!(out, [1, 5, 9]);
1336 }
1337
1338 #[test]
1339 fn position_open_rejects_invalid_headers_directories_and_unaddressed_bytes() {
1340 let (bytes, _) = encode(&[(0..300).collect()]);
1341 let (_, index_start, _) = PositionStream::parse_layout(&bytes).unwrap();
1342 for case in 0..7 {
1343 let mut bad = bytes.clone();
1344 match case {
1345 0 => bad[2] = 7,
1346 1 => bad[3] = 2,
1347 2 => bad[index_start..index_start + 4].copy_from_slice(&1u32.to_le_bytes()),
1348 3 => bad[index_start + 4..index_start + 12].copy_from_slice(&1u64.to_le_bytes()),
1349 4 => bad[index_start + INDEX_ENTRY..index_start + INDEX_ENTRY + 4]
1350 .copy_from_slice(&u32::MAX.to_le_bytes()),
1351 5 => bad[index_start + INDEX_ENTRY + 4..index_start + INDEX_ENTRY + 12]
1352 .copy_from_slice(&0u64.to_le_bytes()),
1353 6 => {
1354 let total_at = bad.len() - FOOTER + 4;
1355 bad[total_at..total_at + 8].copy_from_slice(&u64::MAX.to_le_bytes());
1356 }
1357 _ => unreachable!(),
1358 }
1359 assert!(
1360 PositionStream::open(OwnedBytes::new(bad)).is_err(),
1361 "case {case}"
1362 );
1363 }
1364 let (mut empty, _) = encode(&[]);
1365 assert!(PositionStream::open(OwnedBytes::new(empty.clone())).is_ok());
1366 empty.insert(0, 0);
1367 assert!(PositionStream::open(OwnedBytes::new(empty)).is_err());
1368 let stream = PositionStream::open(OwnedBytes::new(bytes.clone())).unwrap();
1369 assert_eq!(stream.bytes.as_slice(), bytes);
1370 }
1371
1372 #[test]
1373 fn position_codecs_round_trip_every_width_and_short_tail() {
1374 for codec in [PostingCodec::Rounded, PostingCodec::Simd4x] {
1375 for count in [1, 3, 31, 127, 128, 129, 257] {
1376 for width in 0..=32 {
1377 let max = if width == 0 {
1378 0
1379 } else {
1380 u32::MAX >> (32 - width)
1381 };
1382 let values: Vec<_> = (0..count)
1383 .map(|i| if i % 3 == 0 { max } else { 0 })
1384 .collect();
1385 let mut bytes = Vec::new();
1386 let mut encoder = PositionStreamEncoder::with_posting_codec(&mut bytes, codec);
1387 encoder.push_values(&values).unwrap();
1388 encoder.finish().unwrap();
1389 let stream = PositionStream::open(OwnedBytes::new(bytes.clone())).unwrap();
1390 let mut actual = Vec::new();
1391 let mut block = Vec::new();
1392 for idx in 0..stream.num_blocks() {
1393 assert!(stream.decode_block(idx, &mut block));
1394 actual.extend_from_slice(&block);
1395 let (start, end, _) = stream.block_range(idx).unwrap();
1396 assert_eq!(
1397 bytes[start + 3],
1398 u8::from(
1399 codec == PostingCodec::Simd4x
1400 && block.len() == POSITION_STREAM_BLOCK
1401 )
1402 );
1403 if bytes[start + 3] == 0 {
1404 let rounded = simd::RoundedBitWidth::from_exact(simd::bits_needed(
1405 block.iter().copied().max().unwrap(),
1406 ));
1407 let mut expected = Vec::new();
1408 expected.extend_from_slice(&(block.len() as u16).to_le_bytes());
1409 expected.extend_from_slice(&[rounded.as_u8(), 0]);
1410 for value in &block {
1411 expected.extend_from_slice(
1412 &value.to_le_bytes()[..rounded.bytes_per_value()],
1413 );
1414 }
1415 assert_eq!(&bytes[start..end], expected);
1416 }
1417 }
1418 assert_eq!(actual, values);
1419 for (at, replacement) in [(2, 33), (3, 2)] {
1420 let mut corrupt = bytes.clone();
1421 corrupt[at] = replacement;
1422 assert!(PositionStream::open(OwnedBytes::new(corrupt)).is_err());
1423 }
1424 let mut copied = Vec::new();
1425 PositionStream::concatenate_streaming(&[&bytes], &mut copied).unwrap();
1426 assert_eq!(copied, bytes);
1427 }
1428 }
1429 }
1430 }
1431
1432 #[test]
1433 fn mixed_position_codecs_copy_short_interior_blocks_and_preserve_cursors() {
1434 let mut encoded = Vec::new();
1435 let mut docs = Vec::new();
1436 for codec in [
1437 PostingCodec::Simd4x,
1438 PostingCodec::Rounded,
1439 PostingCodec::Simd4x,
1440 ] {
1441 let mut bytes = Vec::new();
1442 let mut encoder = PositionStreamEncoder::with_posting_codec(&mut bytes, codec);
1443 for count in [1, 130, 7, 127] {
1444 let mut positions: Vec<_> = (0..count).map(|i| i * 3 + 4).collect();
1445 encoder.push_doc(&mut positions).unwrap();
1446 docs.push(positions);
1447 }
1448 encoder.finish().unwrap();
1449 encoded.push(bytes);
1450 }
1451 let mut output = Vec::new();
1452 PositionStream::concatenate_streaming(
1453 &encoded.iter().map(Vec::as_slice).collect::<Vec<_>>(),
1454 &mut output,
1455 )
1456 .unwrap();
1457 let stream = PositionStream::open(OwnedBytes::new(output.clone())).unwrap();
1458 assert!(!stream.canonical_blocks);
1459 let mut next = 0;
1460 for source in encoded {
1461 let source = PositionStream::open(OwnedBytes::new(source)).unwrap();
1462 for idx in 0..source.num_blocks() {
1463 let (start, end, _) = source.block_range(idx).unwrap();
1464 let (out_start, out_end, _) = stream.block_range(next).unwrap();
1465 assert_eq!(
1466 &output[out_start..out_end],
1467 &source.bytes.as_slice()[start..end]
1468 );
1469 next += 1;
1470 }
1471 }
1472 let mut starts = vec![0u64];
1473 for doc in &docs {
1474 starts.push(starts.last().unwrap() + doc.len() as u64);
1475 }
1476 let mut cursor = TermPositions(stream).into_cursor();
1477 let mut actual = Vec::new();
1478 for i in (0..docs.len()).chain((0..docs.len()).rev()) {
1479 assert!(cursor.read_into(starts[i], docs[i].len() as u32, &mut actual));
1480 assert_eq!(actual, docs[i]);
1481 }
1482 }
1483
1484 #[test]
1485 fn term_cursor_reuses_blocks_and_isolates_terms_and_backward_seeks() {
1486 let docs: Vec<Vec<u32>> = (0..200).map(|i| vec![i, i + 3]).collect();
1487 let (bytes, _) = encode(&docs);
1488 let mut cursor = TermPositions::open(OwnedBytes::new(bytes.clone()))
1489 .unwrap()
1490 .into_cursor();
1491 let mut out = Vec::new();
1492 for (id, expected) in docs.iter().enumerate() {
1493 assert!(cursor.read_into(id as u64 * 2, 2, &mut out));
1494 assert_eq!(&out, expected);
1495 }
1496 assert_eq!(
1497 cursor.cache.decodes,
1498 400usize.div_ceil(POSITION_STREAM_BLOCK)
1499 );
1500 assert!(cursor.read_into(0, 2, &mut out));
1501 assert_eq!(out, docs[0]);
1502 let decodes = cursor.cache.decodes;
1503 assert!(cursor.read_into(2, 2, &mut out));
1504 assert_eq!(cursor.cache.decodes, decodes);
1505 let (other_bytes, _) = encode(&[vec![99, 100]]);
1506 let mut other = TermPositions::open(OwnedBytes::new(other_bytes))
1507 .unwrap()
1508 .into_cursor();
1509 assert!(other.read_into(0, 2, &mut out));
1510 assert_eq!(out, vec![99, 100]);
1511 assert!(!cursor.read_into(u64::MAX, 2, &mut out));
1512 assert!(cursor.read_into(0, 0, &mut out));
1513 assert!(out.is_empty());
1514 let stream = cursor.positions.0;
1515 assert_eq!(
1516 stream.bytes.as_slice(),
1517 bytes,
1518 "reading must not modify encoded blocks"
1519 );
1520 }
1521
1522 #[test]
1523 fn term_cursor_handles_copied_short_blocks_and_cross_block_documents() {
1524 let first = vec![vec![1, 4, 9]; 13];
1525 let second = vec![vec![10, 20, 30]; 80];
1526 let (a, _) = encode(&first);
1527 let (b, _) = encode(&second);
1528 let mut merged = Vec::new();
1529 PositionStream::concatenate_streaming(&[&a, &b], &mut merged).unwrap();
1530 let mut cursor = TermPositions::open(OwnedBytes::new(merged))
1531 .unwrap()
1532 .into_cursor();
1533 let mut out = Vec::new();
1534 for (doc, expected) in first.iter().chain(&second).enumerate() {
1535 assert!(cursor.read_into(doc as u64 * 3, 3, &mut out));
1536 assert_eq!(&out, expected);
1537 }
1538 assert_eq!(cursor.cache.decodes, 3);
1539 assert!(cursor.cache.values.capacity() <= POSITION_STREAM_BLOCK);
1540 }
1541
1542 fn encode(docs: &[Vec<u32>]) -> (Vec<u8>, u64) {
1543 let mut buf = Vec::new();
1544 let mut encoder = PositionStreamEncoder::new(&mut buf);
1545 for doc in docs {
1546 let mut positions = doc.clone();
1547 encoder.push_doc(&mut positions).unwrap();
1548 }
1549 let (total, bytes) = encoder.finish().unwrap();
1550 assert_eq!(bytes as usize, buf.len());
1551 (buf, total)
1552 }
1553
1554 fn read_all(stream: &PositionStream, docs: &[Vec<u32>]) -> Vec<Vec<u32>> {
1555 let mut cursor = 0u64;
1556 let mut scratch = Vec::new();
1557 let mut out = Vec::new();
1558 let mut result = Vec::new();
1559 for doc in docs {
1560 assert!(stream.read_into(cursor, doc.len() as u32, &mut scratch, &mut out));
1561 result.push(out.clone());
1562 cursor += doc.len() as u64;
1563 }
1564 result
1565 }
1566
1567 #[test]
1568 fn stream_round_trips_sorted_positions_across_blocks() {
1569 let docs: Vec<Vec<u32>> = (0..50)
1570 .map(|d| (0..(d % 7 + 1) * 13).map(|i| i * 3 + d).collect())
1571 .chain(std::iter::once((0..300).map(|i| i * 1000).collect()))
1572 .chain(std::iter::once(vec![70_000, 5, 5, 1 << 21]))
1573 .collect();
1574 let (buf, total) = encode(&docs);
1575 assert_eq!(total, docs.iter().map(|d| d.len() as u64).sum::<u64>());
1576 assert!(PositionStream::is_stream(&buf));
1577 let stream = PositionStream::open(OwnedBytes::new(buf)).unwrap();
1578 assert!(stream.canonical_blocks);
1579 assert_eq!(stream.total_positions(), total);
1580 assert_eq!(stream.num_blocks(), total.div_ceil(128) as usize);
1581 let expected: Vec<Vec<u32>> = docs
1582 .iter()
1583 .map(|d| {
1584 let mut s = d.clone();
1585 s.sort_unstable();
1586 s
1587 })
1588 .collect();
1589 assert_eq!(read_all(&stream, &docs), expected);
1590 let mut scratch = Vec::new();
1592 let mut out = Vec::new();
1593 assert!(!stream.read_into(total - 1, 2, &mut scratch, &mut out));
1594 }
1595
1596 #[test]
1597 fn raw_repacking_preserves_payload_bytes_without_certifying_document_boundaries() {
1598 let a: Vec<Vec<u32>> = (0..40).map(|d| vec![d, d + 2, d + 7]).collect();
1599 let b: Vec<Vec<u32>> = (0..90).map(|d| (0..d % 5 + 1).collect()).collect();
1600 let (buf_a, _) = encode(&a);
1601 let (buf_b, _) = encode(&b);
1602 let mut merged = Vec::new();
1603 let mut encoder = PositionStreamEncoder::new(&mut merged);
1604 let mut values = Vec::new();
1605 for buf in [buf_a, buf_b] {
1606 let stream = PositionStream::open(OwnedBytes::new(buf)).unwrap();
1607 for idx in 0..stream.num_blocks() {
1608 assert!(stream.decode_block(idx, &mut values));
1609 encoder.push_values(&values).unwrap();
1610 }
1611 }
1612 encoder.finish().unwrap();
1613 let all: Vec<Vec<u32>> = a.iter().chain(&b).cloned().collect();
1614 let (mut direct, _) = encode(&all);
1615 assert!(has_unique_positions(&direct));
1616 assert!(!has_unique_positions(&merged));
1617 let at = direct.len() - FOOTER;
1618 let count = u32::from_le_bytes(direct[at..at + 4].try_into().unwrap()) & !UNIQUE_POSITIONS;
1619 direct[at..at + 4].copy_from_slice(&count.to_le_bytes());
1620 assert_eq!(merged, direct);
1621 }
1622
1623 #[test]
1624 fn streaming_concatenation_copies_non_aligned_blocks_verbatim() {
1625 let a: Vec<Vec<u32>> = (0..43)
1628 .map(|doc| {
1629 (0..doc % 5 + 1)
1630 .map(|position| doc + position * 7)
1631 .collect()
1632 })
1633 .collect();
1634 let b: Vec<Vec<u32>> = (0..51)
1635 .map(|doc| {
1636 (0..doc % 4 + 1)
1637 .map(|position| doc * 2 + position)
1638 .collect()
1639 })
1640 .collect();
1641 let (encoded_a, total_a) = encode(&a);
1642 let (encoded_b, total_b) = encode(&b);
1643 assert_ne!(total_a % POSITION_STREAM_BLOCK as u64, 0);
1644 assert_ne!(total_b % POSITION_STREAM_BLOCK as u64, 0);
1645
1646 let source_blocks = |raw: &[u8]| {
1647 let stream = PositionStream::open(OwnedBytes::new(raw.to_vec())).unwrap();
1648 (0..stream.num_blocks())
1649 .map(|idx| {
1650 let (start, end, _) = stream.block_range(idx).unwrap();
1651 raw[start..end].to_vec()
1652 })
1653 .collect::<Vec<_>>()
1654 };
1655 let expected_blocks: Vec<Vec<u8>> = source_blocks(&encoded_a)
1656 .into_iter()
1657 .chain(source_blocks(&encoded_b))
1658 .collect();
1659
1660 let mut merged = Vec::new();
1661 let (total, written) = PositionStream::concatenate_streaming(
1662 &[encoded_a.as_slice(), encoded_b.as_slice()],
1663 &mut merged,
1664 )
1665 .unwrap();
1666 assert_eq!(total, total_a + total_b);
1667 assert_eq!(written as usize, merged.len());
1668
1669 let stream = PositionStream::open(OwnedBytes::new(merged.clone())).unwrap();
1670 assert!(!stream.canonical_blocks);
1671 let actual_blocks: Vec<Vec<u8>> = (0..stream.num_blocks())
1672 .map(|idx| {
1673 let (start, end, _) = stream.block_range(idx).unwrap();
1674 merged[start..end].to_vec()
1675 })
1676 .collect();
1677 assert_eq!(actual_blocks, expected_blocks, "encoded blocks changed");
1678 assert_eq!(stream.num_blocks(), expected_blocks.len());
1679
1680 let all: Vec<Vec<u32>> = a.iter().chain(&b).cloned().collect();
1681 let expected: Vec<Vec<u32>> = all
1682 .iter()
1683 .map(|positions| {
1684 let mut positions = positions.clone();
1685 positions.sort_unstable();
1686 positions
1687 })
1688 .collect();
1689 assert_eq!(read_all(&stream, &all), expected);
1690 }
1691
1692 #[test]
1693 fn single_source_streaming_concatenation_is_an_exact_copy() {
1694 let (encoded, total) = encode(&[(0..137).collect()]);
1695 let mut copied = Vec::new();
1696 let result =
1697 PositionStream::concatenate_streaming(&[encoded.as_slice()], &mut copied).unwrap();
1698 assert_eq!(result, (total, encoded.len() as u64));
1699 assert_eq!(copied, encoded);
1700 }
1701
1702 #[test]
1706 fn short_bitpacker_block_is_rejected_as_corruption() {
1707 fn assemble(count: usize, width: u8, codec: u8, payload_len: usize) -> Vec<u8> {
1708 let mut raw = Vec::new();
1709 raw.extend_from_slice(&(count as u16).to_le_bytes());
1710 raw.extend_from_slice(&[width, codec]);
1711 raw.extend(std::iter::repeat_n(0u8, payload_len));
1712 let block_len = raw.len();
1713 raw.write_u32::<LittleEndian>(0).unwrap();
1714 raw.write_u64::<LittleEndian>(0).unwrap();
1715 raw.write_u32::<LittleEndian>(1).unwrap();
1716 raw.write_u64::<LittleEndian>(count as u64).unwrap();
1717 raw.write_u32::<LittleEndian>(MAGIC).unwrap();
1718 assert_eq!(raw.len(), block_len + INDEX_ENTRY + FOOTER);
1719 raw
1720 }
1721 for count in [1usize, 5, 64, 127] {
1722 for width in [0u8, 3, 8, 32] {
1723 let short = assemble(count, width, 1, bitpacking4x::encoded_len(count, width));
1724 let block = &short[..short.len() - INDEX_ENTRY - FOOTER];
1725 assert_eq!(
1726 PositionStream::block_count(block),
1727 None,
1728 "count={count} width={width}"
1729 );
1730 assert!(
1731 PositionStream::open(OwnedBytes::new(short)).is_err(),
1732 "count={count} width={width}"
1733 );
1734 }
1735 }
1736 for count in [1usize, 5, 64, 127] {
1739 let rounded = assemble(count, 8, 0, count);
1740 assert!(PositionStream::open(OwnedBytes::new(rounded)).is_ok());
1741 }
1742 let full = assemble(
1743 POSITION_STREAM_BLOCK,
1744 3,
1745 1,
1746 bitpacking4x::encoded_len(POSITION_STREAM_BLOCK, 3),
1747 );
1748 let stream = PositionStream::open(OwnedBytes::new(full)).unwrap();
1749 let mut block = Vec::new();
1750 assert!(stream.decode_block(0, &mut block));
1751 assert_eq!(block, vec![0; POSITION_STREAM_BLOCK]);
1752 let mut bytes = Vec::new();
1754 let mut encoder =
1755 PositionStreamEncoder::with_posting_codec(&mut bytes, PostingCodec::Simd4x);
1756 encoder.push_values(&[1; 130]).unwrap();
1757 encoder.finish().unwrap();
1758 let stream = PositionStream::open(OwnedBytes::new(bytes.clone())).unwrap();
1759 let (start, _, _) = stream.block_range(0).unwrap();
1760 assert_eq!(bytes[start + 3], 1);
1761 let (start, _, _) = stream.block_range(1).unwrap();
1762 assert_eq!(bytes[start + 3], 0);
1763 }
1764
1765 #[test]
1766 fn term_positions_reject_old_stream_revisions() {
1767 let (buf, _) = encode(&[vec![1, 4], vec![0]]);
1768 let positions = TermPositions::open(OwnedBytes::new(buf.clone())).unwrap();
1769 assert_eq!(positions.positions(0, 2), Some(vec![1, 4]));
1770 for magic in [b"POS3", b"POS4", b"POS7"] {
1771 let mut bytes = buf.clone();
1772 let at = bytes.len() - 4;
1773 bytes[at..].copy_from_slice(magic);
1774 assert!(TermPositions::open(OwnedBytes::new(bytes)).is_err());
1775 }
1776 }
1777}