antlr4_runtime/
prediction.rs

1use std::cmp::Ordering;
2use std::collections::{BTreeMap, BTreeSet, HashMap};
3use std::hash::{BuildHasherDefault, Hash, Hasher};
4use std::mem::size_of;
5
6pub const EMPTY_RETURN_STATE: usize = usize::MAX;
7const COMPACT_EMPTY_RETURN_STATE: u32 = u32::MAX;
8
9/// Lightweight `FxHash`-style hasher used on prediction hot paths.
10#[derive(Debug, Default)]
11pub struct PredictionFxHasher {
12    hash: u64,
13}
14
15const FX_ROT: u32 = 5;
16const FX_SEED: u64 = 0x51_7c_c1_b7_27_22_0a_95;
17
18impl Hasher for PredictionFxHasher {
19    #[inline]
20    fn write(&mut self, bytes: &[u8]) {
21        let mut bytes = bytes;
22        while bytes.len() >= 8 {
23            let (head, rest) = bytes.split_at(8);
24            let word = u64::from_le_bytes(head.try_into().expect("8-byte chunk"));
25            self.hash = (self.hash.rotate_left(FX_ROT) ^ word).wrapping_mul(FX_SEED);
26            bytes = rest;
27        }
28        for &byte in bytes {
29            self.hash = (self.hash.rotate_left(FX_ROT) ^ u64::from(byte)).wrapping_mul(FX_SEED);
30        }
31    }
32
33    #[inline]
34    fn write_u8(&mut self, value: u8) {
35        self.hash = (self.hash.rotate_left(FX_ROT) ^ u64::from(value)).wrapping_mul(FX_SEED);
36    }
37
38    #[inline]
39    fn write_u32(&mut self, value: u32) {
40        self.hash = (self.hash.rotate_left(FX_ROT) ^ u64::from(value)).wrapping_mul(FX_SEED);
41    }
42
43    #[inline]
44    fn write_u64(&mut self, value: u64) {
45        self.hash = (self.hash.rotate_left(FX_ROT) ^ value).wrapping_mul(FX_SEED);
46    }
47
48    #[inline]
49    fn write_usize(&mut self, value: usize) {
50        self.hash = (self.hash.rotate_left(FX_ROT) ^ value as u64).wrapping_mul(FX_SEED);
51    }
52
53    #[inline]
54    fn write_i32(&mut self, value: i32) {
55        self.write_u32(i32::cast_unsigned(value));
56    }
57
58    #[inline]
59    fn finish(&self) -> u64 {
60        self.hash
61    }
62}
63
64type FxHashMap<K, V> = HashMap<K, V, BuildHasherDefault<PredictionFxHasher>>;
65
66/// Store-local identity for one canonical prediction-context graph node.
67#[repr(transparent)]
68#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
69pub struct ContextId(u32);
70
71pub const EMPTY_CONTEXT: ContextId = ContextId(0);
72
73#[derive(Clone, Copy, Debug, Eq, PartialEq)]
74enum ContextTag {
75    Empty,
76    Singleton,
77    Array,
78}
79
80#[derive(Clone, Copy, Debug)]
81struct ContextRecord {
82    tag: ContextTag,
83    cached_hash: u64,
84    parent_or_start: u32,
85    return_state_or_len: u32,
86}
87
88/// Allocation and interning totals for one prediction-context arena.
89#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
90pub struct PredictionContextStats {
91    pub contexts_created: usize,
92    pub singleton_contexts: usize,
93    pub array_contexts: usize,
94    pub array_entries: usize,
95    pub interner_hits: usize,
96    pub pooled_bytes: usize,
97    /// Element storage implied by retained capacities, excluding allocator and
98    /// hash-table control metadata.
99    pub retained_bytes: usize,
100    pub context_capacity: usize,
101    pub array_parent_capacity: usize,
102    pub array_return_state_capacity: usize,
103    pub interner_capacity: usize,
104    pub workspace_merge_cache_entries: usize,
105    pub workspace_merge_cache_capacity: usize,
106    pub workspace_entry_capacity: usize,
107    pub outer_context_cache_hits: usize,
108    pub outer_context_cache_misses: usize,
109}
110
111/// Canonical compact storage paired with one learned parser DFA store.
112#[derive(Debug)]
113pub(crate) struct ContextArena {
114    records: Vec<ContextRecord>,
115    array_parents: Vec<ContextId>,
116    array_return_states: Vec<u32>,
117    interner_heads: FxHashMap<u64, ContextId>,
118    interner_next: Vec<Option<ContextId>>,
119    interner_hits: usize,
120    #[cfg(debug_assertions)]
121    generation: u64,
122}
123
124#[cfg(debug_assertions)]
125fn next_context_arena_generation() -> u64 {
126    use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
127
128    static NEXT_GENERATION: AtomicU64 = AtomicU64::new(1);
129    NEXT_GENERATION.fetch_add(1, AtomicOrdering::Relaxed)
130}
131
132impl ContextArena {
133    pub(crate) fn new() -> Self {
134        let empty = ContextRecord {
135            tag: ContextTag::Empty,
136            cached_hash: prediction_context_empty_hash(),
137            parent_or_start: 0,
138            return_state_or_len: 0,
139        };
140        let mut interner_heads = FxHashMap::default();
141        interner_heads.insert(empty.cached_hash, EMPTY_CONTEXT);
142        Self {
143            records: vec![empty],
144            array_parents: Vec::new(),
145            array_return_states: Vec::new(),
146            interner_heads,
147            interner_next: vec![None],
148            interner_hits: 0,
149            #[cfg(debug_assertions)]
150            generation: next_context_arena_generation(),
151        }
152    }
153
154    #[cfg(debug_assertions)]
155    pub(crate) const fn generation(&self) -> u64 {
156        self.generation
157    }
158
159    pub(crate) fn stats(&self) -> PredictionContextStats {
160        let mut singleton_contexts = 0;
161        let mut array_contexts = 0;
162        for record in &self.records {
163            match record.tag {
164                ContextTag::Empty => {}
165                ContextTag::Singleton => singleton_contexts += 1,
166                ContextTag::Array => array_contexts += 1,
167            }
168        }
169        PredictionContextStats {
170            contexts_created: self.records.len(),
171            singleton_contexts,
172            array_contexts,
173            array_entries: self.array_parents.len(),
174            interner_hits: self.interner_hits,
175            pooled_bytes: self.records.len() * size_of::<ContextRecord>()
176                + self.array_parents.len() * size_of::<ContextId>()
177                + self.array_return_states.len() * size_of::<u32>()
178                + self.interner_next.len() * size_of::<Option<ContextId>>(),
179            retained_bytes: self.records.capacity() * size_of::<ContextRecord>()
180                + self.array_parents.capacity() * size_of::<ContextId>()
181                + self.array_return_states.capacity() * size_of::<u32>()
182                + self.interner_heads.capacity() * size_of::<(u64, ContextId)>()
183                + self.interner_next.capacity() * size_of::<Option<ContextId>>(),
184            context_capacity: self.records.capacity(),
185            array_parent_capacity: self.array_parents.capacity(),
186            array_return_state_capacity: self.array_return_states.capacity(),
187            interner_capacity: self.interner_heads.capacity(),
188            workspace_merge_cache_entries: 0,
189            workspace_merge_cache_capacity: 0,
190            workspace_entry_capacity: 0,
191            outer_context_cache_hits: 0,
192            outer_context_cache_misses: 0,
193        }
194    }
195
196    pub(crate) fn singleton(&mut self, parent: ContextId, return_state: usize) -> ContextId {
197        self.assert_valid(parent);
198        if return_state == EMPTY_RETURN_STATE {
199            return EMPTY_CONTEXT;
200        }
201        #[cfg(feature = "perf-counters")]
202        crate::perf::record_context_cache_call();
203        let return_state =
204            u32::try_from(return_state).expect("prediction return state must fit in u32");
205        let cached_hash = prediction_context_singleton_hash(self.cached_hash(parent), return_state);
206        if let Some(existing) = self.find_interned(cached_hash, |record| {
207            record.tag == ContextTag::Singleton
208                && record.parent_or_start == parent.0
209                && record.return_state_or_len == return_state
210        }) {
211            self.interner_hits = self.interner_hits.saturating_add(1);
212            #[cfg(feature = "perf-counters")]
213            crate::perf::record_context_cache_hit();
214            return existing;
215        }
216        #[cfg(feature = "perf-counters")]
217        {
218            crate::perf::record_context_cache_miss();
219            crate::perf::record_context_cache_insert();
220        }
221        self.push_record(ContextRecord {
222            tag: ContextTag::Singleton,
223            cached_hash,
224            parent_or_start: parent.0,
225            return_state_or_len: return_state,
226        })
227    }
228
229    fn intern_entries(&mut self, entries: &[(ContextId, u32)]) -> ContextId {
230        match entries {
231            [] => EMPTY_CONTEXT,
232            [(parent, return_state)] => {
233                if *return_state == COMPACT_EMPTY_RETURN_STATE {
234                    EMPTY_CONTEXT
235                } else {
236                    self.singleton(
237                        *parent,
238                        usize::try_from(*return_state).expect("u32 return state fits in usize"),
239                    )
240                }
241            }
242            _ => {
243                debug_assert!(
244                    entries
245                        .windows(2)
246                        .all(|pair| { compare_entries(pair[0], pair[1]) == Ordering::Less })
247                );
248                #[cfg(feature = "perf-counters")]
249                crate::perf::record_context_cache_call();
250                let cached_hash = prediction_context_array_hash(self, entries);
251                if let Some(existing) = self.find_interned(cached_hash, |record| {
252                    if record.tag != ContextTag::Array
253                        || usize::try_from(record.return_state_or_len).ok() != Some(entries.len())
254                    {
255                        return false;
256                    }
257                    let start =
258                        usize::try_from(record.parent_or_start).expect("u32 pool index fits usize");
259                    let end = start + entries.len();
260                    self.array_parents[start..end]
261                        .iter()
262                        .copied()
263                        .zip(self.array_return_states[start..end].iter().copied())
264                        .eq(entries.iter().copied())
265                }) {
266                    self.interner_hits = self.interner_hits.saturating_add(1);
267                    #[cfg(feature = "perf-counters")]
268                    crate::perf::record_context_cache_hit();
269                    return existing;
270                }
271                #[cfg(feature = "perf-counters")]
272                {
273                    crate::perf::record_context_cache_miss();
274                    crate::perf::record_context_cache_insert();
275                }
276                let start = u32::try_from(self.array_parents.len())
277                    .expect("prediction-context parent pool must fit in u32");
278                let len = u32::try_from(entries.len())
279                    .expect("prediction-context array length must fit in u32");
280                self.array_parents
281                    .extend(entries.iter().map(|(parent, _)| *parent));
282                self.array_return_states
283                    .extend(entries.iter().map(|(_, return_state)| *return_state));
284                self.push_record(ContextRecord {
285                    tag: ContextTag::Array,
286                    cached_hash,
287                    parent_or_start: start,
288                    return_state_or_len: len,
289                })
290            }
291        }
292    }
293
294    fn find_interned(
295        &self,
296        cached_hash: u64,
297        matches: impl Fn(&ContextRecord) -> bool,
298    ) -> Option<ContextId> {
299        let mut candidate = self.interner_heads.get(&cached_hash).copied();
300        while let Some(id) = candidate {
301            let index = usize::try_from(id.0).expect("u32 context ID fits in usize");
302            let record = &self.records[index];
303            if matches(record) {
304                return Some(id);
305            }
306            candidate = self.interner_next[index];
307        }
308        None
309    }
310
311    fn push_record(&mut self, record: ContextRecord) -> ContextId {
312        let id = ContextId(
313            u32::try_from(self.records.len()).expect("prediction-context arena must fit in u32"),
314        );
315        let previous = self.interner_heads.insert(record.cached_hash, id);
316        self.records.push(record);
317        self.interner_next.push(previous);
318        id
319    }
320
321    pub(crate) fn merge(
322        &mut self,
323        left: ContextId,
324        right: ContextId,
325        root_is_wildcard: bool,
326        workspace: &mut PredictionWorkspace,
327    ) -> ContextId {
328        self.assert_valid(left);
329        self.assert_valid(right);
330        #[cfg(feature = "perf-counters")]
331        crate::perf::record_context_merge_call();
332        if left == right {
333            #[cfg(feature = "perf-counters")]
334            crate::perf::record_context_merge_identical();
335            return left;
336        }
337        let key = MergeKey::new(left, right, root_is_wildcard);
338        if let Some(merged) = workspace.merge_cache.get(&key).copied() {
339            #[cfg(feature = "perf-counters")]
340            crate::perf::record_context_merge_cache_hit();
341            return merged;
342        }
343        #[cfg(feature = "perf-counters")]
344        {
345            crate::perf::record_context_merge_cache_miss();
346            crate::perf::record_context_merge_uncached();
347        }
348        let merged = if root_is_wildcard && (left == EMPTY_CONTEXT || right == EMPTY_CONTEXT) {
349            EMPTY_CONTEXT
350        } else {
351            self.merge_uncached(left, right, workspace)
352        };
353        workspace.merge_cache.insert(key, merged);
354        merged
355    }
356
357    fn merge_uncached(
358        &mut self,
359        left: ContextId,
360        right: ContextId,
361        workspace: &mut PredictionWorkspace,
362    ) -> ContextId {
363        match (self.tag(left), self.tag(right)) {
364            (ContextTag::Array, ContextTag::Array) => self.merge_arrays(left, right, workspace),
365            (ContextTag::Array, _) => {
366                let entry = self.first_entry(right);
367                self.merge_array_with_entry(left, entry, false, workspace)
368            }
369            (_, ContextTag::Array) => {
370                let entry = self.first_entry(left);
371                self.merge_array_with_entry(right, entry, true, workspace)
372            }
373            _ => self.merge_two_entries(self.first_entry(left), self.first_entry(right), workspace),
374        }
375    }
376
377    fn merge_two_entries(
378        &mut self,
379        left: (ContextId, u32),
380        right: (ContextId, u32),
381        workspace: &mut PredictionWorkspace,
382    ) -> ContextId {
383        if left == right {
384            return self.intern_entries(std::slice::from_ref(&left));
385        }
386        workspace.entries.clear();
387        if compare_entries(right, left) == Ordering::Less {
388            workspace.entries.extend([right, left]);
389        } else {
390            workspace.entries.extend([left, right]);
391        }
392        self.intern_entries(&workspace.entries)
393    }
394
395    fn merge_array_with_entry(
396        &mut self,
397        array: ContextId,
398        entry: (ContextId, u32),
399        entry_on_left: bool,
400        workspace: &mut PredictionWorkspace,
401    ) -> ContextId {
402        let array_len = self.len(array);
403        let mut insert_index = array_len;
404        for index in 0..array_len {
405            let current = self.entry(array, index).expect("array entry in range");
406            let ordering = compare_entries(entry, current);
407            if ordering == Ordering::Equal {
408                return array;
409            }
410            let should_insert = if entry_on_left {
411                ordering != Ordering::Greater
412            } else {
413                ordering == Ordering::Less
414            };
415            if should_insert {
416                insert_index = index;
417                break;
418            }
419        }
420
421        workspace.entries.clear();
422        for index in 0..insert_index {
423            workspace
424                .entries
425                .push(self.entry(array, index).expect("array entry in range"));
426        }
427        workspace.entries.push(entry);
428        for index in insert_index..array_len {
429            workspace
430                .entries
431                .push(self.entry(array, index).expect("array entry in range"));
432        }
433        self.intern_entries(&workspace.entries)
434    }
435
436    fn merge_arrays(
437        &mut self,
438        left: ContextId,
439        right: ContextId,
440        workspace: &mut PredictionWorkspace,
441    ) -> ContextId {
442        workspace.entries.clear();
443        let mut left_index = 0;
444        let mut right_index = 0;
445        while left_index < self.len(left) && right_index < self.len(right) {
446            let left_entry = self.entry(left, left_index).expect("array entry in range");
447            let right_entry = self
448                .entry(right, right_index)
449                .expect("array entry in range");
450            match compare_entries(left_entry, right_entry) {
451                Ordering::Less => {
452                    workspace.entries.push(left_entry);
453                    left_index += 1;
454                }
455                Ordering::Greater => {
456                    workspace.entries.push(right_entry);
457                    right_index += 1;
458                }
459                Ordering::Equal => {
460                    workspace.entries.push(left_entry);
461                    left_index += 1;
462                    right_index += 1;
463                }
464            }
465        }
466        while left_index < self.len(left) {
467            workspace
468                .entries
469                .push(self.entry(left, left_index).expect("array entry in range"));
470            left_index += 1;
471        }
472        while right_index < self.len(right) {
473            workspace.entries.push(
474                self.entry(right, right_index)
475                    .expect("array entry in range"),
476            );
477            right_index += 1;
478        }
479        self.intern_entries(&workspace.entries)
480    }
481
482    pub(crate) fn len(&self, context: ContextId) -> usize {
483        let record = self.record(context);
484        match record.tag {
485            ContextTag::Empty | ContextTag::Singleton => 1,
486            ContextTag::Array => usize::try_from(record.return_state_or_len)
487                .expect("u32 context length fits in usize"),
488        }
489    }
490
491    pub(crate) fn is_empty(&self, context: ContextId) -> bool {
492        self.assert_valid(context);
493        context == EMPTY_CONTEXT
494    }
495
496    pub(crate) fn has_empty_path(&self, context: ContextId) -> bool {
497        if context == EMPTY_CONTEXT {
498            return true;
499        }
500        let record = self.record(context);
501        match record.tag {
502            ContextTag::Empty => true,
503            ContextTag::Singleton => false,
504            ContextTag::Array => {
505                let len = usize::try_from(record.return_state_or_len)
506                    .expect("u32 context length fits in usize");
507                let start = usize::try_from(record.parent_or_start)
508                    .expect("u32 context pool index fits in usize");
509                self.array_return_states[start + len - 1] == COMPACT_EMPTY_RETURN_STATE
510            }
511        }
512    }
513
514    pub(crate) fn return_state(&self, context: ContextId, index: usize) -> Option<usize> {
515        let (_, return_state) = self.entry(context, index)?;
516        Some(expand_return_state(return_state))
517    }
518
519    pub(crate) fn parent(&self, context: ContextId, index: usize) -> Option<ContextId> {
520        if context == EMPTY_CONTEXT {
521            self.assert_valid(context);
522            return None;
523        }
524        self.entry(context, index).map(|(parent, _)| parent)
525    }
526
527    fn first_entry(&self, context: ContextId) -> (ContextId, u32) {
528        self.entry(context, 0)
529            .expect("empty and singleton contexts have one logical entry")
530    }
531
532    fn entry(&self, context: ContextId, index: usize) -> Option<(ContextId, u32)> {
533        let record = self.record(context);
534        match record.tag {
535            ContextTag::Empty if index == 0 => Some((EMPTY_CONTEXT, COMPACT_EMPTY_RETURN_STATE)),
536            ContextTag::Singleton if index == 0 => Some((
537                ContextId(record.parent_or_start),
538                record.return_state_or_len,
539            )),
540            ContextTag::Array => {
541                let len = usize::try_from(record.return_state_or_len).ok()?;
542                if index >= len {
543                    return None;
544                }
545                let start = usize::try_from(record.parent_or_start).ok()?;
546                Some((
547                    self.array_parents[start + index],
548                    self.array_return_states[start + index],
549                ))
550            }
551            ContextTag::Empty | ContextTag::Singleton => None,
552        }
553    }
554
555    fn tag(&self, context: ContextId) -> ContextTag {
556        self.record(context).tag
557    }
558
559    fn cached_hash(&self, context: ContextId) -> u64 {
560        self.record(context).cached_hash
561    }
562
563    fn record(&self, context: ContextId) -> &ContextRecord {
564        self.assert_valid(context);
565        &self.records[usize::try_from(context.0).expect("u32 context ID fits in usize")]
566    }
567
568    pub(crate) fn assert_valid(&self, context: ContextId) {
569        assert!(
570            usize::try_from(context.0).is_ok_and(|index| index < self.records.len()),
571            "prediction ContextId does not belong to this store"
572        );
573    }
574
575    pub(crate) fn import_all(
576        &mut self,
577        source: &Self,
578        workspace: &mut PredictionWorkspace,
579    ) -> Vec<ContextId> {
580        let mut remap = Vec::with_capacity(source.records.len());
581        remap.push(EMPTY_CONTEXT);
582        for source_index in 1..source.records.len() {
583            let source_id = ContextId(
584                u32::try_from(source_index).expect("source prediction-context ID fits in u32"),
585            );
586            let imported = match source.tag(source_id) {
587                ContextTag::Empty => EMPTY_CONTEXT,
588                ContextTag::Singleton => {
589                    let (parent, return_state) = source.first_entry(source_id);
590                    let parent_index =
591                        usize::try_from(parent.0).expect("u32 context ID fits usize");
592                    assert!(
593                        parent_index < remap.len(),
594                        "prediction contexts must reference earlier arena records"
595                    );
596                    self.singleton(remap[parent_index], expand_return_state(return_state))
597                }
598                ContextTag::Array => {
599                    workspace.entries.clear();
600                    for entry_index in 0..source.len(source_id) {
601                        let (parent, return_state) = source
602                            .entry(source_id, entry_index)
603                            .expect("source array entry in range");
604                        let parent_index =
605                            usize::try_from(parent.0).expect("u32 context ID fits usize");
606                        assert!(
607                            parent_index < remap.len(),
608                            "prediction contexts must reference earlier arena records"
609                        );
610                        workspace.entries.push((remap[parent_index], return_state));
611                    }
612                    workspace
613                        .entries
614                        .sort_unstable_by(|left, right| compare_entries(*left, *right));
615                    workspace.entries.dedup();
616                    self.intern_entries(&workspace.entries)
617                }
618            };
619            remap.push(imported);
620        }
621        remap
622    }
623}
624
625impl Default for ContextArena {
626    fn default() -> Self {
627        Self::new()
628    }
629}
630
631fn compare_entries(left: (ContextId, u32), right: (ContextId, u32)) -> Ordering {
632    left.1.cmp(&right.1).then_with(|| left.0.cmp(&right.0))
633}
634
635fn expand_return_state(return_state: u32) -> usize {
636    if return_state == COMPACT_EMPTY_RETURN_STATE {
637        EMPTY_RETURN_STATE
638    } else {
639        usize::try_from(return_state).expect("u32 return state fits in usize")
640    }
641}
642
643fn prediction_context_empty_hash() -> u64 {
644    let mut hasher = PredictionFxHasher::default();
645    hasher.write_u8(0);
646    hasher.finish()
647}
648
649fn prediction_context_singleton_hash(parent_hash: u64, return_state: u32) -> u64 {
650    let mut hasher = PredictionFxHasher::default();
651    hasher.write_u8(1);
652    hasher.write_u64(parent_hash);
653    hasher.write_u32(return_state);
654    hasher.finish()
655}
656
657fn prediction_context_array_hash(arena: &ContextArena, entries: &[(ContextId, u32)]) -> u64 {
658    let mut hasher = PredictionFxHasher::default();
659    hasher.write_u8(2);
660    hasher.write_usize(entries.len());
661    for (parent, _) in entries {
662        hasher.write_u64(arena.cached_hash(*parent));
663    }
664    hasher.write_usize(entries.len());
665    for (_, return_state) in entries {
666        hasher.write_u32(*return_state);
667    }
668    hasher.finish()
669}
670
671#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
672struct MergeKey {
673    left: ContextId,
674    right: ContextId,
675    root_is_wildcard: bool,
676}
677
678impl MergeKey {
679    fn new(left: ContextId, right: ContextId, root_is_wildcard: bool) -> Self {
680        let (left, right) = if right < left {
681            (right, left)
682        } else {
683            (left, right)
684        };
685        Self {
686            left,
687            right,
688            root_is_wildcard,
689        }
690    }
691}
692
693const MAX_RETAINED_MERGE_CACHE_ENTRIES: usize = 65_536;
694const MAX_RETAINED_CONTEXT_ENTRIES: usize = 16_384;
695
696/// Reusable per-prediction merge cache and temporary compact entry storage.
697#[derive(Debug, Default)]
698pub(crate) struct PredictionWorkspace {
699    merge_cache: FxHashMap<MergeKey, ContextId>,
700    entries: Vec<(ContextId, u32)>,
701}
702
703impl PredictionWorkspace {
704    pub(crate) fn reset(&mut self) {
705        if self.merge_cache.capacity() > MAX_RETAINED_MERGE_CACHE_ENTRIES {
706            self.merge_cache = FxHashMap::default();
707        } else {
708            self.merge_cache.clear();
709        }
710        if self.entries.capacity() > MAX_RETAINED_CONTEXT_ENTRIES {
711            self.entries = Vec::new();
712        } else {
713            self.entries.clear();
714        }
715    }
716
717    pub(crate) fn merge_cache_capacity(&self) -> usize {
718        self.merge_cache.capacity()
719    }
720
721    pub(crate) fn merge_cache_len(&self) -> usize {
722        self.merge_cache.len()
723    }
724
725    pub(crate) const fn entry_capacity(&self) -> usize {
726        self.entries.capacity()
727    }
728
729    pub(crate) fn retained_bytes(&self) -> usize {
730        self.merge_cache.capacity() * size_of::<(MergeKey, ContextId)>()
731            + self.entries.capacity() * size_of::<(ContextId, u32)>()
732    }
733}
734
735#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
736pub enum SemanticContext {
737    None,
738    Predicate {
739        rule_index: usize,
740        pred_index: usize,
741        context_dependent: bool,
742    },
743    Precedence {
744        precedence: i32,
745    },
746    And(Vec<Self>),
747    Or(Vec<Self>),
748}
749
750impl SemanticContext {
751    pub const fn none() -> Self {
752        Self::None
753    }
754
755    pub fn and(left: Self, right: Self) -> Self {
756        combine_semantic_context(left, right, true)
757    }
758
759    pub fn or(left: Self, right: Self) -> Self {
760        combine_semantic_context(left, right, false)
761    }
762
763    pub const fn is_none(&self) -> bool {
764        matches!(self, Self::None)
765    }
766}
767
768fn combine_semantic_context(
769    left: SemanticContext,
770    right: SemanticContext,
771    and: bool,
772) -> SemanticContext {
773    if left == right {
774        return left;
775    }
776    if left.is_none() {
777        return right;
778    }
779    if right.is_none() {
780        return left;
781    }
782    let mut entries = Vec::new();
783    for context in [left, right] {
784        match (and, context) {
785            (true, SemanticContext::And(children)) | (false, SemanticContext::Or(children)) => {
786                entries.extend(children);
787            }
788            (_, other) => entries.push(other),
789        }
790    }
791    entries.sort();
792    entries.dedup();
793    if and {
794        SemanticContext::And(entries)
795    } else {
796        SemanticContext::Or(entries)
797    }
798}
799
800#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
801pub(crate) struct AtnConfig {
802    pub(crate) state: usize,
803    pub(crate) alt: usize,
804    pub(crate) context: ContextId,
805    pub(crate) semantic_context: SemanticContext,
806    pub(crate) reaches_into_outer_context: usize,
807    pub(crate) precedence_filter_suppressed: bool,
808    #[cfg(debug_assertions)]
809    context_generation: u64,
810}
811
812impl AtnConfig {
813    pub(crate) fn new(state: usize, alt: usize, context: ContextId, arena: &ContextArena) -> Self {
814        arena.assert_valid(context);
815        Self {
816            state,
817            alt,
818            context,
819            semantic_context: SemanticContext::None,
820            reaches_into_outer_context: 0,
821            precedence_filter_suppressed: false,
822            #[cfg(debug_assertions)]
823            context_generation: arena.generation(),
824        }
825    }
826
827    #[must_use]
828    #[cfg(test)]
829    pub(crate) fn with_semantic_context(mut self, semantic_context: SemanticContext) -> Self {
830        self.semantic_context = semantic_context;
831        self
832    }
833
834    pub(crate) fn set_context(&mut self, context: ContextId, arena: &ContextArena) {
835        arena.assert_valid(context);
836        self.context = context;
837        #[cfg(debug_assertions)]
838        {
839            self.context_generation = arena.generation();
840        }
841    }
842
843    pub(crate) fn moved_to(&self, state: usize, context: ContextId, arena: &ContextArena) -> Self {
844        let mut moved = Self::new(state, self.alt, context, arena);
845        moved.semantic_context = self.semantic_context.clone();
846        moved.reaches_into_outer_context = self.reaches_into_outer_context;
847        moved.precedence_filter_suppressed = self.precedence_filter_suppressed;
848        moved
849    }
850
851    pub(crate) fn assert_store(&self, arena: &ContextArena) {
852        arena.assert_valid(self.context);
853        #[cfg(debug_assertions)]
854        assert_eq!(
855            self.context_generation,
856            arena.generation(),
857            "ATN config carries a ContextId from another prediction store"
858        );
859    }
860}
861
862#[derive(Clone, Debug, Default)]
863pub(crate) struct AtnConfigSet {
864    configs: Vec<AtnConfig>,
865    config_index: FxHashMap<AtnConfigKey, usize>,
866    full_context: bool,
867    unique_alt: Option<usize>,
868    conflicting_alts: BTreeSet<usize>,
869    has_semantic_context: bool,
870    dips_into_outer_context: bool,
871    readonly: bool,
872}
873
874impl AtnConfigSet {
875    pub(crate) fn new() -> Self {
876        Self::default()
877    }
878
879    pub(crate) fn new_full_context(full_context: bool) -> Self {
880        Self {
881            configs: Vec::new(),
882            config_index: FxHashMap::default(),
883            full_context,
884            unique_alt: None,
885            conflicting_alts: BTreeSet::new(),
886            has_semantic_context: false,
887            dips_into_outer_context: false,
888            readonly: false,
889        }
890    }
891
892    /// Adds a configuration, merging contexts for equivalent config keys.
893    pub(crate) fn add(
894        &mut self,
895        config: AtnConfig,
896        arena: &mut ContextArena,
897        workspace: &mut PredictionWorkspace,
898    ) -> bool {
899        assert!(!self.readonly, "cannot mutate readonly ATN config set");
900        config.assert_store(arena);
901        #[cfg(feature = "perf-counters")]
902        crate::perf::record_config_add_call();
903        if !config.semantic_context.is_none() {
904            self.has_semantic_context = true;
905        }
906        if config.reaches_into_outer_context > 0 {
907            self.dips_into_outer_context = true;
908        }
909        let key = AtnConfigKey::from(&config);
910        if let Some(existing_index) = self.config_index.get(&key).copied() {
911            #[cfg(feature = "perf-counters")]
912            crate::perf::record_config_merge();
913            let existing = &mut self.configs[existing_index];
914            existing.assert_store(arena);
915            existing.context = arena.merge(
916                existing.context,
917                config.context,
918                !self.full_context,
919                workspace,
920            );
921            existing.reaches_into_outer_context = existing
922                .reaches_into_outer_context
923                .max(config.reaches_into_outer_context);
924            existing.precedence_filter_suppressed |= config.precedence_filter_suppressed;
925            self.conflicting_alts.clear();
926            false
927        } else {
928            let index = self.configs.len();
929            self.config_index.insert(key, index);
930            self.configs.push(config);
931            #[cfg(feature = "perf-counters")]
932            crate::perf::record_config_insert(self.configs.len());
933            self.unique_alt = None;
934            self.conflicting_alts.clear();
935            true
936        }
937    }
938
939    pub(crate) fn configs(&self) -> &[AtnConfig] {
940        &self.configs
941    }
942
943    pub(crate) fn into_configs(self) -> Vec<AtnConfig> {
944        self.configs
945    }
946
947    pub(crate) const fn is_empty(&self) -> bool {
948        self.configs.is_empty()
949    }
950
951    pub(crate) const fn len(&self) -> usize {
952        self.configs.len()
953    }
954
955    pub(crate) fn set_readonly(&mut self, readonly: bool) {
956        self.readonly = readonly;
957        if readonly {
958            self.config_index = FxHashMap::default();
959            self.conflicting_alts.clear();
960        }
961    }
962
963    pub(crate) const fn full_context(&self) -> bool {
964        self.full_context
965    }
966
967    pub(crate) const fn has_semantic_context(&self) -> bool {
968        self.has_semantic_context
969    }
970
971    pub(crate) fn unique_alt(&mut self) -> Option<usize> {
972        if self.unique_alt.is_none() {
973            self.unique_alt = unique_alt(self.configs());
974        }
975        self.unique_alt
976    }
977
978    pub(crate) fn alts(&self) -> BTreeSet<usize> {
979        self.configs.iter().map(|config| config.alt).collect()
980    }
981
982    pub(crate) fn conflicting_alt_subsets(&self) -> Vec<BTreeSet<usize>> {
983        conflicting_alt_subsets(self.configs())
984    }
985
986    pub(crate) fn conflicting_alts(&mut self) -> BTreeSet<usize> {
987        if self.conflicting_alts.is_empty() {
988            self.conflicting_alts = self
989                .conflicting_alt_subsets()
990                .into_iter()
991                .filter(|alts| alts.len() > 1)
992                .flatten()
993                .collect();
994        }
995        self.conflicting_alts.clone()
996    }
997
998    pub(crate) fn remap_contexts(&mut self, remap: &[ContextId], arena: &ContextArena) {
999        for config in &mut self.configs {
1000            let index = usize::try_from(config.context.0).expect("u32 context ID fits usize");
1001            config.set_context(
1002                *remap
1003                    .get(index)
1004                    .expect("every imported context ID has a remap"),
1005                arena,
1006            );
1007        }
1008        self.config_index.clear();
1009        if !self.readonly {
1010            for (index, config) in self.configs.iter().enumerate() {
1011                self.config_index.insert(AtnConfigKey::from(config), index);
1012            }
1013        }
1014    }
1015
1016    pub(crate) fn fingerprint(&self) -> u64 {
1017        let mut hasher = PredictionFxHasher::default();
1018        self.configs.hash(&mut hasher);
1019        self.full_context.hash(&mut hasher);
1020        self.has_semantic_context.hash(&mut hasher);
1021        self.dips_into_outer_context.hash(&mut hasher);
1022        hasher.finish()
1023    }
1024
1025    pub(crate) fn retained_bytes(&self) -> usize {
1026        self.configs.capacity() * size_of::<AtnConfig>()
1027            + self.config_index.capacity() * size_of::<(AtnConfigKey, usize)>()
1028    }
1029}
1030
1031impl PartialEq for AtnConfigSet {
1032    fn eq(&self, other: &Self) -> bool {
1033        self.configs == other.configs
1034            && self.full_context == other.full_context
1035            && self.has_semantic_context == other.has_semantic_context
1036            && self.dips_into_outer_context == other.dips_into_outer_context
1037    }
1038}
1039
1040impl Eq for AtnConfigSet {}
1041
1042impl Ord for AtnConfigSet {
1043    fn cmp(&self, other: &Self) -> Ordering {
1044        self.configs
1045            .cmp(&other.configs)
1046            .then_with(|| self.full_context.cmp(&other.full_context))
1047            .then_with(|| self.has_semantic_context.cmp(&other.has_semantic_context))
1048            .then_with(|| {
1049                self.dips_into_outer_context
1050                    .cmp(&other.dips_into_outer_context)
1051            })
1052    }
1053}
1054
1055impl PartialOrd for AtnConfigSet {
1056    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1057        Some(self.cmp(other))
1058    }
1059}
1060
1061#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
1062struct AtnConfigKey {
1063    state: usize,
1064    alt: usize,
1065    semantic_context: SemanticContext,
1066}
1067
1068impl From<&AtnConfig> for AtnConfigKey {
1069    fn from(config: &AtnConfig) -> Self {
1070        Self {
1071            state: config.state,
1072            alt: config.alt,
1073            semantic_context: config.semantic_context.clone(),
1074        }
1075    }
1076}
1077
1078pub(crate) fn unique_alt(configs: &[AtnConfig]) -> Option<usize> {
1079    let mut alt = None;
1080    for config in configs {
1081        match alt {
1082            None => alt = Some(config.alt),
1083            Some(existing) if existing == config.alt => {}
1084            Some(_) => return None,
1085        }
1086    }
1087    alt
1088}
1089
1090pub(crate) fn conflicting_alt_subsets(configs: &[AtnConfig]) -> Vec<BTreeSet<usize>> {
1091    let mut by_state_context = FxHashMap::<(usize, ContextId), BTreeSet<usize>>::default();
1092    for config in configs {
1093        by_state_context
1094            .entry((config.state, config.context))
1095            .or_default()
1096            .insert(config.alt);
1097    }
1098    by_state_context.into_values().collect()
1099}
1100
1101pub(crate) fn all_subsets_conflict(alt_subsets: &[BTreeSet<usize>]) -> bool {
1102    alt_subsets.iter().all(|alts| alts.len() > 1)
1103}
1104
1105pub(crate) fn all_subsets_equal(alt_subsets: &[BTreeSet<usize>]) -> bool {
1106    let mut subsets = alt_subsets.iter();
1107    let Some(first) = subsets.next() else {
1108        return true;
1109    };
1110    subsets.all(|alts| alts == first)
1111}
1112
1113pub(crate) fn single_viable_alt(alt_subsets: &[BTreeSet<usize>]) -> Option<usize> {
1114    let mut result = None;
1115    for alts in alt_subsets {
1116        let min_alt = alts.iter().next().copied()?;
1117        match result {
1118            None => result = Some(min_alt),
1119            Some(existing) if existing == min_alt => {}
1120            Some(_) => return None,
1121        }
1122    }
1123    result
1124}
1125
1126pub(crate) fn has_sll_conflict_terminating_prediction(
1127    configs: &AtnConfigSet,
1128    is_rule_stop_state: impl Fn(usize) -> bool,
1129) -> bool {
1130    if configs
1131        .configs()
1132        .iter()
1133        .all(|config| is_rule_stop_state(config.state))
1134    {
1135        return true;
1136    }
1137    let alt_subsets = configs.conflicting_alt_subsets();
1138    alt_subsets.iter().any(|alts| alts.len() > 1)
1139        && !has_state_associated_with_one_alt(configs.configs())
1140}
1141
1142fn has_state_associated_with_one_alt(configs: &[AtnConfig]) -> bool {
1143    let mut by_state = BTreeMap::<usize, BTreeSet<usize>>::new();
1144    for config in configs {
1145        by_state.entry(config.state).or_default().insert(config.alt);
1146    }
1147    by_state.values().any(|alts| alts.len() == 1)
1148}
1149
1150#[cfg(test)]
1151mod tests {
1152    use super::*;
1153
1154    #[test]
1155    fn arena_interns_singletons_without_per_context_objects() {
1156        let mut arena = ContextArena::new();
1157        let first = arena.singleton(EMPTY_CONTEXT, 7);
1158        let second = arena.singleton(EMPTY_CONTEXT, 7);
1159
1160        assert_eq!(first, second);
1161        assert_eq!(arena.stats().singleton_contexts, 1);
1162        assert_eq!(arena.stats().interner_hits, 1);
1163    }
1164
1165    #[test]
1166    fn array_interner_verifies_payload_after_hash_collision() {
1167        let mut arena = ContextArena::new();
1168        let first_parent = arena.singleton(EMPTY_CONTEXT, 1);
1169        let second_parent = arena.singleton(EMPTY_CONTEXT, 2);
1170        let expected = [(first_parent, 10), (second_parent, 20)];
1171        let colliding = [(second_parent, 10), (first_parent, 20)];
1172        let cached_hash = prediction_context_array_hash(&arena, &expected);
1173        let start = u32::try_from(arena.array_parents.len()).expect("pool index fits u32");
1174        arena
1175            .array_parents
1176            .extend(colliding.iter().map(|(parent, _)| *parent));
1177        arena
1178            .array_return_states
1179            .extend(colliding.iter().map(|(_, return_state)| *return_state));
1180        let collision = arena.push_record(ContextRecord {
1181            tag: ContextTag::Array,
1182            cached_hash,
1183            parent_or_start: start,
1184            return_state_or_len: 2,
1185        });
1186
1187        let interned = arena.intern_entries(&expected);
1188
1189        assert_ne!(interned, collision);
1190        assert_eq!(arena.entry(interned, 0), Some(expected[0]));
1191        assert_eq!(arena.entry(interned, 1), Some(expected[1]));
1192    }
1193
1194    #[test]
1195    fn merge_with_empty_preserves_full_context_empty_path() {
1196        let mut arena = ContextArena::new();
1197        let mut workspace = PredictionWorkspace::default();
1198        let singleton = arena.singleton(EMPTY_CONTEXT, 42);
1199
1200        let merged = arena.merge(singleton, EMPTY_CONTEXT, false, &mut workspace);
1201
1202        assert_eq!(arena.len(merged), 2);
1203        assert_eq!(arena.return_state(merged, 0), Some(42));
1204        assert_eq!(arena.parent(merged, 0), Some(EMPTY_CONTEXT));
1205        assert_eq!(arena.return_state(merged, 1), Some(EMPTY_RETURN_STATE));
1206        assert!(arena.has_empty_path(merged));
1207    }
1208
1209    #[test]
1210    fn wildcard_merge_collapses_to_empty() {
1211        let mut arena = ContextArena::new();
1212        let mut workspace = PredictionWorkspace::default();
1213        let singleton = arena.singleton(EMPTY_CONTEXT, 42);
1214
1215        assert_eq!(
1216            arena.merge(singleton, EMPTY_CONTEXT, true, &mut workspace),
1217            EMPTY_CONTEXT
1218        );
1219    }
1220
1221    #[test]
1222    fn merge_is_order_independent() {
1223        let mut arena = ContextArena::new();
1224        let mut workspace = PredictionWorkspace::default();
1225        let left_parent = arena.singleton(EMPTY_CONTEXT, 100);
1226        let right_parent = arena.singleton(EMPTY_CONTEXT, 200);
1227        let left = arena.singleton(left_parent, 7);
1228        let right = arena.singleton(right_parent, 7);
1229
1230        let left_right = arena.merge(left, right, false, &mut workspace);
1231        workspace.reset();
1232        let right_left = arena.merge(right, left, false, &mut workspace);
1233
1234        assert_eq!(left_right, right_left);
1235        assert_eq!(arena.len(left_right), 2);
1236    }
1237
1238    #[test]
1239    fn import_remaps_contexts_into_destination_arena() {
1240        let mut source = ContextArena::new();
1241        let parent = source.singleton(EMPTY_CONTEXT, 3);
1242        let child = source.singleton(parent, 9);
1243        let mut destination = ContextArena::new();
1244        let mut workspace = PredictionWorkspace::default();
1245
1246        let remap = destination.import_all(&source, &mut workspace);
1247        let imported = remap[usize::try_from(child.0).expect("context ID fits usize")];
1248
1249        assert_eq!(destination.return_state(imported, 0), Some(9));
1250        let imported_parent = destination.parent(imported, 0).expect("parent");
1251        assert_eq!(destination.return_state(imported_parent, 0), Some(3));
1252    }
1253
1254    #[test]
1255    fn config_set_merges_context_ids() {
1256        let mut arena = ContextArena::new();
1257        let mut workspace = PredictionWorkspace::default();
1258        let left = arena.singleton(EMPTY_CONTEXT, 1);
1259        let right = arena.singleton(EMPTY_CONTEXT, 2);
1260        let mut set = AtnConfigSet::new_full_context(true);
1261
1262        assert!(set.add(
1263            AtnConfig::new(1, 1, left, &arena),
1264            &mut arena,
1265            &mut workspace
1266        ));
1267        assert!(!set.add(
1268            AtnConfig::new(1, 1, right, &arena),
1269            &mut arena,
1270            &mut workspace
1271        ));
1272        assert_eq!(set.len(), 1);
1273        assert_eq!(arena.len(set.configs()[0].context), 2);
1274    }
1275
1276    #[test]
1277    fn workspace_drops_pathological_capacity() {
1278        let mut workspace = PredictionWorkspace::default();
1279        workspace
1280            .merge_cache
1281            .reserve(MAX_RETAINED_MERGE_CACHE_ENTRIES.saturating_mul(2));
1282        workspace
1283            .entries
1284            .reserve(MAX_RETAINED_CONTEXT_ENTRIES.saturating_mul(2));
1285        workspace.reset();
1286
1287        assert!(workspace.merge_cache.capacity() <= MAX_RETAINED_MERGE_CACHE_ENTRIES);
1288        assert!(workspace.entries.capacity() <= MAX_RETAINED_CONTEXT_ENTRIES);
1289    }
1290}
antlr4_runtime/prediction.rs

antlr4_runtime/
prediction.rs
