melvin_ob/scheduling/atomic_decision_cube.rs
1use super::atomic_decision::AtomicDecision;
2
3/// A flattened 3D data structure to manage atomic decisions for multiple dimensions with good cache performance.
4pub struct AtomicDecisionCube {
5 /// Length of the time dimension.
6 dt_len: usize,
7 /// Length of the energy dimension.
8 e_len: usize,
9 /// Length of the state dimension.
10 s_len: usize,
11 /// Array of atomic decisions.
12 decisions: Box<[AtomicDecision]>,
13}
14
15impl AtomicDecisionCube {
16 /// Creates a new [`AtomicDecisionCube`] with the specified dimensions and initializes all decisions to `StayInCharge`.
17 ///
18 /// # Arguments
19 /// * `dt_len` - The length of the time dimension.
20 /// * `e_len` - The length of the energy dimension.
21 /// * `s_len` - The length of the state dimension.
22 ///
23 /// # Returns
24 /// A new instance of [`AtomicDecisionCube`].
25 pub fn new(dt_len: usize, e_len: usize, s_len: usize) -> Self {
26 Self {
27 dt_len,
28 e_len,
29 s_len,
30 decisions: vec![AtomicDecision::StayInCharge; dt_len * e_len * s_len]
31 .into_boxed_slice(),
32 }
33 }
34
35 /// Retrieves the atomic decision at the specified indices.
36 ///
37 /// # Arguments
38 /// * `dt` - The index along the time dimension.
39 /// * `e` - The index along the energy dimension.
40 /// * `s` - The index along the state dimension.
41 ///
42 /// # Returns
43 /// The [`AtomicDecision`] at the specified indices.
44 pub fn get(&self, dt: usize, e: usize, s: usize) -> AtomicDecision {
45 self.decisions[dt * self.e_len * self.s_len + e * self.s_len + s]
46 }
47
48 /// Sets the atomic decision at the specified indices.
49 ///
50 /// # Arguments
51 /// * `dt` - The index along the time dimension.
52 /// * `e` - The index along the energy dimension.
53 /// * `s` - The index along the state dimension.
54 /// * `decision` - The [`AtomicDecision`] to set at the specified indices.
55 pub fn set(&mut self, dt: usize, e: usize, s: usize, decision: AtomicDecision) {
56 self.decisions[dt * self.e_len * self.s_len + e * self.s_len + s] = decision;
57 }
58
59 /// Returns the length of the time dimension.
60 ///
61 /// # Returns
62 /// The length of the time dimension.
63 pub fn dt_len(&self) -> usize { self.dt_len }
64
65 /// Returns the length of the energy dimension.
66 ///
67 /// # Returns
68 /// The length of the energy dimension.
69 pub fn e_len(&self) -> usize { self.e_len }
70
71 /// Returns the length of the state dimension.
72 ///
73 /// # Returns
74 /// The length of the state dimension.
75 pub fn s_len(&self) -> usize { self.s_len }
76}