1use 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#[derive(Clone, Copy, Debug)]
16pub struct ReservationPolicy {
17 context: &'static str,
18 fix: &'static str,
19}
20
21impl ReservationPolicy {
22 #[must_use]
24 pub const fn new(context: &'static str, fix: &'static str) -> Self {
25 Self { context, fix }
26 }
27
28 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 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 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 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 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 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 pub fn clear_vec_slots<T>(slots: &mut [Vec<T>]) {
133 crate::output_slots::clear_vec_slots(slots);
134 }
135
136 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 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 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 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
226pub type StagingReservationFailureAdapter<E> = fn(&'static str, usize, String) -> E;
228
229pub 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
246pub 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
262pub 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
283pub 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
304pub 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
326pub 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 #[must_use]
360 pub fn new() -> Self {
361 Self::default()
362 }
363
364 pub fn clear(&mut self) {
366 self.seen.clear();
367 self.ordered_indices.clear();
368 }
369
370 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 pub fn insert_seen(&mut self, key: K) -> bool {
396 self.seen.insert(key)
397 }
398
399 pub fn push_index(&mut self, index: usize) {
401 self.ordered_indices.push(index);
402 }
403
404 pub fn ordered_indices_mut(&mut self) -> &mut Vec<usize> {
406 &mut self.ordered_indices
407 }
408
409 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 #[must_use]
427 pub fn ordered_indices(&self) -> &[usize] {
428 &self.ordered_indices
429 }
430
431 #[must_use]
433 pub fn seen_capacity(&self) -> usize {
434 self.seen.capacity()
435 }
436
437 #[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}