hermes_core/structures/postings/
posting_common.rs1use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
9use std::io::{self, Read, Write};
10
11use crate::DocId;
12
13pub use crate::structures::vint::{read_vint, write_vint};
14
15pub const BLOCK_SIZE: usize = 128;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct SkipEntry {
23 pub first_doc: DocId,
25 pub last_doc: DocId,
27 pub offset: u32,
29}
30
31impl SkipEntry {
32 pub fn new(first_doc: DocId, last_doc: DocId, offset: u32) -> Self {
33 Self {
34 first_doc,
35 last_doc,
36 offset,
37 }
38 }
39
40 pub fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
42 writer.write_u32::<LittleEndian>(self.first_doc)?;
43 writer.write_u32::<LittleEndian>(self.last_doc)?;
44 writer.write_u32::<LittleEndian>(self.offset)?;
45 Ok(())
46 }
47
48 pub fn read<R: Read>(reader: &mut R) -> io::Result<Self> {
50 let first_doc = reader.read_u32::<LittleEndian>()?;
51 let last_doc = reader.read_u32::<LittleEndian>()?;
52 let offset = reader.read_u32::<LittleEndian>()?;
53 Ok(Self {
54 first_doc,
55 last_doc,
56 offset,
57 })
58 }
59}
60
61#[derive(Debug, Clone, Default)]
63pub struct SkipList {
64 entries: Vec<SkipEntry>,
65}
66
67impl SkipList {
68 pub fn new() -> Self {
69 Self::default()
70 }
71
72 pub fn with_capacity(capacity: usize) -> Self {
73 Self {
74 entries: Vec::with_capacity(capacity),
75 }
76 }
77
78 pub fn push(&mut self, first_doc: DocId, last_doc: DocId, offset: u32) {
80 self.entries
81 .push(SkipEntry::new(first_doc, last_doc, offset));
82 }
83
84 pub fn len(&self) -> usize {
86 self.entries.len()
87 }
88
89 pub fn is_empty(&self) -> bool {
90 self.entries.is_empty()
91 }
92
93 pub fn get(&self, index: usize) -> Option<&SkipEntry> {
95 self.entries.get(index)
96 }
97
98 pub fn find_block(&self, target: DocId) -> Option<usize> {
103 let idx = self.entries.partition_point(|e| e.last_doc < target);
104 if idx < self.entries.len() {
105 Some(idx)
106 } else {
107 None
108 }
109 }
110
111 pub fn iter(&self) -> impl Iterator<Item = &SkipEntry> {
113 self.entries.iter()
114 }
115
116 pub fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
118 writer.write_u32::<LittleEndian>(self.entries.len() as u32)?;
119 for entry in &self.entries {
120 entry.write(writer)?;
121 }
122 Ok(())
123 }
124
125 pub fn read<R: Read>(reader: &mut R) -> io::Result<Self> {
127 let count = reader.read_u32::<LittleEndian>()? as usize;
128 let mut entries = Vec::with_capacity(count);
129 for _ in 0..count {
130 entries.push(SkipEntry::read(reader)?);
131 }
132 Ok(Self { entries })
133 }
134
135 pub fn from_tuples(tuples: &[(DocId, DocId, u32)]) -> Self {
137 Self {
138 entries: tuples
139 .iter()
140 .map(|(first, last, offset)| SkipEntry::new(*first, *last, *offset))
141 .collect(),
142 }
143 }
144
145 pub fn to_tuples(&self) -> Vec<(DocId, DocId, u32)> {
147 self.entries
148 .iter()
149 .map(|e| (e.first_doc, e.last_doc, e.offset))
150 .collect()
151 }
152}
153
154pub fn write_doc_id_block<W: Write>(writer: &mut W, doc_ids: &[DocId]) -> io::Result<DocId> {
159 if doc_ids.is_empty() {
160 return Ok(0);
161 }
162
163 write_vint(writer, doc_ids.len() as u64)?;
164
165 let mut prev = 0u32;
166 for (i, &doc_id) in doc_ids.iter().enumerate() {
167 if i == 0 {
168 write_vint(writer, doc_id as u64)?;
170 } else {
171 write_vint(writer, (doc_id - prev) as u64)?;
173 }
174 prev = doc_id;
175 }
176
177 Ok(*doc_ids.last().unwrap())
178}
179
180pub fn read_doc_id_block<R: Read>(reader: &mut R) -> io::Result<Vec<DocId>> {
184 let count = read_vint(reader)? as usize;
185 let mut doc_ids = Vec::with_capacity(count);
186
187 let mut prev = 0u32;
188 for i in 0..count {
189 let value = read_vint(reader)? as u32;
190 let doc_id = if i == 0 {
191 value } else {
193 prev + value };
195 doc_ids.push(doc_id);
196 prev = doc_id;
197 }
198
199 Ok(doc_ids)
200}
201
202use crate::structures::simd;
207
208#[derive(Debug, Clone, Copy, PartialEq, Eq)]
212#[repr(u8)]
213pub enum RoundedBitWidth {
214 Zero = 0,
216 Bits8 = 8,
218 Bits16 = 16,
220 Bits32 = 32,
222}
223
224impl RoundedBitWidth {
225 pub fn from_max_value(max_val: u32) -> Self {
227 if max_val == 0 {
228 RoundedBitWidth::Zero
229 } else if max_val <= 255 {
230 RoundedBitWidth::Bits8
231 } else if max_val <= 65535 {
232 RoundedBitWidth::Bits16
233 } else {
234 RoundedBitWidth::Bits32
235 }
236 }
237
238 pub fn bytes_per_value(&self) -> usize {
240 match self {
241 RoundedBitWidth::Zero => 0,
242 RoundedBitWidth::Bits8 => 1,
243 RoundedBitWidth::Bits16 => 2,
244 RoundedBitWidth::Bits32 => 4,
245 }
246 }
247
248 pub fn from_u8(v: u8) -> Option<Self> {
250 match v {
251 0 => Some(RoundedBitWidth::Zero),
252 8 => Some(RoundedBitWidth::Bits8),
253 16 => Some(RoundedBitWidth::Bits16),
254 32 => Some(RoundedBitWidth::Bits32),
255 _ => None,
256 }
257 }
258}
259
260pub fn pack_deltas_fixed(doc_ids: &[DocId]) -> (RoundedBitWidth, Vec<u8>) {
265 if doc_ids.len() <= 1 {
266 return (RoundedBitWidth::Zero, Vec::new());
267 }
268
269 let mut max_delta = 0u32;
271 let mut deltas = Vec::with_capacity(doc_ids.len() - 1);
272
273 for i in 1..doc_ids.len() {
274 let delta = doc_ids[i] - doc_ids[i - 1] - 1; deltas.push(delta);
276 max_delta = max_delta.max(delta);
277 }
278
279 let bit_width = RoundedBitWidth::from_max_value(max_delta);
280 let bytes_per_val = bit_width.bytes_per_value();
281
282 if bytes_per_val == 0 {
283 return (bit_width, Vec::new());
284 }
285
286 let mut packed = Vec::with_capacity(deltas.len() * bytes_per_val);
287
288 match bit_width {
289 RoundedBitWidth::Zero => {}
290 RoundedBitWidth::Bits8 => {
291 for delta in deltas {
292 packed.push(delta as u8);
293 }
294 }
295 RoundedBitWidth::Bits16 => {
296 for delta in deltas {
297 packed.extend_from_slice(&(delta as u16).to_le_bytes());
298 }
299 }
300 RoundedBitWidth::Bits32 => {
301 for delta in deltas {
302 packed.extend_from_slice(&delta.to_le_bytes());
303 }
304 }
305 }
306
307 (bit_width, packed)
308}
309
310pub fn unpack_deltas_fixed(
314 packed: &[u8],
315 bit_width: RoundedBitWidth,
316 first_doc_id: DocId,
317 count: usize,
318 output: &mut [DocId],
319) {
320 if count == 0 {
321 return;
322 }
323
324 output[0] = first_doc_id;
325
326 if count == 1 {
327 return;
328 }
329
330 match bit_width {
331 RoundedBitWidth::Zero => {
332 for (i, out) in output.iter_mut().enumerate().skip(1).take(count - 1) {
334 *out = first_doc_id + i as u32;
335 }
336 }
337 RoundedBitWidth::Bits8 => {
338 simd::unpack_8bit_delta_decode(packed, output, first_doc_id, count);
339 }
340 RoundedBitWidth::Bits16 => {
341 simd::unpack_16bit_delta_decode(packed, output, first_doc_id, count);
342 }
343 RoundedBitWidth::Bits32 => {
344 let mut carry = first_doc_id;
346 for i in 0..count - 1 {
347 let idx = i * 4;
348 let delta = u32::from_le_bytes([
349 packed[idx],
350 packed[idx + 1],
351 packed[idx + 2],
352 packed[idx + 3],
353 ]);
354 carry = carry.wrapping_add(delta).wrapping_add(1);
355 output[i + 1] = carry;
356 }
357 }
358 }
359}
360
361#[cfg(test)]
362mod tests {
363 use super::*;
364
365 #[test]
366 fn test_vint_roundtrip() {
367 let values = [
368 0u64,
369 1,
370 127,
371 128,
372 255,
373 256,
374 16383,
375 16384,
376 u32::MAX as u64,
377 u64::MAX,
378 ];
379
380 for &value in &values {
381 let mut buf = Vec::new();
382 write_vint(&mut buf, value).unwrap();
383 let read_value = read_vint(&mut buf.as_slice()).unwrap();
384 assert_eq!(value, read_value, "Failed for value {}", value);
385 }
386 }
387
388 #[test]
389 fn test_skip_list_roundtrip() {
390 let mut skip_list = SkipList::new();
391 skip_list.push(0, 127, 0);
392 skip_list.push(128, 255, 100);
393 skip_list.push(256, 500, 200);
394
395 let mut buf = Vec::new();
396 skip_list.write(&mut buf).unwrap();
397
398 let restored = SkipList::read(&mut buf.as_slice()).unwrap();
399 assert_eq!(skip_list.len(), restored.len());
400
401 for (a, b) in skip_list.iter().zip(restored.iter()) {
402 assert_eq!(a, b);
403 }
404 }
405
406 #[test]
407 fn test_skip_list_find_block() {
408 let mut skip_list = SkipList::new();
409 skip_list.push(0, 99, 0);
410 skip_list.push(100, 199, 100);
411 skip_list.push(200, 299, 200);
412
413 assert_eq!(skip_list.find_block(0), Some(0));
414 assert_eq!(skip_list.find_block(50), Some(0));
415 assert_eq!(skip_list.find_block(99), Some(0));
416 assert_eq!(skip_list.find_block(100), Some(1));
417 assert_eq!(skip_list.find_block(150), Some(1));
418 assert_eq!(skip_list.find_block(250), Some(2));
419 assert_eq!(skip_list.find_block(300), None);
420 }
421
422 #[test]
423 fn test_doc_id_block_roundtrip() {
424 let doc_ids: Vec<DocId> = vec![0, 5, 10, 100, 1000, 10000];
425
426 let mut buf = Vec::new();
427 let last = write_doc_id_block(&mut buf, &doc_ids).unwrap();
428 assert_eq!(last, 10000);
429
430 let restored = read_doc_id_block(&mut buf.as_slice()).unwrap();
431 assert_eq!(doc_ids, restored);
432 }
433
434 #[test]
435 fn test_doc_id_block_single() {
436 let doc_ids: Vec<DocId> = vec![42];
437
438 let mut buf = Vec::new();
439 write_doc_id_block(&mut buf, &doc_ids).unwrap();
440
441 let restored = read_doc_id_block(&mut buf.as_slice()).unwrap();
442 assert_eq!(doc_ids, restored);
443 }
444}