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}