Skip to main content

apif_source_row/
index.rs

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; // v5: single-file + typed keys + CRC32 checksum
11pub const COMPOSITE_KEY_SEPARATOR: &str = "\x1F"; // Unit separator for composite keys
12
13/// Build a composite key from multiple column values, joined by COMPOSITE_KEY_SEPARATOR.
14pub 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, &current[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}