pugio_lib/
coloring.rs

1use serde::Deserialize;
2
3use crate::graph::Graph;
4
5/// RGB Color.
6pub use colorous::Color;
7
8/// Trait for providing coloring based on input values.
9pub trait Gradient {
10    /// The type of input used for determining color.
11    type Input;
12
13    /// Get the color corresponding to the given input, given the dark mode and inverse settings.
14    fn color(&self, input: Self::Input, dark_mode: bool, inverse: bool) -> Color;
15}
16
17/// Trait for providing values for coloring nodes.
18pub trait Values {
19    /// The type of context shared across all nodes.
20    type Context;
21    /// The type of value for each node.
22    type Value;
23    /// The type of output used for coloring.
24    type Output;
25
26    /// Get the shared context.
27    fn context(&self) -> Self::Context;
28    /// Get the value for the node at the given index.
29    ///
30    /// # Panics
31    /// May panic if the node index does not exist.
32    fn value(&self, index: usize) -> Self::Value;
33    /// Get the output for the node at the given index.
34    ///
35    /// # Panics
36    /// May panic if the node index does not exist.
37    fn output(&self, index: usize) -> Self::Output;
38}
39
40/// Scheme for coloring nodes.
41#[derive(Deserialize, Debug, Clone, Copy, strum::EnumString)]
42#[serde(rename_all = "kebab-case")]
43#[strum(serialize_all = "kebab-case")]
44pub enum NodeColoringScheme {
45    /// Cumulative sum
46    ///
47    /// The cumulative size of a node and all its dependencies.
48    CumSum,
49    /// Dependency count
50    ///
51    /// The number of transitive dependencies a node has.
52    DepCount,
53    /// Reverse dependency count
54    ///
55    /// The number of paths from the root to a node.
56    RevDepCount,
57}
58
59impl From<NodeColoringScheme> for &'static str {
60    fn from(value: NodeColoringScheme) -> Self {
61        match value {
62            NodeColoringScheme::CumSum => "cumulative sum",
63            NodeColoringScheme::DepCount => "dependency count",
64            NodeColoringScheme::RevDepCount => "reverse dependency count",
65        }
66    }
67}
68
69/// Gradient for coloring nodes.
70///
71/// This implements the [`Gradient`] trait to provide colors based on input values of type `Option<f64>`.
72#[derive(Deserialize, Debug, Default, Clone, Copy, strum::EnumString)]
73#[serde(rename_all = "kebab-case")]
74#[strum(serialize_all = "kebab-case")]
75pub enum NodeColoringGradient {
76    /// Reds
77    #[default]
78    Reds,
79    /// Oranges
80    Oranges,
81    /// Purples
82    Purples,
83    /// Greens
84    Greens,
85    /// Blues
86    Blues,
87    /// Blue-Purple
88    BuPu,
89    /// Orange-Red
90    OrRd,
91    /// Purple-Red
92    PuRd,
93    /// Red-Purple
94    RdPu,
95    /// Viridis
96    Viridis,
97    /// Cividis
98    Cividis,
99    /// Plasma
100    Plasma,
101}
102
103impl From<NodeColoringGradient> for colorous::Gradient {
104    fn from(value: NodeColoringGradient) -> Self {
105        use colorous::*;
106        match value {
107            NodeColoringGradient::Reds => REDS,
108            NodeColoringGradient::Oranges => ORANGES,
109            NodeColoringGradient::Purples => PURPLES,
110            NodeColoringGradient::Greens => GREENS,
111            NodeColoringGradient::Blues => BLUES,
112            NodeColoringGradient::BuPu => BLUE_PURPLE,
113            NodeColoringGradient::OrRd => ORANGE_RED,
114            NodeColoringGradient::PuRd => PURPLE_RED,
115            NodeColoringGradient::RdPu => RED_PURPLE,
116            NodeColoringGradient::Viridis => VIRIDIS,
117            NodeColoringGradient::Cividis => CIVIDIS,
118            NodeColoringGradient::Plasma => PLASMA,
119        }
120    }
121}
122
123impl Gradient for NodeColoringGradient {
124    type Input = Option<f64>;
125
126    fn color(&self, input: Self::Input, dark_mode: bool, inverse: bool) -> Color {
127        if let Some(input) = input {
128            let input = input.clamp(0.0, 1.0);
129            let input = if inverse { 1.0 - input } else { input };
130            let mut color = colorous::Gradient::from(*self).eval_continuous(input);
131
132            if dark_mode {
133                let mut hsl: colorsys::Hsl =
134                    colorsys::Rgb::from(&(color.r, color.g, color.b)).into();
135                hsl.set_lightness(100.0 - hsl.lightness());
136                let (r, g, b) = colorsys::Rgb::from(hsl).into();
137                color = Color { r, g, b };
138            }
139            color
140        } else {
141            #[allow(clippy::collapsible_else_if)]
142            if dark_mode {
143                Color {
144                    r: 0x00,
145                    g: 0x00,
146                    b: 0x00,
147                }
148            } else {
149                Color {
150                    r: 0xff,
151                    g: 0xff,
152                    b: 0xff,
153                }
154            }
155        }
156    }
157}
158
159/// Values for coloring nodes.
160///
161/// This implements the [`Values`] trait to provide coloring values and outputs based on the selected
162/// [`NodeColoringScheme`].
163#[derive(Debug)]
164pub struct NodeColoringValues {
165    indices: Vec<usize>,
166    values: Vec<usize>,
167    gamma: f64,
168    max: usize,
169    scheme: NodeColoringScheme,
170}
171
172impl NodeColoringValues {
173    /// Create new coloring values for the given graph and scheme.
174    pub fn new(graph: &Graph, scheme: NodeColoringScheme) -> Self {
175        match scheme {
176            NodeColoringScheme::CumSum => Self::cum_sums(graph),
177            NodeColoringScheme::DepCount => Self::dep_counts(graph),
178            NodeColoringScheme::RevDepCount => Self::rev_dep_counts(graph),
179        }
180    }
181
182    /// Create cumulative sum coloring values for the given graph.
183    fn cum_sums(graph: &Graph) -> Self {
184        let mut values = vec![0; graph.node_capacity()];
185
186        for (index, size) in graph
187            .node_indices()
188            .filter_map(|i| graph.size(i).map(|s| (i, s)))
189        {
190            values[index] = size;
191        }
192
193        let nodes: Vec<usize> = graph.topo().collect();
194
195        for node in nodes.iter().rev() {
196            let sources: Vec<_> = graph.neighbors(*node, false).collect();
197            for source in sources.iter() {
198                values[*source] += values[*node] / sources.len();
199            }
200        }
201
202        let max = *values.iter().max().unwrap();
203        let indices = graph.node_indices().collect();
204
205        Self {
206            indices,
207            values,
208            gamma: 0.25,
209            max,
210            scheme: NodeColoringScheme::CumSum,
211        }
212    }
213
214    /// Create dependency count coloring values for the given graph.
215    fn dep_counts(graph: &Graph) -> Self {
216        let mut values = vec![0; graph.node_capacity()];
217
218        let nodes: Vec<usize> = graph.topo().collect();
219
220        for node in nodes.iter().rev() {
221            for target in graph.neighbors(*node, true) {
222                values[*node] += values[target] + 1;
223            }
224        }
225
226        let max = *values.iter().max().unwrap();
227        let indices = graph.node_indices().collect();
228
229        Self {
230            indices,
231            values,
232            gamma: 0.25,
233            max,
234            scheme: NodeColoringScheme::DepCount,
235        }
236    }
237
238    /// Create reverse dependency count coloring values for the given graph.
239    fn rev_dep_counts(graph: &Graph) -> Self {
240        let mut values = vec![0; graph.node_capacity()];
241
242        for node in graph.topo() {
243            for target in graph.neighbors(node, true) {
244                values[target] += 1;
245            }
246        }
247
248        let max = *values.iter().max().unwrap();
249        let indices = graph.node_indices().collect();
250
251        Self {
252            indices,
253            values,
254            gamma: 0.5,
255            max,
256            scheme: NodeColoringScheme::RevDepCount,
257        }
258    }
259
260    /// Get an iterator over the node indices of the graph as it was in [`new`](Self::new) and their
261    /// corresponding values.
262    pub fn indices_values(&self) -> impl Iterator<Item = (usize, usize)> {
263        self.indices.iter().copied().map(|i| (i, self.values[i]))
264    }
265
266    /// Set the gamma value, clamped between 0.0 and 1.0.
267    pub fn set_gamma(&mut self, gamma: f64) {
268        self.gamma = gamma.clamp(0.0, 1.0);
269    }
270
271    /// Get the gamma value.
272    pub fn gamma(&self) -> f64 {
273        self.gamma
274    }
275
276    /// Get the maximum value.
277    pub fn max(&self) -> usize {
278        self.max
279    }
280
281    /// Get the coloring scheme.
282    pub fn scheme(&self) -> NodeColoringScheme {
283        self.scheme
284    }
285}
286
287impl Values for Option<NodeColoringValues> {
288    type Context = Option<NodeColoringScheme>;
289    type Value = Option<usize>;
290    type Output = Option<f64>;
291
292    fn context(&self) -> Self::Context {
293        self.as_ref().map(|v| v.scheme)
294    }
295
296    fn value(&self, index: usize) -> Self::Value {
297        self.as_ref().map(|v| v.values[index])
298    }
299
300    fn output(&self, index: usize) -> Self::Output {
301        self.as_ref()
302            .map(|v| (v.values[index] as f64 / v.max as f64).powf(v.gamma))
303    }
304}