Skip to main content

cranelift_control/
chaos.rs

1use alloc::vec::Vec;
2use arbitrary::{Arbitrary, Unstructured};
3
4/// The control plane of chaos mode.
5/// Please see the [crate-level documentation](crate).
6#[derive(Debug, Clone, Default)]
7pub struct ControlPlane {
8    data: Vec<u8>,
9    fuel: Option<u8>,
10    /// This is used as a little optimization to avoid additional heap
11    /// allocations when using `Unstructured` internally. See the source of
12    /// [`ControlPlane::shuffle`] for an example.
13    tmp: Vec<u8>,
14}
15
16impl Arbitrary<'_> for ControlPlane {
17    fn arbitrary<'a>(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
18        Ok(Self::new(u.arbitrary()?))
19    }
20}
21
22impl ControlPlane {
23    fn new(data: Vec<u8>) -> Self {
24        Self {
25            data,
26            fuel: None,
27            tmp: Vec::new(),
28        }
29    }
30
31    /// Set the [fuel limit](crate#fuel-limit). Zero is interpreted as the
32    /// fuel limit being deactivated, consistent with the cranelift setting
33    /// `control_plane_fuel`.
34    pub fn set_fuel(&mut self, fuel: u8) {
35        self.fuel = (fuel != 0).then_some(fuel)
36    }
37
38    /// Tries to consume fuel, returning `true` if successful (or if
39    /// fuel-limiting is disabled).
40    fn consume_fuel(&mut self) -> bool {
41        match self.fuel {
42            None => true,               // fuel deactivated
43            Some(f) if f == 0 => false, // no more fuel
44            Some(ref mut f) => {
45                *f -= 1;
46                true
47            }
48        }
49    }
50
51    /// Returns a pseudo-random boolean if the control plane was constructed
52    /// with `arbitrary`.
53    ///
54    /// The default value `false` will always be returned if the
55    /// pseudo-random data is exhausted or the control plane was constructed
56    /// with `default`.
57    pub fn get_decision(&mut self) -> bool {
58        self.consume_fuel() && self.data.pop().unwrap_or_default() & 1 == 1
59    }
60
61    /// Returns an arbitrary value if the control plane was constructed with
62    /// `arbitrary`.
63    ///
64    /// The default value will always be returned if the pseudo-random data is
65    /// exhausted or the control plane was constructed with `default`.
66    pub fn get_arbitrary<T: for<'a> Arbitrary<'a> + Default>(&mut self) -> T {
67        if !self.consume_fuel() || self.data.is_empty() {
68            return T::default();
69        }
70        let mut u = Unstructured::new(&self.data);
71        let res = u.arbitrary().unwrap_or_default();
72
73        // take remaining bytes
74        let rest = u.take_rest();
75        self.tmp.resize(rest.len(), 0); // allocates once per control plane
76        self.tmp.copy_from_slice(rest);
77        core::mem::swap(&mut self.data, &mut self.tmp);
78
79        res
80    }
81
82    /// Shuffles the items in the slice into a pseudo-random permutation if
83    /// the control plane was constructed with `arbitrary`.
84    ///
85    /// The default operation, to leave the slice unchanged, will always be
86    /// performed if the pseudo-random data is exhausted or the control
87    /// plane was constructed with `default`.
88    pub fn shuffle<T>(&mut self, slice: &mut [T]) {
89        if !self.consume_fuel() || self.data.is_empty() {
90            return;
91        }
92        let mut u = Unstructured::new(&self.data);
93
94        // adapted from:
95        // https://docs.rs/arbitrary/1.3.0/arbitrary/struct.Unstructured.html#examples-1
96        let mut to_permute = &mut slice[..];
97
98        while to_permute.len() > 1 {
99            if let Ok(idx) = u.choose_index(to_permute.len()) {
100                to_permute.swap(0, idx);
101                to_permute = &mut to_permute[1..];
102            } else {
103                break;
104            }
105        }
106
107        // take remaining bytes
108        let rest = u.take_rest();
109        self.tmp.resize(rest.len(), 0); // allocates once per control plane
110        self.tmp.copy_from_slice(rest);
111        core::mem::swap(&mut self.data, &mut self.tmp);
112    }
113
114    /// Returns a new iterator over the same items as the input iterator in
115    /// a pseudo-random order if the control plane was constructed with
116    /// `arbitrary`.
117    ///
118    /// The default value, an iterator with an unchanged order, will always
119    /// be returned if the pseudo-random data is exhausted or the control
120    /// plane was constructed with `default`.
121    pub fn shuffled<T>(&mut self, iter: impl Iterator<Item = T>) -> impl Iterator<Item = T> {
122        let mut slice: Vec<_> = iter.collect();
123        self.shuffle(&mut slice);
124        slice.into_iter()
125    }
126}