gpui_component/plot/scale/
ordinal.rs

1// @reference: https://d3js.org/d3-scale/ordinal
2
3#[derive(Clone)]
4pub struct ScaleOrdinal<D, R> {
5    domain: Vec<D>,
6    range: Vec<R>,
7    unknown: Option<R>,
8}
9
10impl<D, R> Default for ScaleOrdinal<D, R> {
11    fn default() -> Self {
12        Self {
13            domain: Vec::new(),
14            range: Vec::new(),
15            unknown: None,
16        }
17    }
18}
19
20impl<D, R> ScaleOrdinal<D, R> {
21    pub fn new(domain: Vec<D>, range: Vec<R>) -> Self {
22        Self {
23            domain,
24            range,
25            unknown: None,
26        }
27    }
28
29    /// Set the domain to the specified array of values.
30    pub fn domain(mut self, domain: Vec<D>) -> Self {
31        self.domain = domain;
32        self
33    }
34
35    /// Set the range of the ordinal scale to the specified array of values.
36    pub fn range(mut self, range: Vec<R>) -> Self {
37        self.range = range;
38        self
39    }
40
41    /// Set the output value of the scale for unknown input values and returns this scale.
42    pub fn unknown(mut self, unknown: R) -> Self {
43        self.unknown = Some(unknown);
44        self
45    }
46}
47
48impl<D, R> ScaleOrdinal<D, R>
49where
50    D: PartialEq,
51    R: Clone,
52{
53    /// Given a value in the input domain, returns the corresponding value in the output range.
54    pub fn map(&self, value: &D) -> Option<R> {
55        if let Some(index) = self.domain.iter().position(|v| v == value) {
56            if self.range.is_empty() {
57                None
58            } else {
59                Some(self.range[index % self.range.len()].clone())
60            }
61        } else {
62            self.unknown.clone()
63        }
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn test_scale_ordinal() {
73        let scale = ScaleOrdinal::new(vec!["a", "b", "c", "d", "e"], vec![10, 20, 30]);
74
75        assert_eq!(scale.map(&"a"), Some(10));
76        assert_eq!(scale.map(&"b"), Some(20));
77        assert_eq!(scale.map(&"c"), Some(30));
78        assert_eq!(scale.map(&"d"), Some(10));
79        assert_eq!(scale.map(&"e"), Some(20));
80        assert_eq!(scale.map(&"f"), None);
81    }
82
83    #[test]
84    fn test_scale_ordinal_unknown() {
85        let scale = ScaleOrdinal::new(vec!["a", "b", "c"], vec![10, 20, 30]).unknown(0);
86
87        assert_eq!(scale.map(&"a"), Some(10));
88        assert_eq!(scale.map(&"f"), Some(0));
89    }
90
91    #[test]
92    fn test_scale_ordinal_colors() {
93        let keys = vec!["a", "b", "c", "d"];
94        let colors = vec!["#1f77b4", "#ff7f0e", "#2ca02c"];
95
96        let scale = ScaleOrdinal::new(keys, colors);
97
98        assert_eq!(scale.map(&"a"), Some("#1f77b4"));
99        assert_eq!(scale.map(&"b"), Some("#ff7f0e"));
100        assert_eq!(scale.map(&"c"), Some("#2ca02c"));
101        // Should cycle back to the first color
102        assert_eq!(scale.map(&"d"), Some("#1f77b4"));
103    }
104}