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}