1use anyhow::{Context, Result};
2use rand::RngExt;
3use rand::SeedableRng;
4use rand::rngs::StdRng;
5use std::collections::BTreeMap;
6use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};
7use std::path::Path;
8
9const INDEX_MAGIC: u32 = 0x47435449;
10const INDEX_VERSION: u32 = 5; pub const COMPOSITE_KEY_SEPARATOR: &str = "\x1F"; pub fn make_composite_key(parts: &[&str]) -> String {
15 parts.join(COMPOSITE_KEY_SEPARATOR)
16}
17
18fn json_value_to_string(v: Option<&serde_json::Value>) -> String {
19 match v {
20 None => String::new(),
21 Some(serde_json::Value::Null) => String::new(),
22 Some(serde_json::Value::String(s)) => s.clone(),
23 Some(serde_json::Value::Number(n)) => n.to_string(),
24 Some(serde_json::Value::Bool(b)) => b.to_string(),
25 Some(other) => other.to_string(),
26 }
27}
28
29const FLAG_HAS_UNICODE: u64 = 1 << 62;
30const FLAG_HAS_METADATA: u64 = 1 << 63;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
33pub enum KeyType {
34 #[default]
35 String,
36 U64,
37 I64,
38 U32,
39 I32,
40 UnixTimestampSec,
41 UnixTimestampMillis,
42 DatePacked,
43 TimePacked,
44 UUID,
45 ULID,
46}
47
48impl KeyType {
49 pub fn parse(&self, value: &str) -> Option<KeyValue> {
50 match self {
51 KeyType::String => Some(KeyValue::String(value.to_string())),
52 KeyType::U64 => value.parse::<u64>().ok().map(KeyValue::U64),
53 KeyType::I64 => value.parse::<i64>().ok().map(KeyValue::I64),
54 KeyType::U32 => value.parse::<u32>().ok().map(KeyValue::U32),
55 KeyType::I32 => value.parse::<i32>().ok().map(KeyValue::I32),
56 KeyType::UnixTimestampSec => value.parse::<i64>().ok().map(KeyValue::I64),
57 KeyType::UnixTimestampMillis => value.parse::<i64>().ok().map(KeyValue::I64),
58 KeyType::DatePacked => parse_date(value).map(KeyValue::U32),
59 KeyType::TimePacked => parse_time(value).map(KeyValue::U32),
60 KeyType::UUID => parse_uuid(value).map(KeyValue::UUID),
61 KeyType::ULID => parse_ulid(value).map(KeyValue::ULID),
62 }
63 }
64
65 pub fn supports_numeric_vec(&self) -> bool {
66 matches!(
67 self,
68 KeyType::U64
69 | KeyType::I64
70 | KeyType::U32
71 | KeyType::I32
72 | KeyType::DatePacked
73 | KeyType::TimePacked
74 | KeyType::UUID
75 | KeyType::ULID
76 )
77 }
78
79 pub fn id(&self) -> u8 {
80 match self {
81 KeyType::String => 0,
82 KeyType::U64 => 1,
83 KeyType::I64 => 2,
84 KeyType::U32 => 3,
85 KeyType::I32 => 4,
86 KeyType::UnixTimestampSec => 5,
87 KeyType::UnixTimestampMillis => 6,
88 KeyType::DatePacked => 7,
89 KeyType::TimePacked => 8,
90 KeyType::UUID => 9,
91 KeyType::ULID => 10,
92 }
93 }
94
95 pub fn from_id(id: u8) -> Option<Self> {
96 match id {
97 0 => Some(KeyType::String),
98 1 => Some(KeyType::U64),
99 2 => Some(KeyType::I64),
100 3 => Some(KeyType::U32),
101 4 => Some(KeyType::I32),
102 5 => Some(KeyType::UnixTimestampSec),
103 6 => Some(KeyType::UnixTimestampMillis),
104 7 => Some(KeyType::DatePacked),
105 8 => Some(KeyType::TimePacked),
106 9 => Some(KeyType::UUID),
107 10 => Some(KeyType::ULID),
108 _ => None,
109 }
110 }
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
114pub enum KeyValue {
115 String(String),
116 U64(u64),
117 I64(i64),
118 U32(u32),
119 I32(i32),
120 UUID(UuidParts),
121 ULID(UlidParts),
122}
123
124#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
125pub struct UuidParts(pub u64, pub u64);
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
128pub struct UlidParts(pub u64, pub u64);
129
130fn parse_uuid(s: &str) -> Option<UuidParts> {
131 let s = s.trim();
132 let bytes = parse_hex_bytes(s)?;
133 if bytes.len() != 16 {
134 return None;
135 }
136 let p0 = u64::from_le_bytes(bytes[0..8].try_into().ok()?);
137 let p1 = u64::from_le_bytes(bytes[8..16].try_into().ok()?);
138 Some(UuidParts(p0, p1))
139}
140
141fn parse_ulid(s: &str) -> Option<UlidParts> {
142 const BASE32_CHARS: &[u8; 32] = b"0123456789ABCDEFGHJKMNPQRSTVWXYZ";
143 let s = s.trim();
144 if s.len() != 26 {
145 return None;
146 }
147 let mut bytes = [0u8; 16];
148 for (i, c) in s.bytes().enumerate() {
149 let c = c.to_ascii_uppercase();
150 let idx = BASE32_CHARS.iter().position(|&x| x == c)?;
151 bytes[i / 2] = bytes[i / 2] * 32 + idx as u8;
152 }
153 let p0 = u64::from_le_bytes(bytes[0..8].try_into().ok()?);
154 let p1 = u64::from_le_bytes(bytes[8..16].try_into().ok()?);
155 Some(UlidParts(p0, p1))
156}
157
158fn parse_hex_bytes(s: &str) -> Option<Vec<u8>> {
159 let s = s.replace(['-', ':'], "");
160 if s.len() != 32 || !s.bytes().all(|b| b.is_ascii_hexdigit()) {
161 return None;
162 }
163 (0..16)
164 .map(|i| {
165 let idx = i * 2;
166 u8::from_str_radix(&s[idx..idx + 2], 16).ok()
167 })
168 .collect()
169}
170
171fn parse_date(s: &str) -> Option<u32> {
172 let parts: Vec<&str> = s.split(&['-', '/'][..]).collect();
173 if parts.len() != 3 {
174 return None;
175 }
176 let year: u32 = parts[0].parse().ok()?;
177 let month: u32 = parts[1].parse().ok()?;
178 let day: u32 = parts[2].parse().ok()?;
179 if !(1900..=2100).contains(&year) || !(1..=12).contains(&month) || !(1..=31).contains(&day) {
180 return None;
181 }
182 Some(year * 10000 + month * 100 + day)
183}
184
185fn parse_time(s: &str) -> Option<u32> {
186 let parts: Vec<&str> = s.split(':').collect();
187 if parts.len() < 2 {
188 return None;
189 }
190 let hour: u32 = parts[0].parse().ok()?;
191 let min: u32 = parts[1].parse().ok()?;
192 let sec = if parts.len() > 2 {
193 parts[2].parse().ok()?
194 } else {
195 0
196 };
197 if hour > 23 || min > 59 || sec > 59 {
198 return None;
199 }
200 Some(hour * 10000 + min * 100 + sec)
201}
202
203pub fn infer_key_type(samples: &[&str]) -> KeyType {
204 if samples.is_empty() {
205 return KeyType::String;
206 }
207 if samples.iter().all(|&s| is_uuid(s)) {
208 return KeyType::UUID;
209 }
210 if samples.iter().all(|&s| is_ulid(s)) {
211 return KeyType::ULID;
212 }
213 if samples.iter().all(|&s| is_date(s)) {
214 return KeyType::DatePacked;
215 }
216 if samples.iter().all(|&s| is_time(s)) {
217 return KeyType::TimePacked;
218 }
219 if samples.iter().all(|&s| s.parse::<u64>().is_ok()) {
220 return KeyType::U64;
221 }
222 if samples.iter().all(|&s| s.parse::<i64>().is_ok()) {
223 return KeyType::I64;
224 }
225 KeyType::String
226}
227
228pub struct InferenceStats {
229 pub samples_taken: usize,
230 pub bytes_scanned: u64,
231 pub all_matched: bool,
232 pub confidence: f32,
233}
234
235pub fn infer_key_type_from_stream<R: Read + Seek + BufRead>(
236 reader: &mut R,
237 key_column_idx: usize,
238 max_samples: usize,
239 max_bytes_scan: u64,
240) -> Result<(KeyType, InferenceStats)> {
241 use std::io::SeekFrom;
242
243 let file_size = reader.seek(SeekFrom::End(0))?;
244 if file_size == 0 {
245 return Ok((
246 KeyType::String,
247 InferenceStats {
248 samples_taken: 0,
249 bytes_scanned: 0,
250 all_matched: false,
251 confidence: 0.0,
252 },
253 ));
254 }
255
256 let mut samples = Vec::with_capacity(max_samples.min(1000));
257 let mut bytes_scanned = 0u64;
258 let mut rng = StdRng::from_rng(&mut rand::rng());
259
260 let num_samples = max_samples.min(1000);
261 let sample_positions: Vec<u64> = (0..num_samples)
262 .map(|_| rng.random_range(0..file_size))
263 .collect();
264
265 for pos in sample_positions {
266 if bytes_scanned >= max_bytes_scan {
267 break;
268 }
269
270 reader.seek(SeekFrom::Start(pos))?;
271 if pos > 0 {
272 let mut dummy = String::new();
273 reader.read_line(&mut dummy).ok();
274 }
275
276 let mut this_line = String::new();
277 match reader.read_line(&mut this_line) {
278 Ok(0) | Err(_) => break,
279 Ok(_) => {}
280 }
281
282 bytes_scanned += this_line.len() as u64;
283
284 let trimmed = this_line.trim();
285 if trimmed.is_empty() || trimmed.starts_with('#') {
286 continue;
287 }
288
289 let parts: Vec<&str> = if trimmed.contains('\t') {
290 trimmed.split('\t').collect()
291 } else {
292 trimmed.split(',').collect()
293 };
294
295 if key_column_idx >= parts.len() {
296 continue;
297 }
298
299 samples.push(parts[key_column_idx].to_string());
300 }
301
302 if samples.is_empty() {
303 return Ok((
304 KeyType::String,
305 InferenceStats {
306 samples_taken: 0,
307 bytes_scanned,
308 all_matched: false,
309 confidence: 0.0,
310 },
311 ));
312 }
313
314 let sample_refs: Vec<&str> = samples.iter().map(|s| s.as_str()).collect();
315 let inferred = infer_key_type(&sample_refs);
316 let all_matched = match inferred {
317 KeyType::String => true,
318 _ => samples.iter().all(|s| inferred.parse(s).is_some()),
319 };
320
321 let confidence = if samples.len() >= max_samples.min(1000) {
322 1.0
323 } else {
324 (samples.len() as f32 / max_samples.min(1000) as f32).min(1.0)
325 };
326
327 Ok((
328 inferred,
329 InferenceStats {
330 samples_taken: samples.len(),
331 bytes_scanned,
332 all_matched,
333 confidence,
334 },
335 ))
336}
337
338pub fn infer_key_type_from_ndjson_stream<R: Read + Seek + BufRead>(
339 reader: &mut R,
340 key_column: &str,
341 max_samples: usize,
342 max_bytes_scan: u64,
343) -> Result<(KeyType, InferenceStats)> {
344 use std::io::SeekFrom;
345
346 let file_size = reader.seek(SeekFrom::End(0))?;
347 if file_size == 0 {
348 return Ok((
349 KeyType::String,
350 InferenceStats {
351 samples_taken: 0,
352 bytes_scanned: 0,
353 all_matched: false,
354 confidence: 0.0,
355 },
356 ));
357 }
358
359 let mut samples = Vec::with_capacity(max_samples.min(1000));
360 let mut bytes_scanned = 0u64;
361 let mut rng = StdRng::from_rng(&mut rand::rng());
362
363 let num_samples = max_samples.min(1000);
364 let sample_positions: Vec<u64> = (0..num_samples)
365 .map(|_| rng.random_range(0..file_size))
366 .collect();
367
368 for pos in sample_positions {
369 if bytes_scanned >= max_bytes_scan {
370 break;
371 }
372
373 reader.seek(SeekFrom::Start(pos))?;
374 if pos > 0 {
375 let mut dummy = String::new();
376 reader.read_line(&mut dummy).ok();
377 }
378
379 let mut this_line = String::new();
380 match reader.read_line(&mut this_line) {
381 Ok(0) | Err(_) => break,
382 Ok(_) => {}
383 }
384
385 bytes_scanned += this_line.len() as u64;
386
387 let trimmed = this_line.trim();
388 if trimmed.is_empty() || trimmed.starts_with('#') {
389 continue;
390 }
391
392 let obj: serde_json::Map<String, serde_json::Value> = match serde_json::from_str(trimmed) {
393 Ok(o) => o,
394 Err(_) => continue,
395 };
396
397 let value = match obj.get(key_column) {
398 Some(v) => json_value_to_string(Some(v)),
399 None => continue,
400 };
401
402 samples.push(value);
403 }
404
405 if samples.is_empty() {
406 return Ok((
407 KeyType::String,
408 InferenceStats {
409 samples_taken: 0,
410 bytes_scanned,
411 all_matched: false,
412 confidence: 0.0,
413 },
414 ));
415 }
416
417 let sample_refs: Vec<&str> = samples.iter().map(|s| s.as_str()).collect();
418 let inferred = infer_key_type(&sample_refs);
419 let all_matched = match inferred {
420 KeyType::String => true,
421 _ => samples.iter().all(|s| inferred.parse(s).is_some()),
422 };
423
424 let confidence = if samples.len() >= max_samples.min(1000) {
425 1.0
426 } else {
427 (samples.len() as f32 / max_samples.min(1000) as f32).min(1.0)
428 };
429
430 Ok((
431 inferred,
432 InferenceStats {
433 samples_taken: samples.len(),
434 bytes_scanned,
435 all_matched,
436 confidence,
437 },
438 ))
439}
440
441fn is_uuid(s: &str) -> bool {
442 let s = s.trim();
443 if s.len() != 36 {
444 return false;
445 }
446 let expected = "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx";
447 let mut si = s.bytes();
448 for c in expected.bytes() {
449 if c == b'x' {
450 if !si.next().is_some_and(|b| b.is_ascii_hexdigit()) {
451 return false;
452 }
453 } else if si.next() != Some(c) {
454 return false;
455 }
456 }
457 true
458}
459
460fn is_ulid(s: &str) -> bool {
461 const VALID_ULID_CHARS: &[u8; 32] = b"0123456789ABCDEFGHJKMNPQRSTVWXYZ";
462 let s = s.trim();
463 if s.len() != 26 {
464 return false;
465 }
466 s.bytes()
467 .all(|b| VALID_ULID_CHARS.contains(&b.to_ascii_uppercase()))
468}
469
470fn is_date(s: &str) -> bool {
471 parse_date(s).is_some()
472}
473
474fn is_time(s: &str) -> bool {
475 parse_time(s).is_some()
476}
477
478#[derive(Debug, Clone)]
479pub struct BloomFilter {
480 bits: Vec<u64>,
481 hash_count: u32,
482 bit_count: usize,
483}
484
485impl BloomFilter {
486 pub fn new(expected_elements: usize, false_positive_rate: f64) -> Self {
487 let m = Self::optimal_bit_count(expected_elements, false_positive_rate);
488 let k = Self::optimal_hash_count(m, expected_elements);
489 let bits = vec![0u64; m.div_ceil(64)];
490 Self {
491 bits,
492 hash_count: k as u32,
493 bit_count: m,
494 }
495 }
496
497 pub fn with_capacity(bit_count: usize, hash_count: u32) -> Self {
498 let bits = vec![0u64; bit_count.div_ceil(64)];
499 Self {
500 bits,
501 hash_count,
502 bit_count,
503 }
504 }
505
506 fn optimal_bit_count(n: usize, p: f64) -> usize {
507 let n = n.max(1) as f64;
508 let p = p.clamp(0.0001, 0.9999);
509 let m = -n * p.ln() / (std::f64::consts::LN_2 * std::f64::consts::LN_2);
510 m.ceil() as usize
511 }
512
513 fn optimal_hash_count(m: usize, n: usize) -> usize {
514 let m = m.max(1) as f64;
515 let n = n.max(1) as f64;
516 let k = (m / n * std::f64::consts::LN_2).ceil();
517 k.max(1.0) as usize
518 }
519
520 pub fn insert(&mut self, key: &str) {
521 for i in 0..self.hash_count {
522 let idx = self.hash(key, i);
523 self.bits[idx / 64] |= 1 << (idx % 64);
524 }
525 }
526
527 #[must_use]
528 pub fn contains(&self, key: &str) -> bool {
529 for i in 0..self.hash_count {
530 let idx = self.hash(key, i);
531 if self.bits[idx / 64] & (1 << (idx % 64)) == 0 {
532 return false;
533 }
534 }
535 true
536 }
537
538 fn hash(&self, key: &str, salt: u32) -> usize {
539 let h1 = Self::fnv1a(key, 0);
540 let h2 = Self::fnv1a(key, 0xdeadbeef);
541 let combined = (h1 as u64).wrapping_add((h2 as u64).wrapping_mul(salt as u64));
542 (combined as usize) % self.bit_count
543 }
544
545 fn fnv1a(data: &str, salt: u32) -> u32 {
546 let mut hash: u32 = 2166136261u32.wrapping_add(salt);
547 for byte in data.bytes() {
548 hash ^= byte as u32;
549 hash = hash.wrapping_mul(16777619);
550 }
551 hash
552 }
553
554 pub fn bit_count(&self) -> usize {
555 self.bit_count
556 }
557
558 pub fn hash_count(&self) -> u32 {
559 self.hash_count
560 }
561
562 pub fn memory_bits(&self) -> usize {
563 self.bits.len() * 64
564 }
565
566 pub fn write_to(&self, writer: &mut impl Write) -> Result<()> {
567 writer.write_all(&self.bit_count.to_le_bytes())?;
568 writer.write_all(&self.hash_count.to_le_bytes())?;
569 for chunk in self.bits.chunks(8192) {
570 let bytes = chunk
571 .iter()
572 .fold(Vec::with_capacity(chunk.len() * 8), |mut acc, &w| {
573 acc.extend_from_slice(&w.to_le_bytes());
574 acc
575 });
576 writer.write_all(&bytes)?;
577 }
578 Ok(())
579 }
580
581 pub fn read_from(reader: &mut impl Read) -> Result<Self> {
582 let mut bit_buf = [0u8; 8];
583 reader.read_exact(&mut bit_buf)?;
584 let bit_count = usize::from_le_bytes(bit_buf);
585
586 let mut hash_buf = [0u8; 4];
587 reader.read_exact(&mut hash_buf)?;
588 let hash_count = u32::from_le_bytes(hash_buf);
589
590 let bits = vec![0u64; bit_count.div_ceil(64)];
591 let mut result = Self {
592 bits,
593 hash_count,
594 bit_count,
595 };
596
597 let byte_count = bit_count.div_ceil(8);
598 let mut buf = vec![0u8; byte_count];
599 reader.read_exact(&mut buf)?;
600 for (i, chunk) in buf.chunks(8).enumerate() {
601 let mut word = [0u8; 8];
602 word[..chunk.len()].copy_from_slice(chunk);
603 result.bits[i] = u64::from_le_bytes(word);
604 }
605 Ok(result)
606 }
607}
608
609#[derive(Debug, Clone)]
610pub struct XorFilter {
611 fingerprint_bits: u32,
612 array_size: usize,
613 seed: u64,
614 fingerprints: Vec<u8>,
615}
616
617impl XorFilter {
618 pub fn new(expected_elements: usize, false_positive_rate: f64) -> Self {
619 let fpr = false_positive_rate.clamp(0.0001, 0.9999);
620 let fingerprint_bits = Self::optimal_fingerprint_bits(fpr);
621 let array_size = Self::optimal_array_size(expected_elements);
622 let fingerprints = vec![0u8; array_size];
623 Self {
624 fingerprint_bits,
625 array_size,
626 seed: 0x9e3779b97f4a7c15,
627 fingerprints,
628 }
629 }
630
631 fn optimal_fingerprint_bits(p: f64) -> u32 {
632 let p = p.clamp(0.0001, 0.9999);
633 ((-p.log2()).ceil() as u32).clamp(4, 16)
634 }
635
636 fn optimal_array_size(n: usize) -> usize {
637 let n = n.max(1);
638 let c = 1.23;
639 let size = ((n as f64) * c).ceil() as usize;
640 size.next_power_of_two()
641 }
642
643 fn murmurhash64(data: &[u8], seed: u64) -> u64 {
644 let c1: u64 = 0x9e3779b97f4a7c15;
645 let c2: u64 = 0x9e3779b97f4a7c15;
646 let mut h: u64 = seed;
647 let len = data.len();
648
649 let mut i = 0;
650 while i + 8 <= len {
651 let mut k = u64::from_le_bytes([
652 data[i],
653 data[i + 1],
654 data[i + 2],
655 data[i + 3],
656 data[i + 4],
657 data[i + 5],
658 data[i + 6],
659 data[i + 7],
660 ]);
661 k = k.wrapping_mul(c1);
662 k = k.rotate_left(31);
663 k = k.wrapping_mul(c2);
664 h ^= k;
665 h = h.rotate_left(27);
666 h = h.wrapping_add(0x9e3779b97f4a7c15);
667 h = h.wrapping_mul(c1);
668 i += 8;
669 }
670
671 let mut k: u64 = 0;
672 match len % 8 {
673 7 => k ^= (data[i + 6] as u64) << 48,
674 6 => k ^= (data[i + 5] as u64) << 40,
675 5 => k ^= (data[i + 4] as u64) << 32,
676 4 => k ^= (data[i + 3] as u64) << 24,
677 3 => k ^= (data[i + 2] as u64) << 16,
678 2 => k ^= (data[i + 1] as u64) << 8,
679 1 => k ^= data[i] as u64,
680 _ => {}
681 }
682 if !len.is_multiple_of(8) {
683 k ^= (len as u64).wrapping_mul(c2);
684 k = k.wrapping_mul(c1);
685 h ^= k;
686 h = h.rotate_left(31);
687 h = h.wrapping_mul(c2);
688 } else {
689 h ^= (len as u64).wrapping_mul(c1);
690 h ^= h.rotate_left(31);
691 h ^= h.rotate_right(33);
692 }
693
694 h = h.wrapping_add(h << 15);
695 h ^= h.rotate_right(41);
696 h = h.wrapping_add(h << 13);
697 h ^= h.rotate_right(35);
698 h = h.wrapping_add(h << 9);
699 h ^= h.rotate_right(49);
700 h = h.wrapping_add(h << 15);
701 h ^= h.rotate_right(33);
702 h = h.wrapping_add(h << 17);
703 h ^= h.rotate_right(41);
704 h
705 }
706
707 fn compute_positions(&self, key: &str) -> [usize; 3] {
708 let h = Self::murmurhash64(key.as_bytes(), self.seed);
709 let h2 = Self::murmurhash64(key.as_bytes(), self.seed ^ 0x9e3779b97f4a7c15);
710
711 let mask = self.array_size - 1;
712 let h0 = (h as usize) & mask;
713 let h1 = ((h >> 32) as usize) & mask;
714 let h2 = (h2 as usize) & mask;
715
716 [h0, h1, h2]
717 }
718
719 fn compute_fingerprint(&self, key: &str) -> u8 {
720 let h = Self::murmurhash64(key.as_bytes(), self.seed ^ 0xdeadbeef);
721 let fb = self.fingerprint_bits;
722 ((h >> 32) ^ h) as u8 & ((1u8 << fb) - 1)
723 }
724
725 pub fn insert(&mut self, key: &str) -> bool {
726 let [h0, h1, h2] = self.compute_positions(key);
727 let f = self.compute_fingerprint(key);
728
729 self.fingerprints[h0] ^= f;
730 self.fingerprints[h1] ^= f;
731 self.fingerprints[h2] ^= f;
732
733 true
734 }
735
736 #[must_use]
737 pub fn contains(&self, key: &str) -> bool {
738 let [h0, h1, h2] = self.compute_positions(key);
739 let f = self.compute_fingerprint(key);
740
741 let expected = self.fingerprints[h0] ^ self.fingerprints[h1] ^ self.fingerprints[h2];
742 expected == f
743 }
744
745 pub fn array_size(&self) -> usize {
746 self.array_size
747 }
748
749 pub fn fingerprint_bits(&self) -> u32 {
750 self.fingerprint_bits
751 }
752
753 pub fn memory_bits(&self) -> usize {
754 self.fingerprints.len() * 8
755 }
756
757 pub fn build_from_keys(keys: &[String]) -> Option<Self> {
758 let n = keys.len();
759 if n == 0 {
760 return Some(Self::new(1, 0.01));
761 }
762
763 let mut filter = Self::new(n, 0.01);
764 let mut retry_count = 0;
765 const MAX_RETRIES: usize = 100;
766
767 for key in keys {
768 if !filter.insert(key) {
769 retry_count += 1;
770 if retry_count >= MAX_RETRIES {
771 return None;
772 }
773 filter = Self::new(n, 0.01);
774 for k in keys {
775 if !filter.insert(k) {
776 retry_count += 1;
777 if retry_count >= MAX_RETRIES {
778 return None;
779 }
780 }
781 }
782 }
783 }
784
785 Some(filter)
786 }
787}
788
789#[derive(Debug, Clone)]
790pub struct IndexEntry {
791 pub offset: u64,
792 pub row_length: u32,
793}
794
795#[derive(Debug, Clone)]
796pub struct IndexEntryV4 {
797 pub offset: u64,
798 pub row_length: u32,
799 pub has_unicode_suffix: bool,
800 pub has_extended_metadata: bool,
801}
802
803impl IndexEntryV4 {
804 pub fn new(offset: u64, row_length: u32) -> Self {
805 Self {
806 offset,
807 row_length,
808 has_unicode_suffix: false,
809 has_extended_metadata: false,
810 }
811 }
812
813 pub fn with_unicode(mut self) -> Self {
814 self.has_unicode_suffix = true;
815 self
816 }
817
818 pub fn with_metadata(mut self) -> Self {
819 self.has_extended_metadata = true;
820 self
821 }
822
823 pub fn encode(&self) -> u64 {
824 let mut bits = self.offset & 0x3FFFFFFFFFFFFFFF;
825 if self.has_unicode_suffix {
826 bits |= FLAG_HAS_UNICODE;
827 }
828 if self.has_extended_metadata {
829 bits |= FLAG_HAS_METADATA;
830 }
831 bits
832 }
833
834 pub fn decode(bits: u64) -> Self {
835 let offset = bits & 0x3FFFFFFFFFFFFFFF;
836 let has_unicode_suffix = (bits & FLAG_HAS_UNICODE) != 0;
837 let has_extended_metadata = (bits & FLAG_HAS_METADATA) != 0;
838 Self {
839 offset,
840 row_length: 0,
841 has_unicode_suffix,
842 has_extended_metadata,
843 }
844 }
845
846 pub fn with_row_length(mut self, row_length: u32) -> Self {
847 self.row_length = row_length;
848 self
849 }
850}
851
852#[derive(Debug)]
853pub struct IndexHeader {
854 pub version: u32,
855 pub key_column: String,
856 pub key_type: KeyType,
857 pub data_offset: u64,
858 pub entry_count: u32,
859}
860
861enum KeyStorage {
862 String(BTreeMap<String, Vec<IndexEntry>>),
863 Numeric(Vec<(KeyValue, String, Vec<IndexEntry>)>),
864}
865
866pub struct SourceIndex {
867 storage: KeyStorage,
868 header: IndexHeader,
869 filter: Option<XorFilter>,
870}
871
872impl SourceIndex {
873 pub fn new(key_column: &str) -> Self {
874 Self {
875 storage: KeyStorage::String(BTreeMap::new()),
876 header: IndexHeader {
877 version: INDEX_VERSION,
878 key_column: key_column.to_string(),
879 key_type: KeyType::String,
880 data_offset: 0,
881 entry_count: 0,
882 },
883 filter: None,
884 }
885 }
886
887 pub fn with_key_type(key_column: &str, key_type: KeyType) -> Self {
888 let storage = if key_type.supports_numeric_vec() {
889 KeyStorage::Numeric(Vec::new())
890 } else {
891 KeyStorage::String(BTreeMap::new())
892 };
893 Self {
894 storage,
895 header: IndexHeader {
896 version: INDEX_VERSION,
897 key_column: key_column.to_string(),
898 key_type,
899 data_offset: 0,
900 entry_count: 0,
901 },
902 filter: None,
903 }
904 }
905
906 pub fn with_filter(mut self, filter: XorFilter) -> Self {
907 self.filter = Some(filter);
908 self
909 }
910
911 pub fn key_type(&self) -> KeyType {
912 self.header.key_type
913 }
914
915 pub fn insert(&mut self, key: String, offset: u64, row_length: u32) -> Result<()> {
916 let entry = IndexEntry { offset, row_length };
917 match &mut self.storage {
918 KeyStorage::String(map) => {
919 map.entry(key).or_default().push(entry);
920 }
921 KeyStorage::Numeric(vec) => {
922 let kv = match self.header.key_type.parse(&key) {
923 Some(kv) => kv,
924 None => anyhow::bail!(
925 "type mismatch at offset {}: key '{}' cannot be parsed as {:?}. \
926 Consider running `grpctestify index --force` to rebuild with correct type inference.",
927 offset,
928 key,
929 self.header.key_type
930 ),
931 };
932 match vec.binary_search_by(|e| e.0.cmp(&kv)) {
933 Ok(pos) => vec[pos].2.push(entry),
934 Err(pos) => vec.insert(pos, (kv, key.clone(), vec![entry])),
935 }
936 }
937 }
938 Ok(())
939 }
940
941 pub fn batch_insert(
942 &mut self,
943 entries: impl IntoIterator<Item = (String, u64, u32)>,
944 ) -> Result<()> {
945 match &mut self.storage {
946 KeyStorage::String(map) => {
947 for (key, offset, row_length) in entries {
948 map.entry(key)
949 .or_default()
950 .push(IndexEntry { offset, row_length });
951 }
952 }
953 KeyStorage::Numeric(vec) => {
954 let key_type = self.header.key_type;
955 let mut batch: Vec<(KeyValue, String, IndexEntry)> = entries
956 .into_iter()
957 .filter_map(|(key, offset, row_length)| {
958 let kv = key_type.parse(&key)?;
959 Some((kv, key, IndexEntry { offset, row_length }))
960 })
961 .collect();
962
963 batch.sort_unstable_by(|a, b| a.0.cmp(&b.0));
964
965 let mut prev_kv: Option<KeyValue> = None;
966 let mut idx = 0usize;
967 for (kv, key, entry) in batch {
968 if let Some(ref pk) = prev_kv
969 && pk == &kv
970 && let Some(ref mut last) = vec.last_mut()
971 && last.0 == kv
972 {
973 last.2.push(entry);
974 continue;
975 }
976 if idx < vec.len() && vec[idx].0 <= kv {
977 idx = vec[idx..]
978 .binary_search_by(|e| e.0.cmp(&kv))
979 .unwrap_or_else(|e| idx + e);
980 }
981 if idx < vec.len() && vec[idx].0 == kv {
982 vec[idx].2.push(entry);
983 } else {
984 vec.insert(idx, (kv.clone(), key.clone(), vec![entry]));
985 }
986 prev_kv = Some(kv);
987 }
988 }
989 }
990 Ok(())
991 }
992
993 fn binary_search(&self, key: &str) -> Option<&IndexEntry> {
994 match &self.storage {
995 KeyStorage::String(map) => map.get(key).and_then(|v| v.first()),
996 KeyStorage::Numeric(vec) => {
997 let key_type = self.header.key_type;
998 let kv = key_type.parse(key)?;
999 vec.binary_search_by(|e| e.0.cmp(&kv))
1000 .ok()
1001 .map(|pos| &vec[pos].2[0])
1002 }
1003 }
1004 }
1005
1006 pub fn lookup(&self, key: &str) -> Option<&IndexEntry> {
1007 if let Some(filter) = &self.filter
1008 && !filter.contains(key)
1009 {
1010 return None;
1011 }
1012 self.binary_search(key)
1013 }
1014
1015 pub fn lookup_all(&self, key: &str) -> Option<&[IndexEntry]> {
1016 if let Some(filter) = &self.filter
1017 && !filter.contains(key)
1018 {
1019 return None;
1020 }
1021 match &self.storage {
1022 KeyStorage::String(map) => map.get(key).map(|v| v.as_slice()),
1023 KeyStorage::Numeric(vec) => {
1024 let key_type = self.header.key_type;
1025 let kv = key_type.parse(key)?;
1026 vec.binary_search_by(|e| e.0.cmp(&kv))
1027 .ok()
1028 .map(|pos| vec[pos].2.as_slice())
1029 }
1030 }
1031 }
1032
1033 #[must_use]
1034 pub fn contains(&self, key: &str) -> bool {
1035 if let Some(filter) = &self.filter
1036 && !filter.contains(key)
1037 {
1038 return false;
1039 }
1040 self.lookup_all(key).is_some()
1041 }
1042
1043 pub fn fast_negative_lookup(&self, key: &str) -> bool {
1044 if let Some(filter) = &self.filter {
1045 return filter.contains(key);
1046 }
1047 true
1048 }
1049
1050 pub fn lookup_range(&self, start: &str, end: &str) -> Vec<&IndexEntry> {
1051 let mut results = Vec::new();
1052 match &self.storage {
1053 KeyStorage::String(map) => {
1054 let start_s = start.to_string();
1055 let end_s = end.to_string();
1056 for (_key, entries) in map.range(start_s..=end_s) {
1057 for entry in entries {
1058 results.push(entry);
1059 }
1060 }
1061 }
1062 KeyStorage::Numeric(vec) => {
1063 let key_type = self.header.key_type;
1064 let Some(start_kv) = key_type.parse(start) else {
1065 return results;
1066 };
1067 let Some(end_kv) = key_type.parse(end) else {
1068 return results;
1069 };
1070 let start_idx = match vec.binary_search_by(|e| e.0.cmp(&start_kv)) {
1071 Ok(idx) => idx,
1072 Err(idx) => idx,
1073 };
1074 for (kv, _, entries) in &vec[start_idx..] {
1075 if kv > &end_kv {
1076 break;
1077 }
1078 for entry in entries {
1079 results.push(entry);
1080 }
1081 }
1082 }
1083 }
1084 results
1085 }
1086
1087 #[must_use]
1088 pub fn len(&self) -> usize {
1089 match &self.storage {
1090 KeyStorage::String(map) => map.values().map(Vec::len).sum(),
1091 KeyStorage::Numeric(vec) => vec.iter().map(|(_, _, entries)| entries.len()).sum(),
1092 }
1093 }
1094
1095 pub fn unique_keys_len(&self) -> usize {
1096 match &self.storage {
1097 KeyStorage::String(map) => map.len(),
1098 KeyStorage::Numeric(vec) => vec.len(),
1099 }
1100 }
1101
1102 #[must_use]
1103 pub fn is_empty(&self) -> bool {
1104 match &self.storage {
1105 KeyStorage::String(map) => map.is_empty(),
1106 KeyStorage::Numeric(vec) => vec.is_empty(),
1107 }
1108 }
1109
1110 pub fn entry_count(&self) -> u32 {
1111 self.header.entry_count
1112 }
1113
1114 pub fn index_version(&self) -> u32 {
1115 self.header.version
1116 }
1117
1118 pub fn iter(&self) -> impl Iterator<Item = (&str, &IndexEntry)> {
1119 match &self.storage {
1120 KeyStorage::String(map) => Box::new(
1121 map.iter()
1122 .flat_map(|(k, vs)| vs.iter().map(move |v| (k.as_str(), v))),
1123 )
1124 as Box<dyn Iterator<Item = (&str, &IndexEntry)>>,
1125 KeyStorage::Numeric(vec) => Box::new(vec.iter().flat_map(|(_, key_str, vs)| {
1126 let k = key_str.as_str();
1127 vs.iter().map(move |v| (k, v))
1128 }))
1129 as Box<dyn Iterator<Item = (&str, &IndexEntry)>>,
1130 }
1131 }
1132
1133 pub fn write_to_file(&mut self, path: &Path) -> Result<()> {
1134 let count = self.len() as u32;
1135 self.header.entry_count = count;
1136
1137 let key_col_bytes = self.header.key_column.as_bytes();
1138 let key_col_len = key_col_bytes.len() as u32;
1139 let key_type_id = self.header.key_type.id();
1140
1141 let header_size = 4 + 4 + 1 + 4 + key_col_len + 8 + 4;
1142 let data_offset = header_size as u64;
1143 self.header.data_offset = data_offset;
1144
1145 let mut buf = Vec::with_capacity(1024 * 1024);
1146 buf.write_all(&INDEX_MAGIC.to_le_bytes())?;
1147 buf.write_all(&INDEX_VERSION.to_le_bytes())?;
1148 buf.write_all(&key_type_id.to_le_bytes())?;
1149 buf.write_all(&key_col_len.to_le_bytes())?;
1150 buf.write_all(key_col_bytes)?;
1151 buf.write_all(&data_offset.to_le_bytes())?;
1152 buf.write_all(&count.to_le_bytes())?;
1153
1154 let mut prev_key = String::new();
1155 match &self.storage {
1156 KeyStorage::String(map) => {
1157 for (key, entries) in map {
1158 let (prefix_len, suffix) = shared_prefix_suffix(&prev_key, key);
1159 write_var_u64(&mut buf, prefix_len as u64)?;
1160 write_var_u64(&mut buf, suffix.len() as u64)?;
1161 buf.write_all(suffix.as_bytes())?;
1162 write_var_u64(&mut buf, entries.len() as u64)?;
1163
1164 let mut prev_offset = 0u64;
1165 for (i, entry) in entries.iter().enumerate() {
1166 if i == 0 {
1167 write_var_u64(&mut buf, entry.offset)?;
1168 } else {
1169 write_var_u64(&mut buf, entry.offset.saturating_sub(prev_offset))?;
1170 }
1171 write_var_u64(&mut buf, entry.row_length as u64)?;
1172 prev_offset = entry.offset;
1173 }
1174 prev_key = key.clone();
1175 }
1176 }
1177 KeyStorage::Numeric(vec) => {
1178 for (_, key, entries) in vec {
1179 let (prefix_len, suffix) = shared_prefix_suffix(&prev_key, key);
1180 write_var_u64(&mut buf, prefix_len as u64)?;
1181 write_var_u64(&mut buf, suffix.len() as u64)?;
1182 buf.write_all(suffix.as_bytes())?;
1183 write_var_u64(&mut buf, entries.len() as u64)?;
1184
1185 let mut prev_offset = 0u64;
1186 for (i, entry) in entries.iter().enumerate() {
1187 if i == 0 {
1188 write_var_u64(&mut buf, entry.offset)?;
1189 } else {
1190 write_var_u64(&mut buf, entry.offset.saturating_sub(prev_offset))?;
1191 }
1192 write_var_u64(&mut buf, entry.row_length as u64)?;
1193 prev_offset = entry.offset;
1194 }
1195 prev_key = key.clone();
1196 }
1197 }
1198 }
1199
1200 let checksum = crc32fast::hash(&buf);
1201 buf.write_all(&checksum.to_le_bytes())?;
1202
1203 std::fs::write(path, &buf)?;
1204 Ok(())
1205 }
1206
1207 pub fn read_from_file(path: &Path) -> Result<Self> {
1208 let file = std::fs::File::open(path)
1209 .with_context(|| format!("failed to open index file: {}", path.display()))?;
1210 let file_len = file.metadata().map(|m| m.len()).unwrap_or(0);
1211 use std::io::Read;
1212
1213 let mut all_data = Vec::with_capacity(file_len as usize);
1214 std::io::BufReader::new(file).read_to_end(&mut all_data)?;
1215
1216 if all_data.len() < 4 {
1217 anyhow::bail!("index file too short: {} bytes", all_data.len());
1218 }
1219
1220 let (payload, stored_crc) = all_data.split_at(all_data.len() - 4);
1221 let expected_crc = u32::from_le_bytes(stored_crc[..4].try_into().unwrap());
1222 let actual_crc = crc32fast::hash(payload);
1223 if actual_crc != expected_crc {
1224 anyhow::bail!(
1225 "CRC32 checksum mismatch in index file: expected 0x{:08X}, got 0x{:08X}",
1226 expected_crc,
1227 actual_crc
1228 );
1229 }
1230
1231 let mut cursor = std::io::Cursor::new(payload);
1232
1233 let magic = read_u32(&mut cursor)?;
1234 if magic != INDEX_MAGIC {
1235 anyhow::bail!(
1236 "invalid index file magic: expected 0x{:08X}, got 0x{:08X}",
1237 INDEX_MAGIC,
1238 magic
1239 );
1240 }
1241
1242 let version = read_u32(&mut cursor)?;
1243 if version != INDEX_VERSION {
1244 anyhow::bail!("unsupported index version: {version}");
1245 }
1246
1247 let key_type_id = read_u8(&mut cursor)?;
1248 let key_type = KeyType::from_id(key_type_id).context("invalid key type in index file")?;
1249
1250 let key_col_len = read_u32(&mut cursor)? as usize;
1251 let mut key_col_buf = vec![0u8; key_col_len];
1252 cursor.read_exact(&mut key_col_buf)?;
1253 let key_column =
1254 String::from_utf8(key_col_buf.clone()).context("invalid UTF-8 in index key_column")?;
1255
1256 let data_offset = read_u64(&mut cursor)?;
1257 let entry_count = read_u32(&mut cursor)?;
1258
1259 let storage = if key_type.supports_numeric_vec() {
1260 KeyStorage::Numeric(Vec::new())
1261 } else {
1262 KeyStorage::String(BTreeMap::new())
1263 };
1264
1265 let mut storage = storage;
1266 let mut read_entries = 0u32;
1267 let mut prev_key = String::new();
1268 while read_entries < entry_count {
1269 let prefix_len = read_var_u64(&mut cursor)? as usize;
1270 let suffix_len = read_var_u64(&mut cursor)? as usize;
1271 let mut suffix_buf = vec![0u8; suffix_len];
1272 cursor.read_exact(&mut suffix_buf)?;
1273 let suffix =
1274 String::from_utf8(suffix_buf).context("invalid UTF-8 in index key suffix")?;
1275 let key = rebuild_key(&prev_key, prefix_len, &suffix)?;
1276 let posting_count = read_var_u64(&mut cursor)? as u32;
1277 let mut postings = Vec::with_capacity(posting_count as usize);
1278 let mut prev_offset = 0u64;
1279 for _ in 0..posting_count {
1280 let raw = read_var_u64(&mut cursor)?;
1281 let offset = if postings.is_empty() {
1282 raw
1283 } else {
1284 prev_offset.saturating_add(raw)
1285 };
1286 let row_length = read_var_u64(&mut cursor)? as u32;
1287 postings.push(IndexEntry { offset, row_length });
1288 read_entries += 1;
1289 prev_offset = offset;
1290 }
1291 prev_key = key.clone();
1292
1293 match &mut storage {
1294 KeyStorage::String(map) => {
1295 map.insert(key, postings);
1296 }
1297 KeyStorage::Numeric(vec) => {
1298 if let Some(kv) = key_type.parse(&key) {
1299 vec.push((kv, key, postings));
1300 }
1301 }
1302 }
1303 }
1304
1305 Ok(Self {
1306 storage,
1307 header: IndexHeader {
1308 version,
1309 key_column,
1310 key_type,
1311 data_offset,
1312 entry_count,
1313 },
1314 filter: None,
1315 })
1316 }
1317
1318 pub fn lookup_row<R: Read + Seek>(&self, reader: &mut R, key: &str) -> Result<Option<String>> {
1319 let entries = self.lookup_all(key);
1320 let entry = match entries.and_then(|e| e.first()) {
1321 Some(e) => e,
1322 None => return Ok(None),
1323 };
1324
1325 reader.seek(SeekFrom::Start(entry.offset))?;
1326 let mut buf = vec![0u8; entry.row_length as usize];
1327 reader.read_exact(&mut buf)?;
1328 let line = String::from_utf8(buf).context("invalid UTF-8 in source row")?;
1329 Ok(Some(line))
1330 }
1331
1332 pub fn lookup_row_from_mmap(&self, mmap_data: &[u8], key: &str) -> Result<Option<String>> {
1333 let entries = self.lookup_all(key);
1334 let entry = match entries.and_then(|e| e.first()) {
1335 Some(e) => e,
1336 None => return Ok(None),
1337 };
1338
1339 let start = entry.offset as usize;
1340 let end = start + entry.row_length as usize;
1341
1342 if end > mmap_data.len() {
1343 anyhow::bail!(
1344 "index entry out of bounds: offset={} len={} mmap_len={}",
1345 entry.offset,
1346 entry.row_length,
1347 mmap_data.len()
1348 );
1349 }
1350
1351 let line = String::from_utf8(mmap_data[start..end].to_vec())
1352 .context("invalid UTF-8 in source row")?;
1353 Ok(Some(line))
1354 }
1355
1356 pub fn key_column(&self) -> &str {
1357 &self.header.key_column
1358 }
1359}
1360
1361pub fn read_index_key_type<R: std::io::Read + std::io::Seek>(reader: &mut R) -> Result<KeyType> {
1362 use std::io::SeekFrom;
1363 reader.seek(SeekFrom::Start(0))?;
1364 let magic = read_u32(reader)?;
1365 if magic != INDEX_MAGIC {
1366 anyhow::bail!("not a valid index file");
1367 }
1368 let version = read_u32(reader)?;
1369 if version != INDEX_VERSION {
1370 anyhow::bail!("unsupported index version: {}", version);
1371 }
1372 let key_type_id = read_u8(reader)?;
1373 KeyType::from_id(key_type_id).context("invalid key type in index file")
1374}
1375
1376pub fn is_index_valid(path: &Path) -> bool {
1377 let file = match std::fs::File::open(path) {
1378 Ok(f) => f,
1379 Err(_) => return false,
1380 };
1381 let mut reader = BufReader::new(file);
1382 if read_u32(&mut reader).ok() != Some(INDEX_MAGIC) {
1383 return false;
1384 }
1385 if read_u32(&mut reader).ok() != Some(INDEX_VERSION) {
1386 return false;
1387 }
1388 if read_u8(&mut reader).is_err() {
1389 return false;
1390 }
1391 let key_col_len = match read_u32(&mut reader) {
1392 Ok(l) => l as usize,
1393 Err(_) => return false,
1394 };
1395 let mut key_col_buf = vec![0u8; key_col_len];
1396 if reader.read_exact(&mut key_col_buf).is_err() {
1397 return false;
1398 }
1399 if read_u64(&mut reader).is_err() {
1400 return false;
1401 }
1402 if read_u32(&mut reader).is_err() {
1403 return false;
1404 }
1405
1406 let file_len = match reader.seek(std::io::SeekFrom::End(0)) {
1407 Ok(l) => l,
1408 Err(_) => return false,
1409 };
1410 let header_end = 4 + 4 + 1 + 4 + key_col_len + 8 + 4;
1411 if file_len < (header_end + 4) as u64 {
1412 return false;
1413 }
1414
1415 let checksum_pos = file_len - 4;
1416 if reader.seek(std::io::SeekFrom::Start(checksum_pos)).is_err() {
1417 return false;
1418 }
1419 let stored_checksum = match read_u32(&mut reader) {
1420 Ok(c) => c,
1421 Err(_) => return false,
1422 };
1423
1424 if reader.seek(std::io::SeekFrom::Start(0)).is_err() {
1425 return false;
1426 }
1427 let data_to_hash_len = checksum_pos as usize;
1428 let mut data_to_hash = vec![0u8; data_to_hash_len];
1429 if reader.read_exact(&mut data_to_hash).is_err() {
1430 return false;
1431 }
1432 let computed_checksum = crc32fast::hash(&data_to_hash);
1433 computed_checksum == stored_checksum
1434}
1435
1436fn read_u32(reader: &mut impl Read) -> Result<u32> {
1437 let mut buf = [0u8; 4];
1438 reader.read_exact(&mut buf)?;
1439 Ok(u32::from_le_bytes(buf))
1440}
1441
1442fn read_u64(reader: &mut impl Read) -> Result<u64> {
1443 let mut buf = [0u8; 8];
1444 reader.read_exact(&mut buf)?;
1445 Ok(u64::from_le_bytes(buf))
1446}
1447
1448fn read_u8(reader: &mut impl Read) -> Result<u8> {
1449 let mut buf = [0u8; 1];
1450 reader.read_exact(&mut buf)?;
1451 Ok(buf[0])
1452}
1453
1454fn shared_prefix_suffix<'a>(prev: &str, current: &'a str) -> (usize, &'a str) {
1455 let max = prev.len().min(current.len());
1456 let mut i = 0usize;
1457 let prev_b = prev.as_bytes();
1458 let cur_b = current.as_bytes();
1459 while i < max && prev_b[i] == cur_b[i] {
1460 i += 1;
1461 }
1462 while i > 0 && !current.is_char_boundary(i) {
1463 i -= 1;
1464 }
1465 (i, ¤t[i..])
1466}
1467
1468fn rebuild_key(prev: &str, prefix_len: usize, suffix: &str) -> Result<String> {
1469 if prefix_len > prev.len() || !prev.is_char_boundary(prefix_len) {
1470 anyhow::bail!("invalid key prefix length in index stream");
1471 }
1472 let mut out = String::with_capacity(prefix_len + suffix.len());
1473 out.push_str(&prev[..prefix_len]);
1474 out.push_str(suffix);
1475 Ok(out)
1476}
1477
1478fn write_var_u64(writer: &mut impl Write, mut value: u64) -> Result<()> {
1479 while value >= 0x80 {
1480 writer.write_all(&[((value as u8) & 0x7F) | 0x80])?;
1481 value >>= 7;
1482 }
1483 writer.write_all(&[value as u8])?;
1484 Ok(())
1485}
1486
1487fn read_var_u64(reader: &mut impl Read) -> Result<u64> {
1488 let mut shift = 0u32;
1489 let mut out = 0u64;
1490 loop {
1491 if shift > 63 {
1492 anyhow::bail!("varint too long in index stream");
1493 }
1494 let mut b = [0u8; 1];
1495 reader.read_exact(&mut b)?;
1496 let byte = b[0];
1497 out |= ((byte & 0x7F) as u64) << shift;
1498 if (byte & 0x80) == 0 {
1499 break;
1500 }
1501 shift += 7;
1502 }
1503 Ok(out)
1504}
1505
1506#[cfg(test)]
1507mod tests {
1508 use super::*;
1509 use std::io::Cursor;
1510
1511 #[cfg(not(miri))]
1512 #[test]
1513 fn write_and_read_roundtrip() {
1514 let dir = std::env::temp_dir().join("gctf_index_test");
1515 std::fs::create_dir_all(&dir).unwrap();
1516 let path = dir.join("test.gcti");
1517
1518 let mut idx = SourceIndex::new("pvz_id");
1519 idx.insert("pvz_001".into(), 0, 50).unwrap();
1520 idx.insert("pvz_002".into(), 51, 60).unwrap();
1521 idx.insert("pvz_003".into(), 112, 45).unwrap();
1522 idx.write_to_file(&path).unwrap();
1523
1524 let loaded = SourceIndex::read_from_file(&path).unwrap();
1525 assert_eq!(loaded.len(), 3);
1526 assert_eq!(loaded.key_column(), "pvz_id");
1527
1528 let e1 = loaded.lookup("pvz_001").unwrap();
1529 assert_eq!(e1.offset, 0);
1530 assert_eq!(e1.row_length, 50);
1531
1532 let e2 = loaded.lookup("pvz_002").unwrap();
1533 assert_eq!(e2.offset, 51);
1534
1535 assert!(loaded.lookup("missing").is_none());
1536
1537 std::fs::remove_file(&path).ok();
1538 }
1539
1540 #[cfg(not(miri))]
1541 #[test]
1542 fn invalid_magic_fails() {
1543 let dir = std::env::temp_dir().join("gctf_index_test");
1544 std::fs::create_dir_all(&dir).unwrap();
1545 let path = dir.join("bad_magic.gcti");
1546
1547 let mut f = std::fs::File::create(&path).unwrap();
1548 f.write_all(&0xDEADBEEFu32.to_le_bytes()).unwrap();
1549
1550 let result = SourceIndex::read_from_file(&path);
1551 assert!(result.is_err());
1552
1553 std::fs::remove_file(&path).ok();
1554 }
1555
1556 #[test]
1557 fn entries_sorted_by_key() {
1558 let mut idx = SourceIndex::new("id");
1559 idx.insert("c".into(), 200, 10).unwrap();
1560 idx.insert("a".into(), 0, 10).unwrap();
1561 idx.insert("b".into(), 100, 10).unwrap();
1562
1563 let keys: Vec<&str> = idx.iter().map(|(k, _)| k).collect();
1564 assert_eq!(keys, vec!["a", "b", "c"]);
1565 }
1566
1567 #[test]
1568 fn lookup_row_from_source() {
1569 let source_data = "id,name,age\n1,Alice,30\n2,Bob,25\n";
1570
1571 let mut idx = SourceIndex::new("id");
1572 let header_line = "id,name,age\n";
1573 let row1_offset = header_line.len() as u64;
1574 let row1 = "1,Alice,30";
1575 idx.insert("1".into(), row1_offset, row1.len() as u32)
1576 .unwrap();
1577
1578 let row2_offset = (header_line.len() + row1.len() + 1) as u64;
1579 let row2 = "2,Bob,25";
1580 idx.insert("2".into(), row2_offset, row2.len() as u32)
1581 .unwrap();
1582
1583 let mut cursor = Cursor::new(source_data);
1584 let line1 = idx.lookup_row(&mut cursor, "1").unwrap().unwrap();
1585 assert_eq!(line1, "1,Alice,30");
1586
1587 let line2 = idx.lookup_row(&mut cursor, "2").unwrap().unwrap();
1588 assert_eq!(line2, "2,Bob,25");
1589
1590 assert!(idx.lookup_row(&mut cursor, "99").unwrap().is_none());
1591 }
1592
1593 #[cfg(not(miri))]
1594 #[test]
1595 fn empty_index_roundtrip() {
1596 let dir = std::env::temp_dir().join("gctf_index_test");
1597 std::fs::create_dir_all(&dir).unwrap();
1598 let path = dir.join("empty.gcti");
1599
1600 let mut idx = SourceIndex::new("id");
1601 idx.write_to_file(&path).unwrap();
1602
1603 let loaded = SourceIndex::read_from_file(&path).unwrap();
1604 assert!(loaded.is_empty());
1605 assert_eq!(loaded.len(), 0);
1606
1607 std::fs::remove_file(&path).ok();
1608 }
1609
1610 #[cfg(not(miri))]
1611 #[test]
1612 fn unicode_keys() {
1613 let dir = std::env::temp_dir().join("gctf_index_test");
1614 std::fs::create_dir_all(&dir).unwrap();
1615 let path = dir.join("unicode.gcti");
1616
1617 let mut idx = SourceIndex::new("город");
1618 idx.insert("Москва".into(), 0, 10).unwrap();
1619 idx.insert("Санкт-Петербург".into(), 10, 20).unwrap();
1620 idx.write_to_file(&path).unwrap();
1621
1622 let loaded = SourceIndex::read_from_file(&path).unwrap();
1623 assert_eq!(loaded.key_column(), "город");
1624 assert!(loaded.contains("Москва"));
1625 assert!(loaded.contains("Санкт-Петербург"));
1626
1627 std::fs::remove_file(&path).ok();
1628 }
1629
1630 #[cfg(not(miri))]
1631 #[test]
1632 fn duplicate_keys_are_preserved() {
1633 let dir = std::env::temp_dir().join("gctf_index_dup_test");
1634 std::fs::create_dir_all(&dir).unwrap();
1635 let path = dir.join("dup.gcti");
1636
1637 let mut idx = SourceIndex::new("zone_id");
1638 idx.insert("z1".into(), 10, 20).unwrap();
1639 idx.insert("z1".into(), 31, 22).unwrap();
1640 idx.insert("z2".into(), 54, 18).unwrap();
1641 idx.write_to_file(&path).unwrap();
1642
1643 let loaded = SourceIndex::read_from_file(&path).unwrap();
1644 assert_eq!(loaded.len(), 3);
1645 assert_eq!(loaded.unique_keys_len(), 2);
1646
1647 let all = loaded.lookup_all("z1").unwrap();
1648 assert_eq!(all.len(), 2);
1649 assert_eq!(all[0].offset, 10);
1650 assert_eq!(all[1].offset, 31);
1651
1652 std::fs::remove_file(&path).ok();
1653 }
1654
1655 #[test]
1656 fn index_entry_v4_encode_decode() {
1657 let entry = IndexEntryV4::new(0x123456789ABC, 100);
1658 let encoded = entry.encode();
1659 let decoded = IndexEntryV4::decode(encoded);
1660 assert_eq!(decoded.offset, 0x123456789ABC);
1661 assert!(!decoded.has_unicode_suffix);
1662 assert!(!decoded.has_extended_metadata);
1663 }
1664
1665 #[test]
1666 fn index_entry_v4_with_unicode_flag() {
1667 let entry = IndexEntryV4::new(0x1000, 50).with_unicode();
1668 let encoded = entry.encode();
1669 assert!(encoded & FLAG_HAS_UNICODE != 0);
1670
1671 let decoded = IndexEntryV4::decode(encoded);
1672 assert!(decoded.has_unicode_suffix);
1673 }
1674
1675 #[test]
1676 fn index_entry_v4_with_metadata_flag() {
1677 let entry = IndexEntryV4::new(0x1000, 50).with_metadata();
1678 let encoded = entry.encode();
1679 assert!(encoded & FLAG_HAS_METADATA != 0);
1680
1681 let decoded = IndexEntryV4::decode(encoded);
1682 assert!(decoded.has_extended_metadata);
1683 }
1684
1685 #[test]
1686 fn index_entry_v4_both_flags() {
1687 let entry = IndexEntryV4::new(0x1000, 50).with_unicode().with_metadata();
1688 let encoded = entry.encode();
1689 assert!(encoded & FLAG_HAS_UNICODE != 0);
1690 assert!(encoded & FLAG_HAS_METADATA != 0);
1691
1692 let decoded = IndexEntryV4::decode(encoded);
1693 assert!(decoded.has_unicode_suffix);
1694 assert!(decoded.has_extended_metadata);
1695 }
1696
1697 #[test]
1698 fn index_entry_v4_max_offset() {
1699 let max_offset: u64 = 0x3FFFFFFFFFFFFFFF;
1700 let entry = IndexEntryV4::new(max_offset, 1000);
1701 let encoded = entry.encode();
1702 let decoded = IndexEntryV4::decode(encoded);
1703 assert_eq!(decoded.offset, max_offset);
1704 }
1705
1706 #[test]
1707 fn lookup_range_string_keys() {
1708 let mut idx = SourceIndex::new("zone_id");
1709 idx.insert("zone_a".into(), 0, 10).unwrap();
1710 idx.insert("zone_b".into(), 20, 15).unwrap();
1711 idx.insert("zone_c".into(), 50, 20).unwrap();
1712 idx.insert("zone_d".into(), 100, 25).unwrap();
1713
1714 let results = idx.lookup_range("zone_a", "zone_c");
1715 assert_eq!(results.len(), 3);
1716 assert_eq!(results[0].offset, 0);
1717 assert_eq!(results[1].offset, 20);
1718 assert_eq!(results[2].offset, 50);
1719 }
1720
1721 #[test]
1722 fn lookup_range_numeric_keys() {
1723 let mut idx = SourceIndex::with_key_type("date_id", KeyType::DatePacked);
1724 idx.insert("2024-01-01".into(), 0, 10).unwrap();
1725 idx.insert("2024-01-15".into(), 20, 15).unwrap();
1726 idx.insert("2024-01-31".into(), 50, 20).unwrap();
1727 idx.insert("2024-02-01".into(), 100, 25).unwrap();
1728
1729 let results = idx.lookup_range("2024-01-01", "2024-01-31");
1730 assert_eq!(results.len(), 3);
1731
1732 let results2 = idx.lookup_range("2024-01-10", "2024-01-20");
1733 assert_eq!(results2.len(), 1);
1734 assert_eq!(results2[0].offset, 20);
1735 }
1736
1737 #[test]
1738 fn lookup_range_single_key() {
1739 let mut idx = SourceIndex::new("id");
1740 idx.insert("a".into(), 0, 10).unwrap();
1741 idx.insert("b".into(), 20, 15).unwrap();
1742 idx.insert("c".into(), 50, 20).unwrap();
1743
1744 let results = idx.lookup_range("b", "b");
1745 assert_eq!(results.len(), 1);
1746 assert_eq!(results[0].offset, 20);
1747 }
1748
1749 #[test]
1750 fn lookup_range_no_match() {
1751 let mut idx = SourceIndex::new("id");
1752 idx.insert("a".into(), 0, 10).unwrap();
1753 idx.insert("c".into(), 50, 20).unwrap();
1754
1755 let results = idx.lookup_range("b", "b");
1756 assert!(results.is_empty());
1757 }
1758}