qmc/util/
allocator.rs

1#[cfg(feature = "serialize")]
2use serde::{Deserialize, Serialize};
3use std::collections::BinaryHeap;
4
5/// A factory which produces Ts.
6pub trait Factory<T> {
7    /// Get an instance of T
8    fn get_instance(&mut self) -> T;
9    /// Return an instance of T
10    fn return_instance(&mut self, t: T);
11}
12
13/// Reset an instance while preserving its util.
14pub trait Reset {
15    /// Reset the instance, keep the util location.
16    fn reset(&mut self);
17}
18
19impl<T> Reset for Vec<T> {
20    fn reset(&mut self) {
21        self.clear()
22    }
23}
24
25impl<T> Reset for BinaryHeap<T> {
26    fn reset(&mut self) {
27        self.clear()
28    }
29}
30
31#[derive(Debug, Clone)]
32#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
33pub(crate) struct Allocator<T: Default + Reset> {
34    // Only save length, that's really all we care about.
35    #[cfg_attr(feature = "serialize", serde(with = "numeric_serialize"))]
36    instances: Vec<T>,
37    gen_more: bool,
38}
39
40impl<T: Default + Reset> Allocator<T> {
41    pub(crate) fn new_with_max_in_flight(max_in_flight: usize) -> Self {
42        let mut instances = Vec::with_capacity(max_in_flight);
43        instances.resize_with(max_in_flight, T::default);
44        Self {
45            instances,
46            gen_more: false,
47        }
48    }
49}
50
51impl<T: Default + Reset> Factory<T> for Allocator<T> {
52    #[track_caller]
53    fn get_instance(&mut self) -> T {
54        match self.instances.pop() {
55            None => {
56                if self.gen_more {
57                    T::default()
58                } else {
59                    panic!("Out of instances.")
60                }
61            }
62            Some(t) => t,
63        }
64    }
65
66    fn return_instance(&mut self, mut t: T) {
67        t.reset();
68        self.instances.push(t)
69    }
70}
71
72#[derive(Debug)]
73pub(crate) struct StackTuplizer<A, B> {
74    a_vec: Vec<A>,
75    b_vec: Vec<B>,
76}
77
78impl<A, B> Extend<(A, B)> for StackTuplizer<A, B> {
79    fn extend<T: IntoIterator<Item = (A, B)>>(&mut self, iter: T) {
80        iter.into_iter().for_each(|(a, b)| {
81            self.a_vec.push(a);
82            self.b_vec.push(b);
83        })
84    }
85}
86
87impl<A, B> StackTuplizer<A, B> {
88    pub(crate) fn new<T: ?Sized + Factory<Vec<A>> + Factory<Vec<B>>>(t: &mut T) -> Self {
89        let a_vec: Vec<A> = t.get_instance();
90        let b_vec: Vec<B> = t.get_instance();
91        Self { a_vec, b_vec }
92    }
93
94    pub(crate) fn dissolve<T: ?Sized + Factory<Vec<A>> + Factory<Vec<B>>>(self, t: &mut T) {
95        t.return_instance(self.a_vec);
96        t.return_instance(self.b_vec);
97    }
98
99    pub(crate) fn push(&mut self, tup: (A, B)) {
100        let (a, b) = tup;
101        self.a_vec.push(a);
102        self.b_vec.push(b);
103    }
104
105    pub(crate) fn pop(&mut self) -> Option<(A, B)> {
106        self.a_vec.pop().zip(self.b_vec.pop())
107    }
108
109    pub(crate) fn get(&self, index: usize) -> Option<(&A, &B)> {
110        self.a_vec.get(index).zip(self.b_vec.get(index))
111    }
112
113    pub(crate) fn at(&self, index: usize) -> (&A, &B) {
114        (&self.a_vec[index], &self.b_vec[index])
115    }
116
117    pub(crate) fn set(&mut self, index: usize, tup: (A, B)) {
118        let (a, b) = tup;
119        self.a_vec[index] = a;
120        self.b_vec[index] = b;
121    }
122
123    pub(crate) fn resize_each<F, G>(&mut self, size: usize, f: F, g: G)
124    where
125        F: Fn() -> A,
126        G: Fn() -> B,
127    {
128        self.a_vec.resize_with(size, f);
129        self.b_vec.resize_with(size, g);
130    }
131
132    pub(crate) fn iter(&self) -> impl Iterator<Item = (&A, &B)> {
133        self.a_vec.iter().zip(self.b_vec.iter())
134    }
135
136    pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = (&mut A, &mut B)> {
137        self.a_vec.iter_mut().zip(self.b_vec.iter_mut())
138    }
139}
140
141#[cfg(feature = "serialize")]
142mod numeric_serialize {
143    use serde::{Deserialize, Deserializer, Serializer};
144
145    pub(crate) fn serialize<T, S>(t: &[T], s: S) -> Result<S::Ok, S::Error>
146    where
147        S: Serializer,
148        T: Default,
149    {
150        let n = t.len();
151        s.serialize_u64(n as u64)
152    }
153
154    pub(crate) fn deserialize<'de, T, D>(d: D) -> Result<Vec<T>, D::Error>
155    where
156        D: Deserializer<'de>,
157        T: Default,
158    {
159        let s: u64 = Deserialize::deserialize(d)?;
160        Ok((0..s).map(|_| T::default()).collect())
161    }
162}