Skip to main content

vyre_driver/
reservation_policy.rs

1//! Backend-neutral reservation policy adapters.
2//!
3//! Concrete backends own their wording, but hot dispatch paths should share one
4//! reservation policy for Vec, SmallVec, hash collections, and output slots.
5
6use std::collections::hash_map::RandomState;
7use std::collections::{HashMap, HashSet};
8use std::hash::{BuildHasher, Hash};
9
10use smallvec::{Array, SmallVec};
11
12use crate::BackendError;
13
14/// Domain wording for a family of bounded reservations.
15#[derive(Clone, Copy, Debug)]
16pub struct ReservationPolicy {
17    context: &'static str,
18    fix: &'static str,
19}
20
21impl ReservationPolicy {
22    /// Create a reservation policy with a stable error context and fix.
23    #[must_use]
24    pub const fn new(context: &'static str, fix: &'static str) -> Self {
25        Self { context, fix }
26    }
27
28    /// Ensure a Vec reaches `target_capacity` without changing length.
29    ///
30    /// # Errors
31    ///
32    /// Returns [`BackendError`] when the Vec cannot reserve memory.
33    pub fn reserve_vec_to_capacity<T>(
34        self,
35        vec: &mut Vec<T>,
36        target_capacity: usize,
37        item: &'static str,
38    ) -> Result<(), BackendError> {
39        crate::allocation::reserve_vec_to_capacity(
40            vec,
41            target_capacity,
42            self.context,
43            item,
44            self.fix,
45        )
46    }
47
48    /// Allocate an empty Vec with `target_capacity` reserved.
49    ///
50    /// # Errors
51    ///
52    /// Returns [`BackendError`] when the Vec cannot reserve memory.
53    pub fn reserved_vec<T>(
54        self,
55        target_capacity: usize,
56        item: &'static str,
57    ) -> Result<Vec<T>, BackendError> {
58        let mut vec = Vec::new();
59        self.reserve_vec_to_capacity(&mut vec, target_capacity, item)?;
60        Ok(vec)
61    }
62
63    /// Reserve `additional` more Vec elements without changing length.
64    ///
65    /// # Errors
66    ///
67    /// Returns [`BackendError`] when the Vec cannot reserve memory.
68    pub fn reserve_vec_additional<T>(
69        self,
70        vec: &mut Vec<T>,
71        additional: usize,
72        item: &'static str,
73    ) -> Result<(), BackendError> {
74        crate::allocation::reserve_vec_additional(vec, additional, self.context, item, self.fix)
75    }
76
77    /// Reserve enough Vec storage for `target_len` elements without resizing.
78    ///
79    /// # Errors
80    ///
81    /// Returns [`BackendError`] when the Vec cannot reserve memory.
82    pub fn reserve_vec_exact_for_len<T>(
83        self,
84        vec: &mut Vec<T>,
85        target_len: usize,
86        item: &'static str,
87    ) -> Result<(), BackendError> {
88        crate::output_slots::reserve_vec_exact_for_len(
89            vec,
90            target_len,
91            self.context,
92            item,
93            self.fix,
94        )
95    }
96
97    /// Ensure a Vec of output slots has at least `slot_count` slots.
98    ///
99    /// # Errors
100    ///
101    /// Returns [`BackendError`] when the outer Vec cannot reserve memory.
102    pub fn ensure_vec_slots_at_least<T>(
103        self,
104        slots: &mut Vec<Vec<T>>,
105        slot_count: usize,
106        item: &'static str,
107    ) -> Result<(), BackendError> {
108        crate::output_slots::ensure_vec_slots_at_least(
109            slots,
110            slot_count,
111            self.context,
112            item,
113            self.fix,
114        )
115    }
116
117    /// Resize a Vec of output slots while preserving existing prefixes.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`BackendError`] when the outer Vec cannot reserve memory.
122    pub fn resize_vec_slots<T>(
123        self,
124        slots: &mut Vec<Vec<T>>,
125        slot_count: usize,
126        item: &'static str,
127    ) -> Result<(), BackendError> {
128        crate::output_slots::resize_vec_slots(slots, slot_count, self.context, item, self.fix)
129    }
130
131    /// Clear inner output buffers without changing slot count.
132    pub fn clear_vec_slots<T>(slots: &mut [Vec<T>]) {
133        crate::output_slots::clear_vec_slots(slots);
134    }
135
136    /// Ensure a SmallVec reaches `target_capacity` without changing length.
137    ///
138    /// # Errors
139    ///
140    /// Returns [`BackendError`] when the SmallVec cannot reserve memory.
141    pub fn reserve_smallvec_to_capacity<A>(
142        self,
143        vec: &mut SmallVec<A>,
144        target_capacity: usize,
145        item: &'static str,
146    ) -> Result<(), BackendError>
147    where
148        A: Array,
149    {
150        let additional = target_capacity.saturating_sub(vec.len());
151        self.reserve_smallvec_additional(vec, additional, item)
152    }
153
154    /// Reserve `additional` more SmallVec elements without changing length.
155    ///
156    /// # Errors
157    ///
158    /// Returns [`BackendError`] when the SmallVec cannot reserve memory.
159    pub fn reserve_smallvec_additional<A>(
160        self,
161        vec: &mut SmallVec<A>,
162        additional: usize,
163        item: &'static str,
164    ) -> Result<(), BackendError>
165    where
166        A: Array,
167    {
168        crate::allocation::reserve_smallvec_additional(
169            vec,
170            additional,
171            self.context,
172            item,
173            self.fix,
174        )
175    }
176
177    /// Ensure a HashSet reaches `target_capacity` without changing length.
178    ///
179    /// # Errors
180    ///
181    /// Returns [`BackendError`] when the HashSet cannot reserve memory.
182    pub fn reserve_hash_set_to_capacity<T, S>(
183        self,
184        set: &mut HashSet<T, S>,
185        target_capacity: usize,
186        item: &'static str,
187    ) -> Result<(), BackendError>
188    where
189        T: Eq + Hash,
190        S: BuildHasher,
191    {
192        crate::allocation::reserve_hash_set_to_capacity(
193            set,
194            target_capacity,
195            self.context,
196            item,
197            self.fix,
198        )
199    }
200
201    /// Ensure a HashMap reaches `target_capacity` without changing length.
202    ///
203    /// # Errors
204    ///
205    /// Returns [`BackendError`] when the HashMap cannot reserve memory.
206    pub fn reserve_hash_map_to_capacity<K, V, S>(
207        self,
208        map: &mut HashMap<K, V, S>,
209        target_capacity: usize,
210        item: &'static str,
211    ) -> Result<(), BackendError>
212    where
213        K: Eq + Hash,
214        S: BuildHasher,
215    {
216        crate::allocation::reserve_hash_map_to_capacity(
217            map,
218            target_capacity,
219            self.context,
220            item,
221            self.fix,
222        )
223    }
224}
225
226/// Convert a shared reservation failure into a caller-domain error.
227pub type StagingReservationFailureAdapter<E> = fn(&'static str, usize, String) -> E;
228
229/// Reserve Vec capacity and map failures into a caller-domain typed error.
230///
231/// # Errors
232///
233/// Returns `E` when the Vec cannot reserve memory.
234pub fn reserve_typed_vec_to_capacity<T, E>(
235    policy: ReservationPolicy,
236    vec: &mut Vec<T>,
237    target_capacity: usize,
238    item: &'static str,
239    failure: StagingReservationFailureAdapter<E>,
240) -> Result<(), E> {
241    policy
242        .reserve_vec_to_capacity(vec, target_capacity, item)
243        .map_err(|error| failure(item, target_capacity, error.to_string()))
244}
245
246/// Allocate an empty Vec with reserved capacity and typed failure mapping.
247///
248/// # Errors
249///
250/// Returns `E` when the Vec cannot reserve memory.
251pub fn reserved_typed_vec<T, E>(
252    policy: ReservationPolicy,
253    target_capacity: usize,
254    item: &'static str,
255    failure: StagingReservationFailureAdapter<E>,
256) -> Result<Vec<T>, E> {
257    let mut vec = Vec::new();
258    reserve_typed_vec_to_capacity(policy, &mut vec, target_capacity, item, failure)?;
259    Ok(vec)
260}
261
262/// Reserve HashSet capacity and map failures into a caller-domain typed error.
263///
264/// # Errors
265///
266/// Returns `E` when the HashSet cannot reserve memory.
267pub fn reserve_typed_hash_set_to_capacity<T, S, E>(
268    policy: ReservationPolicy,
269    set: &mut HashSet<T, S>,
270    target_capacity: usize,
271    item: &'static str,
272    failure: StagingReservationFailureAdapter<E>,
273) -> Result<(), E>
274where
275    T: Eq + Hash,
276    S: BuildHasher,
277{
278    policy
279        .reserve_hash_set_to_capacity(set, target_capacity, item)
280        .map_err(|error| failure(item, target_capacity, error.to_string()))
281}
282
283/// Reserve HashMap capacity and map failures into a caller-domain typed error.
284///
285/// # Errors
286///
287/// Returns `E` when the HashMap cannot reserve memory.
288pub fn reserve_typed_hash_map_to_capacity<K, V, S, E>(
289    policy: ReservationPolicy,
290    map: &mut HashMap<K, V, S>,
291    target_capacity: usize,
292    item: &'static str,
293    failure: StagingReservationFailureAdapter<E>,
294) -> Result<(), E>
295where
296    K: Eq + Hash,
297    S: BuildHasher,
298{
299    policy
300        .reserve_hash_map_to_capacity(map, target_capacity, item)
301        .map_err(|error| failure(item, target_capacity, error.to_string()))
302}
303
304/// Reserve paired duplicate-detection and stable-order buffers with one typed failure adapter.
305///
306/// # Errors
307///
308/// Returns `E` when either staging collection cannot reserve memory.
309pub fn reserve_typed_hash_set_and_vec_to_capacity<K, V, S, E>(
310    policy: ReservationPolicy,
311    set: &mut HashSet<K, S>,
312    vec: &mut Vec<V>,
313    target_capacity: usize,
314    set_item: &'static str,
315    vec_item: &'static str,
316    failure: StagingReservationFailureAdapter<E>,
317) -> Result<(), E>
318where
319    K: Eq + Hash,
320    S: BuildHasher,
321{
322    reserve_typed_hash_set_to_capacity(policy, set, target_capacity, set_item, failure)?;
323    reserve_typed_vec_to_capacity(policy, vec, target_capacity, vec_item, failure)
324}
325
326/// Reusable duplicate-detection plus stable-order scratch for planner hot paths.
327pub struct ReusableIndexScratch<K, S = RandomState> {
328    seen: HashSet<K, S>,
329    ordered_indices: Vec<usize>,
330}
331
332impl<K, S> std::fmt::Debug for ReusableIndexScratch<K, S> {
333    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
334        f.debug_struct("ReusableIndexScratch")
335            .field("seen_capacity", &self.seen.capacity())
336            .field("ordered_index_capacity", &self.ordered_indices.capacity())
337            .finish()
338    }
339}
340
341impl<K, S> Default for ReusableIndexScratch<K, S>
342where
343    S: Default,
344{
345    fn default() -> Self {
346        Self {
347            seen: HashSet::with_hasher(S::default()),
348            ordered_indices: Vec::new(),
349        }
350    }
351}
352
353impl<K, S> ReusableIndexScratch<K, S>
354where
355    K: Eq + Hash,
356    S: BuildHasher + Default,
357{
358    /// Create empty reusable index scratch.
359    #[must_use]
360    pub fn new() -> Self {
361        Self::default()
362    }
363
364    /// Clear retained scratch entries without releasing retained capacity.
365    pub fn clear(&mut self) {
366        self.seen.clear();
367        self.ordered_indices.clear();
368    }
369
370    /// Reserve duplicate-detection and ordering scratch to the requested capacity.
371    ///
372    /// # Errors
373    ///
374    /// Returns `E` when either retained scratch collection cannot reserve memory.
375    pub fn try_reserve_with<E>(
376        &mut self,
377        policy: ReservationPolicy,
378        capacity: usize,
379        seen_item: &'static str,
380        ordered_indices_item: &'static str,
381        failure: StagingReservationFailureAdapter<E>,
382    ) -> Result<(), E> {
383        reserve_typed_hash_set_and_vec_to_capacity(
384            policy,
385            &mut self.seen,
386            &mut self.ordered_indices,
387            capacity,
388            seen_item,
389            ordered_indices_item,
390            failure,
391        )
392    }
393
394    /// Insert a duplicate-detection key.
395    pub fn insert_seen(&mut self, key: K) -> bool {
396        self.seen.insert(key)
397    }
398
399    /// Append an input index to the reusable ordering buffer.
400    pub fn push_index(&mut self, index: usize) {
401        self.ordered_indices.push(index);
402    }
403
404    /// Mutable ordering buffer for planner-specific sort keys.
405    pub fn ordered_indices_mut(&mut self) -> &mut Vec<usize> {
406        &mut self.ordered_indices
407    }
408
409    /// Sort ordered indices only when the current key order is not already monotonic.
410    pub fn sort_indices_unstable_by_key_if_needed<Key, F>(&mut self, mut key: F)
411    where
412        Key: Ord,
413        F: FnMut(usize) -> Key,
414    {
415        let needs_sort = self
416            .ordered_indices
417            .windows(2)
418            .any(|pair| key(pair[0]) > key(pair[1]));
419        if needs_sort {
420            self.ordered_indices
421                .sort_unstable_by_key(|&index| key(index));
422        }
423    }
424
425    /// Ordered input indices after planner-specific sorting.
426    #[must_use]
427    pub fn ordered_indices(&self) -> &[usize] {
428        &self.ordered_indices
429    }
430
431    /// Retained duplicate-detection capacity.
432    #[must_use]
433    pub fn seen_capacity(&self) -> usize {
434        self.seen.capacity()
435    }
436
437    /// Retained ordering capacity.
438    #[must_use]
439    pub fn ordered_index_capacity(&self) -> usize {
440        self.ordered_indices.capacity()
441    }
442}
443
444#[cfg(test)]
445mod tests {
446    use std::cell::Cell;
447    use std::collections::{HashMap, HashSet};
448
449    use smallvec::SmallVec;
450
451    use super::{
452        reserve_typed_hash_map_to_capacity, reserve_typed_hash_set_and_vec_to_capacity,
453        reserve_typed_hash_set_to_capacity, reserve_typed_vec_to_capacity, reserved_typed_vec,
454        ReservationPolicy, ReusableIndexScratch,
455    };
456
457    const TEST_POLICY: ReservationPolicy =
458        ReservationPolicy::new("generated staging reserve", "split generated dispatch");
459
460    #[derive(Debug, Eq, PartialEq)]
461    enum TypedReserveError {
462        Reserve {
463            field: &'static str,
464            requested: usize,
465            message: String,
466        },
467    }
468
469    fn typed_reserve_error(
470        field: &'static str,
471        requested: usize,
472        message: String,
473    ) -> TypedReserveError {
474        TypedReserveError::Reserve {
475            field,
476            requested,
477            message,
478        }
479    }
480
481    #[test]
482    fn policy_reserves_vec_smallvec_and_hash_collections_to_target_capacity() {
483        let mut vec = Vec::<u8>::with_capacity(4);
484        let mut small = SmallVec::<[u8; 2]>::new();
485        let mut map = HashMap::<u32, u32>::with_capacity(4);
486        let mut set = HashSet::<u32>::with_capacity(4);
487
488        vec.extend_from_slice(&[1, 2, 3, 4]);
489        small.extend_from_slice(&[1, 2, 3, 4]);
490        for value in 0..4 {
491            map.insert(value, value);
492            set.insert(value);
493        }
494        vec.clear();
495        small.clear();
496        map.clear();
497        set.clear();
498
499        TEST_POLICY
500            .reserve_vec_to_capacity(&mut vec, 32, "byte")
501            .expect("Fix: Vec target reservation should grow");
502        TEST_POLICY
503            .reserve_smallvec_to_capacity(&mut small, 32, "byte")
504            .expect("Fix: SmallVec target reservation should grow");
505        TEST_POLICY
506            .reserve_hash_map_to_capacity(&mut map, 32, "entry")
507            .expect("Fix: HashMap target reservation should grow");
508        TEST_POLICY
509            .reserve_hash_set_to_capacity(&mut set, 32, "entry")
510            .expect("Fix: HashSet target reservation should grow");
511
512        assert!(vec.capacity() >= 32);
513        assert!(small.capacity() >= 32);
514        assert!(map.capacity() >= 32);
515        assert!(set.capacity() >= 32);
516        assert!(vec.is_empty());
517        assert!(small.is_empty());
518        assert!(map.is_empty());
519        assert!(set.is_empty());
520    }
521
522    #[test]
523    fn policy_manages_output_slot_vectors_without_dropping_live_prefixes() {
524        let mut slots = vec![vec![1_u8], vec![2, 3]];
525
526        TEST_POLICY
527            .ensure_vec_slots_at_least(&mut slots, 4, "slot")
528            .expect("Fix: slot reservation should grow");
529        assert_eq!(slots.len(), 4);
530        assert_eq!(slots[0], vec![1]);
531        assert_eq!(slots[1], vec![2, 3]);
532
533        TEST_POLICY
534            .resize_vec_slots(&mut slots, 1, "slot")
535            .expect("Fix: slot resize should truncate without allocation");
536        assert_eq!(slots, vec![vec![1]]);
537
538        ReservationPolicy::clear_vec_slots(&mut slots);
539        assert_eq!(slots, vec![Vec::<u8>::new()]);
540    }
541
542    #[test]
543    fn typed_policy_reservations_share_vec_set_and_map_growth() {
544        let mut vec = Vec::<u8>::new();
545        let mut set = HashSet::<u32>::new();
546        let mut map = HashMap::<u32, u32>::new();
547
548        reserve_typed_vec_to_capacity(TEST_POLICY, &mut vec, 32, "typed byte", typed_reserve_error)
549            .expect("Fix: typed Vec reservation should grow");
550        reserve_typed_hash_set_to_capacity(
551            TEST_POLICY,
552            &mut set,
553            32,
554            "typed set entry",
555            typed_reserve_error,
556        )
557        .expect("Fix: typed HashSet reservation should grow");
558        reserve_typed_hash_map_to_capacity(
559            TEST_POLICY,
560            &mut map,
561            32,
562            "typed map entry",
563            typed_reserve_error,
564        )
565        .expect("Fix: typed HashMap reservation should grow");
566        reserve_typed_hash_set_and_vec_to_capacity(
567            TEST_POLICY,
568            &mut set,
569            &mut vec,
570            64,
571            "paired set entry",
572            "paired byte",
573            typed_reserve_error,
574        )
575        .expect("Fix: paired typed reservations should share one adapter");
576        let reserved =
577            reserved_typed_vec::<u16, _>(TEST_POLICY, 16, "typed word", typed_reserve_error)
578                .expect("Fix: typed Vec allocation should reserve");
579
580        assert!(vec.capacity() >= 64);
581        assert!(set.capacity() >= 64);
582        assert!(map.capacity() >= 32);
583        assert!(reserved.capacity() >= 16);
584        assert!(reserved.is_empty());
585    }
586
587    #[test]
588    fn typed_policy_reservation_reports_domain_failure_on_overflow() {
589        let mut bytes = Vec::<u8>::new();
590        let err = reserve_typed_vec_to_capacity(
591            TEST_POLICY,
592            &mut bytes,
593            usize::MAX,
594            "oversized typed byte",
595            typed_reserve_error,
596        )
597        .expect_err("oversized typed reservation should fail");
598
599        match err {
600            TypedReserveError::Reserve {
601                field,
602                requested,
603                message,
604            } => {
605                assert_eq!(field, "oversized typed byte");
606                assert_eq!(requested, usize::MAX);
607                assert!(message.contains("oversized typed byte"));
608                assert!(message.contains("Fix:"));
609            }
610        }
611    }
612
613    #[test]
614    fn reusable_index_scratch_preserves_capacity_and_orders_only_when_needed() {
615        let mut scratch = ReusableIndexScratch::<u32>::new();
616
617        scratch
618            .try_reserve_with(
619                TEST_POLICY,
620                64,
621                "scratch seen",
622                "scratch ordered",
623                typed_reserve_error,
624            )
625            .expect("Fix: reusable scratch should reserve through shared policy");
626        assert!(scratch.insert_seen(7));
627        assert!(!scratch.insert_seen(7));
628        scratch.push_index(0);
629        scratch.push_index(1);
630        scratch.push_index(2);
631
632        let key_calls = Cell::new(0);
633        scratch.sort_indices_unstable_by_key_if_needed(|index| {
634            key_calls.set(key_calls.get() + 1);
635            [10_u32, 20, 30][index]
636        });
637        assert_eq!(scratch.ordered_indices(), &[0, 1, 2]);
638        assert_eq!(
639            key_calls.get(),
640            4,
641            "Fix: monotonic planner indices must skip sort_unstable_by_key."
642        );
643
644        let seen_capacity = scratch.seen_capacity();
645        let ordered_capacity = scratch.ordered_index_capacity();
646        scratch.clear();
647        scratch.push_index(2);
648        scratch.push_index(0);
649        scratch.push_index(1);
650        scratch.sort_indices_unstable_by_key_if_needed(|index| [10_u32, 20, 30][index]);
651
652        assert_eq!(scratch.ordered_indices(), &[0, 1, 2]);
653        assert!(scratch.seen_capacity() >= seen_capacity);
654        assert!(scratch.ordered_index_capacity() >= ordered_capacity);
655    }
656
657    #[test]
658    fn generated_reusable_index_scratch_matrix_keeps_exact_order_contract() {
659        for len in 0..=96 {
660            let mut scratch = ReusableIndexScratch::<usize>::new();
661            scratch
662                .try_reserve_with(
663                    TEST_POLICY,
664                    len,
665                    "generated seen",
666                    "generated ordered",
667                    typed_reserve_error,
668                )
669                .expect("Fix: generated scratch reservation should succeed");
670
671            for index in (0..len).rev() {
672                assert!(scratch.insert_seen(index));
673                scratch.push_index(index);
674            }
675            scratch.sort_indices_unstable_by_key_if_needed(|index| index);
676
677            assert_eq!(scratch.ordered_indices().len(), len);
678            for (expected, actual) in scratch.ordered_indices().iter().copied().enumerate() {
679                assert_eq!(actual, expected);
680            }
681            let seen_capacity = scratch.seen_capacity();
682            let ordered_capacity = scratch.ordered_index_capacity();
683            scratch.clear();
684            scratch
685                .try_reserve_with(
686                    TEST_POLICY,
687                    len / 2,
688                    "generated seen shrink",
689                    "generated ordered shrink",
690                    typed_reserve_error,
691                )
692                .expect("Fix: generated scratch reuse should keep retained storage");
693            assert!(scratch.seen_capacity() >= seen_capacity);
694            assert!(scratch.ordered_index_capacity() >= ordered_capacity);
695            assert!(scratch.ordered_indices().is_empty());
696        }
697    }
698}