1 use alloc::vec::Vec;
2 use arbitrary::{Arbitrary, Unstructured};
3 
4 /// The control plane of chaos mode.
5 /// Please see the [crate-level documentation](crate).
6 #[derive(Debug, Clone, Default)]
7 pub 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 
16 impl Arbitrary<'_> for ControlPlane {
arbitrary<'a>(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self>17     fn arbitrary<'a>(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
18         Ok(Self::new(u.arbitrary()?))
19     }
20 }
21 
22 impl ControlPlane {
new(data: Vec<u8>) -> Self23     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`.
set_fuel(&mut self, fuel: u8)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).
consume_fuel(&mut self) -> bool40     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`.
get_decision(&mut self) -> bool57     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`.
get_arbitrary<T: for<'a> Arbitrary<'a> + Default>(&mut self) -> T66     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`.
shuffle<T>(&mut self, slice: &mut [T])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`.
shuffled<T>(&mut self, iter: impl Iterator<Item = T>) -> impl Iterator<Item = T>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 }
127