gpui_component/chart/
sankey_chart.rs

1use std::{
2    hash::{DefaultHasher, Hash, Hasher},
3    rc::Rc,
4};
5
6use gpui::{
7    AnyElement, App, Bounds, Corners, ElementId, Hsla, IntoElement, Pixels, Point, SharedString,
8    TextAlign, Window, fill, linear_color_stop, linear_gradient, point, prelude::FluentBuilder, px,
9};
10use gpui_component_macros::IntoPlot;
11
12use crate::{
13    ActiveTheme,
14    plot::{
15        Plot,
16        label::{PlotLabel, TEXT_GAP, TEXT_SIZE, Text, measure_text_width, truncate_text_to_width},
17        origin_point,
18        shape::{
19            Sankey, SankeyAlign, SankeyGraph, SankeyLink, SankeyLinkLayout, SankeyValueScale,
20            sankey_link_path,
21        },
22        tooltip::{PlotHover, Tooltip, TooltipState},
23    },
24};
25
26const DEFAULT_NODE_WIDTH: f32 = 10.;
27const DEFAULT_NODE_PADDING: f32 = 16.;
28const DEFAULT_LINK_OPACITY: f32 = 0.3;
29const DEFAULT_MIN_LINK_WIDTH: f32 = 1.;
30const DEFAULT_LABEL_GAP: f32 = 6.;
31/// Cap each side's label margin (as a fraction of width) so a long label is
32/// truncated to a modest column beside the flow instead of dominating it.
33const MAX_LABEL_WIDTH_RATIO: f32 = 0.2;
34/// Cap the reserved top+bottom label band as a fraction of height.
35const MAX_LABEL_MARGIN_RATIO: f32 = 0.6;
36/// How much the links not attached to the hovered node fade, as a share of
37/// their opacity.
38const HOVER_DIM: f32 = 0.7;
39
40/// The placement of a sankey chart for one bounds size: the graph plus the
41/// label lines and margins it was laid out with.
42///
43/// Placing the graph relaxes the node order over several iterations, and the
44/// label margins need every label measured, so a chart keeps the frame in
45/// element state and reuses it while its key is unchanged.
46struct SankeyFrame {
47    graph: SankeyGraph,
48    layer_count: usize,
49    node_labels: Vec<Vec<SankeyLabel>>,
50    /// The label margins reserved on the left and right of the flow.
51    left: f32,
52    right: f32,
53}
54
55/// The frame of the last placement with the key it was placed for.
56#[derive(Default)]
57struct SankeyFrameCache {
58    key: Option<u64>,
59    frame: Option<Rc<SankeyFrame>>,
60}
61
62/// The hover a sankey chart paints, sampled once per frame in [`Plot::hover`].
63#[derive(Clone, Copy)]
64struct SankeyHover {
65    /// The hovered node.
66    node: usize,
67    /// How far the hover has faded in.
68    focus: f32,
69}
70
71/// A styled line of a sankey node label.
72#[derive(Clone)]
73pub struct SankeyLabel {
74    text: SharedString,
75    color: Option<Hsla>,
76    font_size: Option<f32>,
77}
78
79impl SankeyLabel {
80    /// Create a label line with the default color and font size.
81    pub fn new(text: impl Into<SharedString>) -> Self {
82        Self {
83            text: text.into(),
84            color: None,
85            font_size: None,
86        }
87    }
88
89    /// Set the text color. Defaults to the theme foreground.
90    pub fn color(mut self, color: impl Into<Hsla>) -> Self {
91        self.color = Some(color.into());
92        self
93    }
94
95    /// Set the font size. Defaults to 10.
96    pub fn font_size(mut self, font_size: f32) -> Self {
97        self.font_size = Some(font_size);
98        self
99    }
100
101    fn line_height(&self) -> f32 {
102        self.font_size.unwrap_or(TEXT_SIZE) + TEXT_GAP
103    }
104}
105
106fn block_height(lines: &[SankeyLabel]) -> f32 {
107    lines.iter().map(|line| line.line_height()).sum()
108}
109
110/// A Sankey diagram, layout modeled after [d3-sankey](https://github.com/d3/d3-sankey).
111///
112/// Links reference nodes by their index in the node list; map string ids to
113/// indices before constructing.
114#[derive(IntoPlot)]
115pub struct SankeyChart<T: 'static> {
116    nodes: Vec<T>,
117    links: Vec<SankeyLink>,
118    node_width: f32,
119    node_padding: f32,
120    align: SankeyAlign,
121    iterations: usize,
122    value_scale: SankeyValueScale,
123    node_corner_radius: Option<Pixels>,
124    node_color: Option<Rc<dyn Fn(&T) -> Hsla>>,
125    node_label: Option<Rc<dyn Fn(&T) -> SharedString>>,
126    value_label: Option<Rc<dyn Fn(&T, f64) -> SharedString>>,
127    labels: Option<Rc<dyn Fn(&T, f64) -> Vec<SankeyLabel>>>,
128    link_opacity: f32,
129    min_link_width: f32,
130    label_gap: f32,
131    id: Option<ElementId>,
132    /// The placement for this frame, resolved in `prepaint` (measuring labels
133    /// needs the window) and read by `tooltip_state` and `paint`.
134    frame: Option<Rc<SankeyFrame>>,
135    hover: Option<SankeyHover>,
136}
137
138impl<T> SankeyChart<T> {
139    /// Create a chart from nodes and links; links reference nodes by their
140    /// index in `nodes` (map string ids to indices before constructing).
141    pub fn new<I, L>(nodes: I, links: L) -> Self
142    where
143        I: IntoIterator<Item = T>,
144        L: IntoIterator<Item = SankeyLink>,
145    {
146        Self {
147            nodes: nodes.into_iter().collect(),
148            links: links.into_iter().collect(),
149            node_width: DEFAULT_NODE_WIDTH,
150            node_padding: DEFAULT_NODE_PADDING,
151            align: SankeyAlign::default(),
152            iterations: 6,
153            value_scale: SankeyValueScale::default(),
154            node_corner_radius: None,
155            node_color: None,
156            node_label: None,
157            value_label: None,
158            labels: None,
159            link_opacity: DEFAULT_LINK_OPACITY,
160            min_link_width: DEFAULT_MIN_LINK_WIDTH,
161            label_gap: DEFAULT_LABEL_GAP,
162            id: None,
163            frame: None,
164            hover: None,
165        }
166    }
167
168    /// Enable an interactive hover tooltip for this chart: the links of the
169    /// hovered node stand out from the rest and the tooltip shows its label and
170    /// throughput.
171    ///
172    /// The `id` must be unique among sibling elements. Without it, the chart
173    /// stays a non-interactive plot.
174    pub fn id(mut self, id: impl Into<ElementId>) -> Self {
175        self.id = Some(id.into());
176        self
177    }
178
179    /// Set the node rectangle width. Defaults to 10.
180    pub fn node_width(mut self, node_width: f32) -> Self {
181        self.node_width = node_width;
182        self
183    }
184
185    /// Set the vertical gap between nodes in a column. Defaults to 16.
186    pub fn node_padding(mut self, node_padding: f32) -> Self {
187        self.node_padding = node_padding;
188        self
189    }
190
191    /// Set the node alignment. Defaults to [`SankeyAlign::Justify`].
192    pub fn node_align(mut self, align: SankeyAlign) -> Self {
193        self.align = align;
194        self
195    }
196
197    /// Set the number of relaxation passes. Defaults to 6.
198    pub fn iterations(mut self, iterations: usize) -> Self {
199        self.iterations = iterations;
200        self
201    }
202
203    /// Set how flow values map to node heights.
204    ///
205    /// Defaults to [`SankeyValueScale::Linear`]. Use [`SankeyValueScale::Sqrt`]
206    /// to keep a dominant flow from dwarfing the small ones without
207    /// pre-transforming the data; labels still receive the raw values.
208    pub fn value_scale(mut self, value_scale: SankeyValueScale) -> Self {
209        self.value_scale = value_scale;
210        self
211    }
212
213    /// Set the corner radius of the node rectangles. Defaults to 0.
214    pub fn node_corner_radius(mut self, radius: impl Into<Pixels>) -> Self {
215        self.node_corner_radius = Some(radius.into());
216        self
217    }
218
219    /// Set the color of each node.
220    ///
221    /// Defaults to cycling the theme chart palette by node index.
222    pub fn node_color<H>(mut self, color: impl Fn(&T) -> H + 'static) -> Self
223    where
224        H: Into<Hsla> + 'static,
225    {
226        self.node_color = Some(Rc::new(move |t| color(t).into()));
227        self
228    }
229
230    /// Set the name label of each node, drawn in muted foreground. No name
231    /// label is drawn unless set.
232    pub fn node_label(mut self, label: impl Fn(&T) -> SharedString + 'static) -> Self {
233        self.node_label = Some(Rc::new(label));
234        self
235    }
236
237    /// Set the value label of each node, drawn above the name label. No value
238    /// label is drawn unless set.
239    ///
240    /// The closure receives the datum and the node's raw computed throughput
241    /// (max of incoming and outgoing flow, in unscaled units).
242    pub fn value_label(mut self, label: impl Fn(&T, f64) -> SharedString + 'static) -> Self {
243        self.value_label = Some(Rc::new(label));
244        self
245    }
246
247    /// Set fully custom node labels, one [`SankeyLabel`] per line, top to
248    /// bottom. Takes precedence over `node_label`/`value_label` when set;
249    /// unset by default.
250    ///
251    /// The closure receives the datum and the node's raw computed throughput
252    /// (max of incoming and outgoing flow, in unscaled units).
253    pub fn labels(mut self, labels: impl Fn(&T, f64) -> Vec<SankeyLabel> + 'static) -> Self {
254        self.labels = Some(Rc::new(labels));
255        self
256    }
257
258    /// Set the opacity of the link ribbons. Defaults to 0.3.
259    pub fn link_opacity(mut self, opacity: f32) -> Self {
260        self.link_opacity = opacity;
261        self
262    }
263
264    /// Set the minimum ribbon thickness, so tiny flows stay visible. Defaults to 1.
265    pub fn min_link_width(mut self, width: f32) -> Self {
266        self.min_link_width = width;
267        self
268    }
269
270    /// Set the gap between a node and its labels. Defaults to 6.
271    pub fn label_gap(mut self, gap: f32) -> Self {
272        self.label_gap = gap;
273        self
274    }
275
276    fn sankey(&self) -> Sankey {
277        Sankey::new()
278            .node_width(self.node_width)
279            .node_padding(self.node_padding)
280            .node_align(self.align)
281            .iterations(self.iterations)
282            .value_scale(self.value_scale)
283    }
284
285    /// Raw per-node throughput (max of raw incoming and outgoing sums), for
286    /// labels — the layout's `node.value` is in scaled units under a
287    /// non-linear value scale, so labels must not use it.
288    fn raw_throughput(&self) -> Vec<f64> {
289        let mut incoming = vec![0f64; self.nodes.len()];
290        let mut outgoing = vec![0f64; self.nodes.len()];
291        for link in &self.links {
292            if let (Some(o), Some(i)) =
293                (outgoing.get_mut(link.source), incoming.get_mut(link.target))
294            {
295                *o += link.value;
296                *i += link.value;
297            }
298        }
299        incoming
300            .into_iter()
301            .zip(outgoing)
302            .map(|(i, o)| i.max(o))
303            .collect()
304    }
305}
306
307impl<T> SankeyChart<T> {
308    /// Each node's label lines: the custom `labels` closure wins, otherwise the
309    /// value/name lines with the default styles. Labels get the raw throughput,
310    /// not the layout's (possibly scaled) value.
311    fn node_labels(&self, cx: &App) -> Vec<Vec<SankeyLabel>> {
312        let raw_value = self.raw_throughput();
313        self.nodes
314            .iter()
315            .zip(raw_value)
316            .map(|(datum, value)| {
317                if let Some(labels) = &self.labels {
318                    labels(datum, value)
319                } else {
320                    let mut lines = Vec::new();
321                    if let Some(value_label) = &self.value_label {
322                        lines.push(SankeyLabel::new(value_label(datum, value)));
323                    }
324                    if let Some(node_label) = &self.node_label {
325                        lines.push(
326                            SankeyLabel::new(node_label(datum)).color(cx.theme().muted_foreground),
327                        );
328                    }
329                    lines
330                }
331            })
332            .collect()
333    }
334
335    /// The key a placement is reused under: everything that shapes it, which is
336    /// the bounds size, the graph, the placement settings and the label lines.
337    fn frame_key(&self, bounds: Bounds<Pixels>, node_labels: &[Vec<SankeyLabel>]) -> u64 {
338        let mut hasher = DefaultHasher::new();
339        bounds.size.width.as_f32().to_bits().hash(&mut hasher);
340        bounds.size.height.as_f32().to_bits().hash(&mut hasher);
341        self.nodes.len().hash(&mut hasher);
342        for link in &self.links {
343            link.source.hash(&mut hasher);
344            link.target.hash(&mut hasher);
345            link.value.to_bits().hash(&mut hasher);
346        }
347        self.node_width.to_bits().hash(&mut hasher);
348        self.node_padding.to_bits().hash(&mut hasher);
349        self.align.hash(&mut hasher);
350        self.iterations.hash(&mut hasher);
351        self.value_scale.hash(&mut hasher);
352        self.label_gap.to_bits().hash(&mut hasher);
353        for lines in node_labels {
354            lines.len().hash(&mut hasher);
355            for line in lines {
356                line.text.hash(&mut hasher);
357                line.font_size.map(f32::to_bits).hash(&mut hasher);
358                line.color
359                    .map(|color| [color.h, color.s, color.l, color.a].map(f32::to_bits))
360                    .hash(&mut hasher);
361            }
362        }
363        hasher.finish()
364    }
365
366    /// Place the graph within `bounds`, reserving margins for the labels.
367    fn place(
368        &self,
369        bounds: Bounds<Pixels>,
370        node_labels: Vec<Vec<SankeyLabel>>,
371        window: &mut Window,
372    ) -> Option<SankeyFrame> {
373        let width = bounds.size.width.as_f32();
374        let height = bounds.size.height.as_f32();
375
376        // First pass: only the topology (each node's `layer`) is needed to
377        // measure the label margins.
378        let topology = self.sankey().topology(self.nodes.len(), &self.links).ok()?;
379        let layer_count = topology.layer_count();
380        let has_labels = node_labels.iter().any(|lines| !lines.is_empty());
381
382        // Reserve margins so the labels beside the first/last columns and
383        // above the middle columns are not clipped.
384        let mut left = 0f32;
385        let mut right = 0f32;
386        if has_labels {
387            for node in &topology.nodes {
388                if node.layer != 0 && node.layer + 1 != layer_count {
389                    continue;
390                }
391                let mut label_width = 0f32;
392                for line in &node_labels[node.index] {
393                    label_width = label_width.max(measure_text_width(
394                        &line.text,
395                        px(line.font_size.unwrap_or(TEXT_SIZE)),
396                        window,
397                    ));
398                }
399                if node.layer == 0 {
400                    left = left.max(label_width + self.label_gap);
401                } else {
402                    right = right.max(label_width + self.label_gap);
403                }
404            }
405
406            // Cap each side independently so one long label is truncated to a
407            // modest column rather than eating into the flow area.
408            let side_cap = width * MAX_LABEL_WIDTH_RATIO;
409            left = left.min(side_cap);
410            right = right.min(side_cap);
411        }
412        // Above-node labels are only emitted for middle columns, so reserve
413        // the top band for the tallest such label block. Cap the vertical
414        // margins like the horizontal ones so a short chart doesn't collapse
415        // the flow.
416        let mut top = 0f32;
417        if has_labels && layer_count > 2 {
418            for node in &topology.nodes {
419                if node.layer == 0 || node.layer + 1 == layer_count {
420                    continue;
421                }
422                let block = block_height(&node_labels[node.index]);
423                if block > 0. {
424                    top = top.max(block + TEXT_GAP);
425                }
426            }
427        }
428        let mut bottom = if has_labels { TEXT_GAP } else { 0. };
429        let max_vertical = height * MAX_LABEL_MARGIN_RATIO;
430        if top + bottom > max_vertical {
431            let k = max_vertical / (top + bottom);
432            top *= k;
433            bottom *= k;
434        }
435
436        // Second pass: complete the placement on the final extent, reusing
437        // the first pass's topology.
438        let graph = self
439            .sankey()
440            .extent(
441                left,
442                top,
443                (width - right).max(left + 1.),
444                (height - bottom).max(top + 1.),
445            )
446            .layout_from(topology);
447
448        Some(SankeyFrame {
449            graph,
450            layer_count,
451            node_labels,
452            left,
453            right,
454        })
455    }
456
457    /// Whether `link` starts or ends at `node`.
458    fn is_attached(link: &SankeyLinkLayout, node: usize) -> bool {
459        link.source == node || link.target == node
460    }
461}
462
463impl<T> Plot for SankeyChart<T> {
464    /// Resolve the placement for the frame, reusing the last one while nothing
465    /// that shapes it has changed. Measuring the labels needs the window, which
466    /// `tooltip_state` does not have, so this runs here rather than in `paint`.
467    fn prepaint(
468        &mut self,
469        bounds: Bounds<Pixels>,
470        window: &mut Window,
471        cx: &mut App,
472    ) -> Vec<AnyElement> {
473        self.frame = None;
474        let width = bounds.size.width.as_f32();
475        let height = bounds.size.height.as_f32();
476        if self.nodes.is_empty() || self.links.is_empty() || width <= 0. || height <= 0. {
477            return vec![];
478        }
479
480        let node_labels = self.node_labels(cx);
481
482        // An identified chart keeps its placement across frames; without an id,
483        // sibling charts would share one cache and thrash it.
484        self.frame = if self.id.is_some() {
485            let key = self.frame_key(bounds, &node_labels);
486            let cache =
487                window.use_keyed_state("sankey-frame", cx, |_, _| SankeyFrameCache::default());
488            let cached = cache.read(cx);
489            if cached.key == Some(key) {
490                cached.frame.clone()
491            } else {
492                let frame = self.place(bounds, node_labels, window).map(Rc::new);
493                cache.update(cx, |cache, _| {
494                    cache.key = Some(key);
495                    cache.frame = frame.clone();
496                });
497                frame
498            }
499        } else {
500            self.place(bounds, node_labels, window).map(Rc::new)
501        };
502
503        vec![]
504    }
505
506    fn paint(&mut self, bounds: Bounds<Pixels>, window: &mut Window, cx: &mut App) {
507        let Some(frame) = self.frame.clone() else {
508            return;
509        };
510        let SankeyFrame {
511            graph,
512            layer_count,
513            node_labels,
514            left,
515            right,
516        } = &*frame;
517        let (layer_count, left, right) = (*layer_count, *left, *right);
518        let width = bounds.size.width.as_f32();
519        let height = bounds.size.height.as_f32();
520
521        let palette = [
522            cx.theme().chart_1,
523            cx.theme().chart_2,
524            cx.theme().chart_3,
525            cx.theme().chart_4,
526            cx.theme().chart_5,
527        ];
528        let colors: Vec<Hsla> = self
529            .nodes
530            .iter()
531            .enumerate()
532            .map(|(index, datum)| match &self.node_color {
533                Some(color) => color(datum),
534                None => palette[index % palette.len()],
535            })
536            .collect();
537
538        // Links first, under the nodes. The links of the hovered node keep their
539        // opacity while the rest fade behind them.
540        for link in &graph.links {
541            if link.value <= 0. {
542                continue;
543            }
544            let source = &graph.nodes[link.source];
545            let target = &graph.nodes[link.target];
546            let Some(path) =
547                sankey_link_path(source, target, link, self.min_link_width, bounds.origin)
548            else {
549                continue;
550            };
551            let opacity = match self.hover {
552                Some(hover) if !Self::is_attached(link, hover.node) => {
553                    self.link_opacity * (1. - HOVER_DIM * hover.focus)
554                }
555                _ => self.link_opacity,
556            };
557            window.paint_path(
558                path,
559                linear_gradient(
560                    90.,
561                    linear_color_stop(colors[link.source].opacity(opacity), 0.),
562                    linear_color_stop(colors[link.target].opacity(opacity), 1.),
563                ),
564            );
565        }
566
567        let corner_radii = Corners::all(self.node_corner_radius.unwrap_or_default());
568        for node in &graph.nodes {
569            let node_bounds = Bounds::from_corners(
570                origin_point(px(node.x0), px(node.y0), bounds.origin),
571                // Keep tiny nodes visible with a minimum 1px height.
572                origin_point(px(node.x1), px(node.y1.max(node.y0 + 1.)), bounds.origin),
573            );
574            window.paint_quad(fill(node_bounds, colors[node.index]).corner_radii(corner_radii));
575        }
576
577        let mut texts = Vec::new();
578        for node in &graph.nodes {
579            let lines = &node_labels[node.index];
580            if lines.is_empty() {
581                continue;
582            }
583
584            let is_first = node.layer == 0;
585            let is_last = node.layer + 1 == layer_count;
586            // `x`/`align` place the label beside (first/last) or centered above
587            // (middle) the node, and `max_width` bounds it so a long label is
588            // truncated with an ellipsis instead of drawn outside the plot:
589            // first/last to their reserved margin, middle to twice the smaller
590            // gap to the plot edge (generous for interior nodes, only bites a
591            // label long enough to actually run off-plot).
592            let (x, align, max_width) = if is_first {
593                (
594                    node.x0 - self.label_gap,
595                    TextAlign::Right,
596                    left - self.label_gap,
597                )
598            } else if is_last {
599                (
600                    node.x1 + self.label_gap,
601                    TextAlign::Left,
602                    right - self.label_gap,
603                )
604            } else {
605                let center = (node.x0 + node.x1) / 2.;
606                let edge_budget = 2. * center.min(width - center);
607                (center, TextAlign::Center, edge_budget)
608            };
609
610            let block = block_height(lines);
611            let mut y = if is_first || is_last {
612                // Block vertically centered beside the node, clamped into
613                // the plot area so labels of nodes near the top or bottom
614                // edge are not clipped.
615                ((node.y0 + node.y1) / 2. - block / 2.)
616                    .min(height - block)
617                    .max(0.)
618            } else {
619                // Block above the node.
620                node.y0 - block - TEXT_GAP
621            };
622
623            for line in lines {
624                let font_size = px(line.font_size.unwrap_or(TEXT_SIZE));
625                let text = truncate_text_to_width(&line.text, font_size, max_width, window);
626                texts.push(
627                    Text::new(
628                        text,
629                        point(px(x), px(y)),
630                        line.color.unwrap_or(cx.theme().foreground),
631                    )
632                    .font_size(font_size)
633                    .align(align),
634                );
635                y += line.line_height();
636            }
637        }
638        PlotLabel::new(texts).paint(&bounds, window, cx);
639    }
640
641    fn id(&self) -> Option<ElementId> {
642        self.id.clone()
643    }
644
645    fn tooltip_state(
646        &self,
647        position: Point<Pixels>,
648        _bounds: Bounds<Pixels>,
649        _cx: &App,
650    ) -> Option<TooltipState> {
651        let frame = self.frame.as_ref()?;
652        let (x, y) = (position.x.as_f32(), position.y.as_f32());
653        let node = frame.graph.nodes.iter().find(|node| {
654            (node.x0..=node.x1).contains(&x) && (node.y0..=node.y1.max(node.y0 + 1.)).contains(&y)
655        })?;
656        Some(TooltipState::new(node.index, position, vec![]))
657    }
658
659    fn hover(&mut self, hover: Option<&PlotHover>, _window: &mut Window, _cx: &mut App) {
660        self.hover = hover.map(|hover| SankeyHover {
661            node: hover.state().index,
662            focus: hover.focus(),
663        });
664    }
665
666    fn tooltip(
667        &self,
668        state: &TooltipState,
669        cursor: Point<Pixels>,
670        bounds: Bounds<Pixels>,
671        _window: &mut Window,
672        cx: &mut App,
673    ) -> Option<AnyElement> {
674        let datum = self.nodes.get(state.index)?;
675        let value = self.raw_throughput().get(state.index).copied()?;
676        let color = match &self.node_color {
677            Some(color) => color(datum),
678            None => {
679                let palette = [
680                    cx.theme().chart_1,
681                    cx.theme().chart_2,
682                    cx.theme().chart_3,
683                    cx.theme().chart_4,
684                    cx.theme().chart_5,
685                ];
686                palette[state.index % palette.len()]
687            }
688        };
689        let value_text = match &self.value_label {
690            Some(value_label) => value_label(datum, value),
691            None => format!("{}", value).into(),
692        };
693
694        Some(
695            // Follow the cursor; the node's links mark it.
696            Tooltip::new(cursor, bounds.size)
697                .gap(px(8.))
698                .when_some(self.node_label.as_ref(), |this, label| {
699                    this.title(label(datum))
700                })
701                .row(color, SharedString::default(), value_text)
702                .into_any_element(),
703        )
704    }
705}
706
707#[cfg(test)]
708mod tests {
709    use super::*;
710
711    #[test]
712    fn test_sankey_chart_builder() {
713        let chart = SankeyChart::new(vec!["a", "b"], vec![SankeyLink::new(0, 1, 5.)]);
714        assert_eq!(chart.nodes.len(), 2);
715        assert_eq!(chart.links.len(), 1);
716        assert_eq!(chart.node_width, DEFAULT_NODE_WIDTH);
717        assert_eq!(chart.node_padding, DEFAULT_NODE_PADDING);
718        assert_eq!(chart.align, SankeyAlign::Justify);
719        assert_eq!(chart.iterations, 6);
720        assert_eq!(chart.node_corner_radius, None);
721        assert_eq!(chart.link_opacity, DEFAULT_LINK_OPACITY);
722        assert_eq!(chart.min_link_width, DEFAULT_MIN_LINK_WIDTH);
723        assert_eq!(chart.label_gap, DEFAULT_LABEL_GAP);
724        assert!(chart.node_color.is_none());
725        assert!(chart.node_label.is_none());
726        assert!(chart.value_label.is_none());
727        assert!(chart.labels.is_none());
728
729        let chart = chart
730            .node_width(8.)
731            .node_padding(20.)
732            .node_align(SankeyAlign::Left)
733            .iterations(10)
734            .node_corner_radius(px(2.))
735            .node_color(|_| gpui::red())
736            .node_label(|d| SharedString::from(d.to_string()))
737            .value_label(|_, value| SharedString::from(format!("{}", value)))
738            .labels(|d, value| {
739                vec![
740                    SankeyLabel::new(format!("{}", value)),
741                    SankeyLabel::new(d.to_string()),
742                ]
743            })
744            .link_opacity(0.5)
745            .min_link_width(2.)
746            .label_gap(10.);
747        assert_eq!(chart.node_width, 8.);
748        assert_eq!(chart.node_padding, 20.);
749        assert_eq!(chart.align, SankeyAlign::Left);
750        assert_eq!(chart.iterations, 10);
751        assert_eq!(chart.node_corner_radius, Some(px(2.)));
752        assert_eq!(chart.link_opacity, 0.5);
753        assert_eq!(chart.min_link_width, 2.);
754        assert_eq!(chart.label_gap, 10.);
755        assert!(chart.node_color.is_some());
756        assert!(chart.node_label.is_some());
757        assert!(chart.value_label.is_some());
758        assert!(chart.labels.is_some());
759    }
760
761    #[test]
762    fn test_sankey_label_builder() {
763        let label = SankeyLabel::new("a");
764        assert_eq!(label.text, "a");
765        assert_eq!(label.color, None);
766        assert_eq!(label.font_size, None);
767        assert_eq!(label.line_height(), TEXT_SIZE + TEXT_GAP);
768
769        let label = SankeyLabel::new("b").color(gpui::red()).font_size(14.);
770        assert_eq!(label.color, Some(gpui::red()));
771        assert_eq!(label.font_size, Some(14.));
772        assert_eq!(label.line_height(), 14. + TEXT_GAP);
773
774        assert_eq!(
775            block_height(&[SankeyLabel::new("a"), SankeyLabel::new("b").font_size(14.)]),
776            TEXT_SIZE + TEXT_GAP + 14. + TEXT_GAP
777        );
778        assert_eq!(block_height(&[]), 0.);
779    }
780
781    #[test]
782    fn test_sankey_chart_raw_throughput() {
783        // A(out 30) -> B, B -> C(20) + D(10): B's throughput is max(in, out).
784        let chart = SankeyChart::new(
785            vec!["a", "b", "c", "d"],
786            vec![
787                SankeyLink::new(0, 1, 30.),
788                SankeyLink::new(1, 2, 20.),
789                SankeyLink::new(1, 3, 10.),
790            ],
791        );
792        let raw = chart.raw_throughput();
793        assert_eq!(raw, vec![30., 30., 20., 10.]);
794
795        // Under Sqrt the layout's node value is scaled, but raw_throughput
796        // (used for labels) must stay in raw units — the two must differ.
797        let sqrt = chart
798            .value_scale(SankeyValueScale::Sqrt)
799            .sankey()
800            .layout(4, &chart_links())
801            .unwrap();
802        // Node A: layout value is sqrt-scaled (30 -> sqrt(30)), raw is 30.
803        assert!((sqrt.nodes[0].value - 30f64.sqrt()).abs() < 1e-6);
804        assert!((raw[0] - 30.).abs() < 1e-6);
805        assert!(raw[0] != sqrt.nodes[0].value);
806    }
807
808    fn chart_links() -> Vec<SankeyLink> {
809        vec![
810            SankeyLink::new(0, 1, 30.),
811            SankeyLink::new(1, 2, 20.),
812            SankeyLink::new(1, 3, 10.),
813        ]
814    }
815}
gpui_component/chart/sankey_chart.rs

gpui_component/chart/
sankey_chart.rs
