Skip to main content

muse2/asset/
capacity.rs

1//! Represents the capacity of an asset
2use crate::units::{Capacity, Dimensionless};
3use std::cmp::Ordering;
4use std::ops::{Add, Sub};
5
6/// Capacity of an asset, which may be continuous or a discrete number of indivisible units
7#[derive(Clone, PartialEq, Copy, Debug)]
8pub enum AssetCapacity {
9    /// Continuous capacity
10    Continuous(Capacity),
11    /// Discrete capacity represented by a number of indivisible units
12    /// Stores: (number of units, unit size)
13    Discrete(u32, Capacity),
14}
15
16impl AssetCapacity {
17    /// Return the smaller of `self` or `other`.
18    ///
19    /// # Panics
20    ///
21    /// Panics if the comparison is not meaningful. This happens if either `AssetCapacity` contains
22    /// a NaN value, one is discrete and the other continuous or if both are discrete and the unit
23    /// size differs.
24    pub fn min(self, other: AssetCapacity) -> AssetCapacity {
25        match self.partial_cmp(&other) {
26            None => panic!("Comparing invalid AssetCapacity values ({self:?} and {other:?})"),
27            Some(Ordering::Greater) => other,
28            _ => self,
29        }
30    }
31}
32
33impl Add for AssetCapacity {
34    type Output = Self;
35
36    // Add two AssetCapacity values together
37    fn add(self, rhs: AssetCapacity) -> Self {
38        match (self, rhs) {
39            (AssetCapacity::Continuous(cap1), AssetCapacity::Continuous(cap2)) => {
40                AssetCapacity::Continuous(cap1 + cap2)
41            }
42            (AssetCapacity::Discrete(units1, size1), AssetCapacity::Discrete(units2, size2)) => {
43                Self::check_same_unit_size(size1, size2);
44                AssetCapacity::Discrete(units1 + units2, size1)
45            }
46            _ => panic!("Cannot add different types of AssetCapacity ({self:?} and {rhs:?})"),
47        }
48    }
49}
50
51impl Sub for AssetCapacity {
52    type Output = Self;
53
54    // Subtract rhs from self, ensuring that the result is non-negative
55    fn sub(self, rhs: AssetCapacity) -> Self {
56        match (self, rhs) {
57            (AssetCapacity::Continuous(cap1), AssetCapacity::Continuous(cap2)) => {
58                AssetCapacity::Continuous((cap1 - cap2).max(Capacity(0.0)))
59            }
60            (AssetCapacity::Discrete(units1, size1), AssetCapacity::Discrete(units2, size2)) => {
61                Self::check_same_unit_size(size1, size2);
62                AssetCapacity::Discrete(units1 - units2.min(units1), size1)
63            }
64            _ => panic!("Cannot subtract different types of AssetCapacity ({self:?} and {rhs:?})"),
65        }
66    }
67}
68
69impl PartialOrd for AssetCapacity {
70    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
71        match (self, other) {
72            (AssetCapacity::Continuous(a), AssetCapacity::Continuous(b)) => a.partial_cmp(b),
73            (AssetCapacity::Discrete(units1, size1), AssetCapacity::Discrete(units2, size2)) => {
74                // NB: Also returns `None` if either is NaN
75                (*size1 == *size2).then(|| units1.cmp(units2))
76            }
77            _ => None,
78        }
79    }
80}
81
82impl AssetCapacity {
83    /// Validates that two discrete capacities have the same unit size.
84    fn check_same_unit_size(size1: Capacity, size2: Capacity) {
85        assert_eq!(
86            size1, size2,
87            "Can't perform operation on capacities with different unit sizes ({size1} and {size2})",
88        );
89    }
90
91    /// Create an `AssetCapacity` from a total capacity and optional unit size
92    ///
93    /// If a unit size is provided, the capacity is represented as a discrete number of units,
94    /// calculated as the ceiling of (capacity / `unit_size`). If no unit size is provided, the
95    /// capacity is represented as continuous.
96    pub fn from_capacity(capacity: Capacity, unit_size: Option<Capacity>) -> Self {
97        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
98        match unit_size {
99            Some(size) => {
100                let num_units = (capacity / size).value().ceil() as u32;
101                AssetCapacity::Discrete(num_units, size)
102            }
103            None => AssetCapacity::Continuous(capacity),
104        }
105    }
106
107    /// Create an `AssetCapacity` from a total capacity and optional unit size
108    ///
109    /// If a unit size is provided, the capacity is represented as a discrete number of units,
110    /// calculated as the floor of (capacity / `unit_size`). If no unit size is provided, the
111    /// capacity is represented as continuous.
112    pub fn from_capacity_floor(capacity: Capacity, unit_size: Option<Capacity>) -> Self {
113        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
114        match unit_size {
115            Some(size) => {
116                let num_units = (capacity / size).value().floor() as u32;
117                AssetCapacity::Discrete(num_units, size)
118            }
119            None => AssetCapacity::Continuous(capacity),
120        }
121    }
122
123    /// Returns the total capacity represented by this `AssetCapacity`.
124    pub fn total_capacity(&self) -> Capacity {
125        match self {
126            AssetCapacity::Continuous(cap) => *cap,
127            AssetCapacity::Discrete(units, size) => *size * Dimensionless(*units as f64),
128        }
129    }
130
131    /// Returns the number of units if this is a discrete capacity, or `None` if continuous.
132    pub fn n_units(&self) -> Option<u32> {
133        match self {
134            AssetCapacity::Continuous(_) => None,
135            AssetCapacity::Discrete(units, _) => Some(*units),
136        }
137    }
138
139    /// Asserts that both capacities are the same type (both continuous or both discrete).
140    pub fn assert_same_type(&self, other: AssetCapacity) {
141        assert!(
142            matches!(self, AssetCapacity::Continuous(_))
143                == matches!(other, AssetCapacity::Continuous(_)),
144            "Cannot change capacity type"
145        );
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::units::Capacity;
153    use rstest::rstest;
154
155    #[rstest]
156    #[case::exact_multiple(Capacity(12.0), Some(Capacity(4.0)), Some(3), Capacity(12.0))]
157    #[case::rounded_up(Capacity(11.0), Some(Capacity(4.0)), Some(3), Capacity(12.0))]
158    #[case::unit_size_greater_than_capacity(
159        Capacity(3.0),
160        Some(Capacity(4.0)),
161        Some(1),
162        Capacity(4.0)
163    )]
164    #[case::continuous(Capacity(5.5), None, None, Capacity(5.5))]
165    fn from_capacity(
166        #[case] capacity: Capacity,
167        #[case] unit_size: Option<Capacity>,
168        #[case] expected_n: Option<u32>,
169        #[case] expected_total: Capacity,
170    ) {
171        let got = AssetCapacity::from_capacity(capacity, unit_size);
172        assert_eq!(got.n_units(), expected_n);
173        assert_eq!(got.total_capacity(), expected_total);
174    }
175
176    #[rstest]
177    #[case::exact_multiple(Capacity(12.0), Some(Capacity(4.0)), Some(3), Capacity(12.0))]
178    #[case::rounded_down(Capacity(11.0), Some(Capacity(4.0)), Some(2), Capacity(8.0))]
179    #[case::unit_size_greater_than_capacity(
180        Capacity(3.0),
181        Some(Capacity(4.0)),
182        Some(0),
183        Capacity(0.0)
184    )]
185    #[case::continuous(Capacity(5.5), None, None, Capacity(5.5))]
186    fn from_capacity_floor(
187        #[case] capacity: Capacity,
188        #[case] unit_size: Option<Capacity>,
189        #[case] expected_n: Option<u32>,
190        #[case] expected_total: Capacity,
191    ) {
192        let got = AssetCapacity::from_capacity_floor(capacity, unit_size);
193        assert_eq!(got.n_units(), expected_n);
194        assert_eq!(got.total_capacity(), expected_total);
195    }
196
197    #[rstest]
198    #[case::less(
199        AssetCapacity::Continuous(Capacity(4.0)),
200        AssetCapacity::Continuous(Capacity(6.0)),
201        Some(Ordering::Less)
202    )]
203    #[case::equal(
204        AssetCapacity::Continuous(Capacity(4.0)),
205        AssetCapacity::Continuous(Capacity(4.0)),
206        Some(Ordering::Equal)
207    )]
208    #[case::greater(
209        AssetCapacity::Continuous(Capacity(6.0)),
210        AssetCapacity::Continuous(Capacity(4.0)),
211        Some(Ordering::Greater)
212    )]
213    fn partial_cmp_continuous(
214        #[case] left: AssetCapacity,
215        #[case] right: AssetCapacity,
216        #[case] expected: Option<Ordering>,
217    ) {
218        assert_eq!(left.partial_cmp(&right), expected);
219        assert_eq!(left == right, expected == Some(Ordering::Equal));
220    }
221
222    #[rstest]
223    #[case::less(
224        AssetCapacity::Discrete(2, Capacity(3.0)),
225        AssetCapacity::Discrete(4, Capacity(3.0)),
226        Some(Ordering::Less)
227    )]
228    #[case::equal(
229        AssetCapacity::Discrete(4, Capacity(3.0)),
230        AssetCapacity::Discrete(4, Capacity(3.0)),
231        Some(Ordering::Equal)
232    )]
233    #[case::greater(
234        AssetCapacity::Discrete(5, Capacity(3.0)),
235        AssetCapacity::Discrete(4, Capacity(3.0)),
236        Some(Ordering::Greater)
237    )]
238    fn partial_cmp_discrete_with_matching_unit_size(
239        #[case] left: AssetCapacity,
240        #[case] right: AssetCapacity,
241        #[case] expected: Option<Ordering>,
242    ) {
243        assert_eq!(left.partial_cmp(&right), expected);
244        assert_eq!(left == right, expected == Some(Ordering::Equal));
245    }
246
247    #[rstest]
248    #[case::mixed_types(
249        AssetCapacity::Continuous(Capacity(4.0)),
250        AssetCapacity::Discrete(4, Capacity(1.0))
251    )]
252    #[case::different_unit_sizes(
253        AssetCapacity::Discrete(4, Capacity(1.0)),
254        AssetCapacity::Discrete(4, Capacity(2.0))
255    )]
256    #[case::nan_continuous(
257        AssetCapacity::Continuous(Capacity(f64::NAN)),
258        AssetCapacity::Continuous(Capacity(4.0))
259    )]
260    fn partial_cmp_returns_none_for_invalid_comparisons(
261        #[case] left: AssetCapacity,
262        #[case] right: AssetCapacity,
263    ) {
264        assert_eq!(left.partial_cmp(&right), None);
265        assert_ne!(left, right);
266    }
267
268    #[rstest]
269    #[case::continuous(
270        AssetCapacity::Continuous(Capacity(4.0)),
271        AssetCapacity::Continuous(Capacity(6.0)),
272        AssetCapacity::Continuous(Capacity(4.0))
273    )]
274    #[case::discrete(
275        AssetCapacity::Discrete(2, Capacity(3.0)),
276        AssetCapacity::Discrete(4, Capacity(3.0)),
277        AssetCapacity::Discrete(2, Capacity(3.0))
278    )]
279    fn min_returns_smaller_capacity(
280        #[case] left: AssetCapacity,
281        #[case] right: AssetCapacity,
282        #[case] expected: AssetCapacity,
283    ) {
284        assert_eq!(left.min(right), expected);
285    }
286
287    #[rstest]
288    #[case::mixed_types(
289        AssetCapacity::Continuous(Capacity(4.0)),
290        AssetCapacity::Discrete(4, Capacity(1.0))
291    )]
292    #[case::different_unit_sizes(
293        AssetCapacity::Discrete(4, Capacity(1.0)),
294        AssetCapacity::Discrete(4, Capacity(2.0))
295    )]
296    #[should_panic(expected = "Comparing invalid AssetCapacity values")]
297    fn min_panics_for_invalid_comparisons(
298        #[case] left: AssetCapacity,
299        #[case] right: AssetCapacity,
300    ) {
301        let _ = left.min(right);
302    }
303}