Skip to main content

formualizer_eval/engine/
lookup_index_cache.rs

1use std::hash::{Hash, Hasher};
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::{Arc, RwLock};
4
5use formualizer_common::{ExcelError, LiteralValue, SheetId};
6use rustc_hash::FxHashMap;
7use smallvec::SmallVec;
8
9use crate::builtins::lookup::lookup_utils::cmp_for_lookup;
10use crate::engine::{DateSystem, range_view::RangeView};
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
13pub struct LookupIndexKey {
14    pub(crate) sheet_id: SheetId,
15    pub(crate) start_row: u32,
16    pub(crate) start_col: u32,
17    pub(crate) end_row: u32,
18    pub(crate) end_col: u32,
19    pub(crate) axis: LookupAxis,
20    pub(crate) snapshot_id: u64,
21}
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
24pub enum LookupAxis {
25    ColumnInView(usize),
26    RowInView(usize),
27}
28
29#[derive(Debug, Eq, PartialEq)]
30pub enum LookupHashKey {
31    Number(u64),
32    Text(Box<str>),
33    Boolean(bool),
34    Empty,
35}
36
37impl Hash for LookupHashKey {
38    fn hash<H: Hasher>(&self, state: &mut H) {
39        match self {
40            Self::Number(bits) => {
41                0u8.hash(state);
42                bits.hash(state);
43            }
44            Self::Text(text) => {
45                1u8.hash(state);
46                text.hash(state);
47            }
48            Self::Boolean(value) => {
49                2u8.hash(state);
50                value.hash(state);
51            }
52            Self::Empty => {
53                3u8.hash(state);
54            }
55        }
56    }
57}
58
59impl LookupHashKey {
60    fn from_needle(value: &LiteralValue, date_system: DateSystem) -> Option<Self> {
61        // Blank needles retain numeric-zero coercion; blank candidates do not.
62        if matches!(value, LiteralValue::Empty) {
63            Some(Self::Number(0.0f64.to_bits()))
64        } else {
65            Self::from_literal(value, date_system)
66        }
67    }
68
69    pub(crate) fn from_literal(value: &LiteralValue, date_system: DateSystem) -> Option<Self> {
70        match value {
71            LiteralValue::Number(n) => Some(Self::Number(normalize_f64_bits(*n))),
72            LiteralValue::Int(i) => Some(Self::Number(normalize_f64_bits(*i as f64))),
73            LiteralValue::Text(s) => Some(Self::Text(s.to_lowercase().into_boxed_str())),
74            LiteralValue::Boolean(b) => Some(Self::Boolean(*b)),
75            LiteralValue::Empty => None,
76            // Temporal values are numbers in Excel: key them by their serial so
77            // an exact lookup finds them whether the needle or the cell (or
78            // both) carry a temporal type rather than a plain numeric.
79            LiteralValue::Date(_) | LiteralValue::DateTime(_) | LiteralValue::Time(_) => value
80                .as_serial_number_for(date_system)
81                .map(|serial| Self::Number(normalize_f64_bits(serial))),
82            LiteralValue::Error(_)
83            | LiteralValue::Array(_)
84            | LiteralValue::Duration(_)
85            | LiteralValue::Pending => None,
86        }
87    }
88}
89
90fn normalize_f64_bits(n: f64) -> u64 {
91    if n.is_nan() {
92        return f64::NAN.to_bits();
93    }
94    let rounded = n.round();
95    if (n - rounded).abs() < 1e-12 {
96        // Exact comparisons identify signed zero; the index must do so too.
97        if rounded == 0.0 {
98            0.0f64.to_bits()
99        } else {
100            rounded.to_bits()
101        }
102    } else {
103        n.to_bits()
104    }
105}
106
107#[derive(Debug, Clone, Default)]
108pub struct DuplicateIndices {
109    pub(crate) first: usize,
110    pub(crate) last: usize,
111    pub(crate) all: SmallVec<[usize; 1]>,
112}
113
114pub struct LookupIndex {
115    pub(crate) len: usize,
116    date_system: DateSystem,
117    pub(crate) bytes: usize,
118    pub(crate) entries: FxHashMap<LookupHashKey, DuplicateIndices>,
119    pub(crate) cell_values: Box<[LiteralValue]>,
120}
121
122#[cfg(test)]
123thread_local! {
124    static BUILD_ATTEMPTS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
125}
126
127#[cfg(test)]
128pub(crate) fn take_build_attempts() -> usize {
129    BUILD_ATTEMPTS.with(|c| c.replace(0))
130}
131
132impl LookupIndex {
133    pub(crate) fn build(
134        view: &RangeView<'_>,
135        axis: LookupAxis,
136        date_system: DateSystem,
137    ) -> Result<BuildOutcome, ExcelError> {
138        #[cfg(test)]
139        BUILD_ATTEMPTS.with(|c| c.set(c.get() + 1));
140        let (rows, cols) = view.dims();
141        let len = match axis {
142            LookupAxis::ColumnInView(col) => {
143                if col >= cols {
144                    return Ok(BuildOutcome::Degenerate);
145                }
146                rows
147            }
148            LookupAxis::RowInView(row) => {
149                if row >= rows {
150                    return Ok(BuildOutcome::Degenerate);
151                }
152                cols
153            }
154        };
155        if len == 0 {
156            return Ok(BuildOutcome::Degenerate);
157        }
158
159        let mut entries: FxHashMap<LookupHashKey, DuplicateIndices> = FxHashMap::default();
160        let mut cell_values = Vec::with_capacity(len);
161        let mut error_count = 0usize;
162
163        for idx in 0..len {
164            let value = match axis {
165                LookupAxis::ColumnInView(col) => view.get_cell(idx, col),
166                LookupAxis::RowInView(row) => view.get_cell(row, idx),
167            };
168            if matches!(value, LiteralValue::Error(_)) {
169                error_count += 1;
170            }
171            if let Some(key) = LookupHashKey::from_literal(&value, date_system) {
172                let dups = entries.entry(key).or_insert_with(|| DuplicateIndices {
173                    first: idx,
174                    last: idx,
175                    all: SmallVec::new(),
176                });
177                if dups.all.is_empty() {
178                    dups.first = idx;
179                }
180                dups.last = idx;
181                dups.all.push(idx);
182            }
183            cell_values.push(value);
184        }
185
186        if error_count > 0 {
187            return Ok(BuildOutcome::ErrorInLookupAxis);
188        }
189
190        let bytes = retained_bytes(&cell_values, &entries);
191        Ok(BuildOutcome::Built(Self {
192            len,
193            date_system,
194            bytes,
195            entries,
196            cell_values: cell_values.into_boxed_slice(),
197        }))
198    }
199
200    pub(crate) fn find_first_exact(&self, needle: &LiteralValue) -> Option<usize> {
201        let hash_key = LookupHashKey::from_needle(needle, self.date_system)?;
202        if let Some(dups) = self.entries.get(&hash_key) {
203            for &idx in &dups.all {
204                if cmp_for_lookup(needle, &self.cell_values[idx], self.date_system) == Some(0) {
205                    return Some(idx);
206                }
207            }
208        }
209        None
210    }
211
212    pub(crate) fn find_last_exact(&self, needle: &LiteralValue) -> Option<usize> {
213        let hash_key = LookupHashKey::from_needle(needle, self.date_system)?;
214        if let Some(dups) = self.entries.get(&hash_key) {
215            for &idx in dups.all.iter().rev() {
216                if cmp_for_lookup(needle, &self.cell_values[idx], self.date_system) == Some(0) {
217                    return Some(idx);
218                }
219            }
220        }
221        None
222    }
223}
224
225fn retained_bytes(
226    values: &[LiteralValue],
227    entries: &FxHashMap<LookupHashKey, DuplicateIndices>,
228) -> usize {
229    // HashMap capacity excludes control bytes and vacant buckets. Round up to
230    // the backing power-of-two bucket count; allocator bookkeeping is estimated.
231    let buckets = if entries.capacity() == 0 {
232        0
233    } else {
234        entries.capacity().saturating_add(1).next_power_of_two()
235    };
236    let mut bytes = values
237        .len()
238        .saturating_mul(std::mem::size_of::<LiteralValue>())
239        .saturating_add(
240            buckets.saturating_mul(std::mem::size_of::<(LookupHashKey, DuplicateIndices)>() + 1),
241        )
242        .saturating_add(256);
243    for value in values {
244        bytes = bytes.saturating_add(literal_payload_bytes(value));
245    }
246    for (key, indices) in entries {
247        if let LookupHashKey::Text(text) = key {
248            bytes = bytes.saturating_add(text.len());
249        }
250        if indices.all.spilled() {
251            bytes = bytes.saturating_add(
252                indices
253                    .all
254                    .capacity()
255                    .saturating_mul(std::mem::size_of::<usize>()),
256            );
257        }
258    }
259    bytes
260}
261
262fn literal_payload_bytes(value: &LiteralValue) -> usize {
263    match value {
264        LiteralValue::Text(text) => text.capacity(),
265        LiteralValue::Array(rows) => rows.iter().fold(
266            rows.capacity()
267                .saturating_mul(std::mem::size_of::<Vec<LiteralValue>>()),
268            |bytes, row| {
269                row.iter().fold(
270                    bytes.saturating_add(
271                        row.capacity()
272                            .saturating_mul(std::mem::size_of::<LiteralValue>()),
273                    ),
274                    |bytes, value| bytes.saturating_add(literal_payload_bytes(value)),
275                )
276            },
277        ),
278        _ => 0,
279    }
280}
281
282pub(crate) fn estimate_bytes(len: usize, entries: usize) -> usize {
283    len.saturating_mul(std::mem::size_of::<LiteralValue>().saturating_add(8))
284        .saturating_add(entries.saturating_mul(96))
285        .saturating_add(256)
286}
287
288pub(crate) enum BuildOutcome {
289    Built(LookupIndex),
290    ErrorInLookupAxis,
291    Degenerate,
292}
293
294const LOOKUP_INDEX_BUILD_THRESHOLD: u32 = 3;
295const CAP_REJECTED: u32 = u32::MAX;
296const CALL_COUNT_PRUNE_LIMIT: usize = 4096;
297
298#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
299pub struct LookupIndexCacheReport {
300    pub(crate) builds: usize,
301    pub(crate) hits: usize,
302    pub(crate) misses: usize,
303    pub(crate) skipped_volatile: usize,
304    pub(crate) skipped_error: usize,
305    pub(crate) skipped_tiny: usize,
306    pub(crate) skipped_cap: usize,
307    pub(crate) skipped_below_threshold: usize,
308    pub(crate) bytes_in_cache: usize,
309    pub(crate) entries_count: usize,
310}
311
312pub struct LookupIndexCache {
313    inner: RwLock<FxHashMap<LookupIndexKey, Arc<LookupIndex>>>,
314    call_counts: RwLock<FxHashMap<LookupIndexKey, u32>>,
315    volatile_keys: RwLock<FxHashMap<LookupIndexKey, ()>>,
316    build_threshold: u32,
317    bytes_in_use: AtomicUsize,
318    max_bytes: usize,
319    builds: AtomicUsize,
320    hits: AtomicUsize,
321    misses: AtomicUsize,
322    skipped_volatile: AtomicUsize,
323    skipped_error: AtomicUsize,
324    skipped_tiny: AtomicUsize,
325    skipped_cap: AtomicUsize,
326    skipped_below_threshold: AtomicUsize,
327}
328
329fn volatile_key(mut key: LookupIndexKey) -> LookupIndexKey {
330    key.snapshot_id = 0;
331    key
332}
333
334impl LookupIndexCache {
335    pub(crate) fn new(max_bytes: usize) -> Self {
336        Self {
337            inner: RwLock::new(FxHashMap::default()),
338            call_counts: RwLock::new(FxHashMap::default()),
339            volatile_keys: RwLock::new(FxHashMap::default()),
340            build_threshold: LOOKUP_INDEX_BUILD_THRESHOLD,
341            bytes_in_use: AtomicUsize::new(0),
342            max_bytes,
343            builds: AtomicUsize::new(0),
344            hits: AtomicUsize::new(0),
345            misses: AtomicUsize::new(0),
346            skipped_volatile: AtomicUsize::new(0),
347            skipped_error: AtomicUsize::new(0),
348            skipped_tiny: AtomicUsize::new(0),
349            skipped_cap: AtomicUsize::new(0),
350            skipped_below_threshold: AtomicUsize::new(0),
351        }
352    }
353
354    // Called only at an exclusive Engine mutation boundary, after all evaluation
355    // workers have joined. No old-generation builder can race this reclamation.
356    pub(crate) fn clear(&mut self) {
357        self.inner
358            .get_mut()
359            .unwrap_or_else(|p| p.into_inner())
360            .clear();
361        self.call_counts
362            .get_mut()
363            .unwrap_or_else(|p| p.into_inner())
364            .clear();
365        self.volatile_keys
366            .get_mut()
367            .unwrap_or_else(|p| p.into_inner())
368            .clear();
369        self.bytes_in_use.store(0, Ordering::Relaxed);
370    }
371
372    pub(crate) fn get(&self, key: &LookupIndexKey) -> Option<Arc<LookupIndex>> {
373        let found = self
374            .inner
375            .read()
376            .ok()
377            .and_then(|guard| guard.get(key).cloned());
378        if found.is_some() {
379            self.hits.fetch_add(1, Ordering::Relaxed);
380        } else {
381            self.misses.fetch_add(1, Ordering::Relaxed);
382        }
383        found
384    }
385
386    pub(crate) fn should_build(&self, key: LookupIndexKey) -> bool {
387        let Ok(mut guard) = self.call_counts.write() else {
388            self.skipped_below_threshold.fetch_add(1, Ordering::Relaxed);
389            return false;
390        };
391        if guard.len() > CALL_COUNT_PRUNE_LIMIT {
392            guard.retain(|existing_key, _| existing_key.snapshot_id == key.snapshot_id);
393        }
394        let count = guard.entry(key).or_insert(0);
395        if *count == CAP_REJECTED {
396            self.skipped_cap.fetch_add(1, Ordering::Relaxed);
397            return false;
398        }
399        *count = count.saturating_add(1).min(CAP_REJECTED - 1);
400        if *count <= self.build_threshold {
401            self.skipped_below_threshold.fetch_add(1, Ordering::Relaxed);
402            return false;
403        }
404        true
405    }
406
407    pub(crate) fn would_exceed_cap(&self, bytes: usize) -> bool {
408        self.bytes_in_use
409            .load(Ordering::Relaxed)
410            .saturating_add(bytes)
411            > self.max_bytes
412    }
413
414    pub(crate) fn is_known_volatile(&self, key: &LookupIndexKey) -> bool {
415        let volatile_key = volatile_key(*key);
416        self.volatile_keys
417            .read()
418            .map(|guard| guard.contains_key(&volatile_key))
419            .unwrap_or(false)
420    }
421
422    pub(crate) fn note_volatile_key(&self, key: LookupIndexKey) {
423        if let Ok(mut guard) = self.volatile_keys.write() {
424            if guard.len() > CALL_COUNT_PRUNE_LIMIT {
425                guard.clear();
426            }
427            guard.insert(volatile_key(key), ());
428        }
429    }
430
431    pub(crate) fn insert_if_room(
432        &self,
433        key: LookupIndexKey,
434        index: LookupIndex,
435    ) -> Option<Arc<LookupIndex>> {
436        let bytes = index.bytes;
437        if let Ok(mut guard) = self.inner.write() {
438            if let Some(existing) = guard.get(&key) {
439                self.hits.fetch_add(1, Ordering::Relaxed);
440                return Some(existing.clone());
441            }
442            // Serialize duplicate detection, admission and accounting. The earlier
443            // preflight estimate is only a hint and cannot reserve concurrent space.
444            let current = self.bytes_in_use.load(Ordering::Relaxed);
445            if bytes > self.max_bytes.saturating_sub(current) {
446                self.skipped_cap.fetch_add(1, Ordering::Relaxed);
447                // Actual payloads (especially text) can exceed the preflight
448                // estimate. Do not rebuild and discard them on every later call.
449                if let Ok(mut counts) = self.call_counts.write() {
450                    counts.insert(key, CAP_REJECTED);
451                }
452                return None;
453            }
454            let index = Arc::new(index);
455            guard.insert(key, index.clone());
456            self.bytes_in_use.fetch_add(bytes, Ordering::Relaxed);
457            self.builds.fetch_add(1, Ordering::Relaxed);
458            Some(index)
459        } else {
460            None
461        }
462    }
463
464    pub(crate) fn note_skipped_volatile(&self) {
465        self.skipped_volatile.fetch_add(1, Ordering::Relaxed);
466    }
467
468    pub(crate) fn note_skipped_error(&self) {
469        self.skipped_error.fetch_add(1, Ordering::Relaxed);
470    }
471
472    pub(crate) fn note_skipped_tiny(&self) {
473        self.skipped_tiny.fetch_add(1, Ordering::Relaxed);
474    }
475
476    pub(crate) fn note_skipped_cap(&self) {
477        self.skipped_cap.fetch_add(1, Ordering::Relaxed);
478    }
479
480    pub(crate) fn reset_counters(&self) {
481        self.builds.store(0, Ordering::Relaxed);
482        self.hits.store(0, Ordering::Relaxed);
483        self.misses.store(0, Ordering::Relaxed);
484        self.skipped_volatile.store(0, Ordering::Relaxed);
485        self.skipped_error.store(0, Ordering::Relaxed);
486        self.skipped_tiny.store(0, Ordering::Relaxed);
487        self.skipped_cap.store(0, Ordering::Relaxed);
488        self.skipped_below_threshold.store(0, Ordering::Relaxed);
489    }
490
491    pub(crate) fn report(&self) -> LookupIndexCacheReport {
492        let guard = self.inner.read().ok();
493        let entries_count = guard.as_ref().map_or(0, |guard| guard.len());
494        LookupIndexCacheReport {
495            builds: self.builds.load(Ordering::Relaxed),
496            hits: self.hits.load(Ordering::Relaxed),
497            misses: self.misses.load(Ordering::Relaxed),
498            skipped_volatile: self.skipped_volatile.load(Ordering::Relaxed),
499            skipped_error: self.skipped_error.load(Ordering::Relaxed),
500            skipped_tiny: self.skipped_tiny.load(Ordering::Relaxed),
501            skipped_cap: self.skipped_cap.load(Ordering::Relaxed),
502            skipped_below_threshold: self.skipped_below_threshold.load(Ordering::Relaxed),
503            bytes_in_cache: self.bytes_in_use.load(Ordering::Relaxed),
504            entries_count,
505        }
506    }
507}
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512    use chrono::NaiveTime;
513
514    fn key(col: u32) -> LookupIndexKey {
515        LookupIndexKey {
516            sheet_id: 0,
517            start_row: 0,
518            start_col: col,
519            end_row: 128,
520            end_col: col,
521            axis: LookupAxis::ColumnInView(0),
522            snapshot_id: 1,
523        }
524    }
525    fn index(bytes: usize) -> LookupIndex {
526        LookupIndex {
527            len: 0,
528            date_system: DateSystem::Excel1900,
529            bytes,
530            entries: FxHashMap::default(),
531            cell_values: Box::new([]),
532        }
533    }
534
535    #[test]
536    fn exact_index_selects_signed_zero_and_temporal_zero_duplicates() {
537        assert_eq!(
538            LookupHashKey::from_literal(&LiteralValue::Empty, DateSystem::Excel1900),
539            None
540        );
541        assert_eq!(
542            LookupHashKey::from_needle(&LiteralValue::Empty, DateSystem::Excel1900),
543            Some(LookupHashKey::Number(0.0f64.to_bits()))
544        );
545        let midnight = LiteralValue::Time(NaiveTime::from_hms_opt(0, 0, 0).unwrap());
546        let values = vec![
547            LiteralValue::Empty,
548            LiteralValue::Boolean(false),
549            LiteralValue::Text("0".into()),
550            LiteralValue::Number(-0.0),
551            midnight.clone(),
552            LiteralValue::Number(0.0),
553            LiteralValue::Boolean(false),
554            LiteralValue::Text("0".into()),
555        ];
556        let view = RangeView::from_owned_rows(
557            values.into_iter().map(|value| vec![value]).collect(),
558            DateSystem::Excel1900,
559        );
560        let BuildOutcome::Built(index) =
561            LookupIndex::build(&view, LookupAxis::ColumnInView(0), DateSystem::Excel1900).unwrap()
562        else {
563            panic!("expected a lookup index");
564        };
565
566        assert_eq!(index.find_first_exact(&LiteralValue::Number(0.0)), Some(3));
567        assert_eq!(index.find_last_exact(&LiteralValue::Number(-0.0)), Some(5));
568        assert_eq!(index.find_first_exact(&midnight), Some(3));
569        assert_eq!(index.find_last_exact(&midnight), Some(5));
570        assert_eq!(index.find_first_exact(&LiteralValue::Empty), Some(3));
571        assert_eq!(index.find_last_exact(&LiteralValue::Empty), Some(5));
572
573        let temporal_only_view = RangeView::from_owned_rows(
574            vec![
575                vec![LiteralValue::Empty],
576                vec![LiteralValue::Boolean(false)],
577                vec![LiteralValue::Text("0".into())],
578                vec![midnight],
579            ],
580            DateSystem::Excel1900,
581        );
582        let BuildOutcome::Built(temporal_only) = LookupIndex::build(
583            &temporal_only_view,
584            LookupAxis::ColumnInView(0),
585            DateSystem::Excel1900,
586        )
587        .unwrap() else {
588            panic!("expected a temporal lookup index");
589        };
590        assert_eq!(
591            temporal_only.find_first_exact(&LiteralValue::Number(0.0)),
592            Some(3)
593        );
594    }
595
596    #[test]
597    fn concurrent_admission_and_duplicate_races_obey_cap() {
598        for duplicate in [false, true] {
599            let mut cache = LookupIndexCache::new(if duplicate { 1024 } else { 4096 });
600            let barrier = std::sync::Barrier::new(16);
601            std::thread::scope(|scope| {
602                let handles: Vec<_> = (0..16)
603                    .map(|i| {
604                        let cache = &cache;
605                        let barrier = &barrier;
606                        scope.spawn(move || {
607                            barrier.wait();
608                            cache.insert_if_room(key(if duplicate { 0 } else { i }), index(1024))
609                        })
610                    })
611                    .collect();
612                let admitted: Vec<_> = handles
613                    .into_iter()
614                    .filter_map(|h| h.join().unwrap())
615                    .collect();
616                assert_eq!(admitted.len(), if duplicate { 16 } else { 4 });
617                if duplicate {
618                    assert!(admitted.iter().all(|item| Arc::ptr_eq(item, &admitted[0])));
619                }
620            });
621            let report = cache.report();
622            assert_eq!(report.bytes_in_cache, cache.max_bytes);
623            assert_eq!(report.builds, if duplicate { 1 } else { 4 });
624            assert_eq!(report.entries_count, report.builds);
625            assert_eq!(report.skipped_cap, if duplicate { 0 } else { 12 });
626            cache.clear();
627            assert_eq!(cache.report().bytes_in_cache, 0);
628            assert_eq!(cache.report().entries_count, 0);
629            assert!(cache.insert_if_room(key(0), index(1024)).is_some());
630        }
631    }
632
633    #[test]
634    fn retained_text_and_duplicate_heap_payloads_are_charged() {
635        let mut entries = FxHashMap::default();
636        let mut text = String::with_capacity(4096);
637        text.push_str("LONG KEY");
638        let values = [LiteralValue::Text(text)];
639        let empty_bytes = retained_bytes(&[], &entries);
640        assert_eq!(
641            retained_bytes(&values, &entries) - empty_bytes,
642            std::mem::size_of::<LiteralValue>() + 4096
643        );
644        let mut dups = DuplicateIndices::default();
645        dups.all.extend(0..100);
646        let heap_bytes = dups.all.capacity() * std::mem::size_of::<usize>();
647        entries.insert(LookupHashKey::Text("long key".into()), dups);
648        let charged = retained_bytes(&values, &entries);
649        assert!(charged >= empty_bytes + 4096 + 8 + heap_bytes);
650        entries
651            .get_mut(&LookupHashKey::Text("long key".into()))
652            .unwrap()
653            .all = SmallVec::new();
654        assert_eq!(charged - retained_bytes(&values, &entries), heap_bytes);
655        let cache = LookupIndexCache::new(charged - 1);
656        assert!(cache.insert_if_room(key(0), index(charged)).is_none());
657    }
658}