gpui_component/chart/
area_chart.rs

1use std::rc::Rc;
2
3use gpui::{px, App, Background, Bounds, Hsla, Pixels, SharedString, TextAlign, Window};
4use gpui_component_macros::IntoPlot;
5use num_traits::{Num, ToPrimitive};
6
7use crate::{
8    plot::{
9        scale::{Scale, ScaleLinear, ScalePoint, Sealed},
10        shape::Area,
11        Axis, AxisText, Grid, Plot, StrokeStyle, AXIS_GAP,
12    },
13    ActiveTheme, PixelsExt,
14};
15
16#[derive(IntoPlot)]
17pub struct AreaChart<T, X, Y>
18where
19    T: 'static,
20    X: Clone + PartialEq + Into<SharedString> + 'static,
21    Y: Clone + Copy + PartialOrd + Num + ToPrimitive + Sealed + 'static,
22{
23    data: Vec<T>,
24    x: Option<Rc<dyn Fn(&T) -> X>>,
25    y: Vec<Rc<dyn Fn(&T) -> Y>>,
26    stroke: Vec<Hsla>,
27    stroke_style: StrokeStyle,
28    fill: Vec<Background>,
29    tick_margin: usize,
30}
31
32impl<T, X, Y> AreaChart<T, X, Y>
33where
34    X: Clone + PartialEq + Into<SharedString> + 'static,
35    Y: Clone + Copy + PartialOrd + Num + ToPrimitive + Sealed + 'static,
36{
37    pub fn new<I>(data: I) -> Self
38    where
39        I: IntoIterator<Item = T>,
40    {
41        Self {
42            data: data.into_iter().collect(),
43            stroke_style: Default::default(),
44            stroke: vec![],
45            fill: vec![],
46            tick_margin: 1,
47            x: None,
48            y: vec![],
49        }
50    }
51
52    pub fn x(mut self, x: impl Fn(&T) -> X + 'static) -> Self {
53        self.x = Some(Rc::new(x));
54        self
55    }
56
57    pub fn y(mut self, y: impl Fn(&T) -> Y + 'static) -> Self {
58        self.y.push(Rc::new(y));
59        self
60    }
61
62    pub fn stroke(mut self, stroke: impl Into<Hsla>) -> Self {
63        self.stroke.push(stroke.into());
64        self
65    }
66
67    pub fn fill(mut self, fill: impl Into<Background>) -> Self {
68        self.fill.push(fill.into());
69        self
70    }
71
72    pub fn linear(mut self) -> Self {
73        self.stroke_style = StrokeStyle::Linear;
74        self
75    }
76
77    pub fn tick_margin(mut self, tick_margin: usize) -> Self {
78        self.tick_margin = tick_margin;
79        self
80    }
81}
82
83impl<T, X, Y> Plot for AreaChart<T, X, Y>
84where
85    X: Clone + PartialEq + Into<SharedString> + 'static,
86    Y: Clone + Copy + PartialOrd + Num + ToPrimitive + Sealed + 'static,
87{
88    fn paint(&mut self, bounds: Bounds<Pixels>, window: &mut Window, cx: &mut App) {
89        let Some(x_fn) = self.x.as_ref() else {
90            return;
91        };
92
93        if self.y.len() == 0 {
94            return;
95        }
96
97        let width = bounds.size.width.as_f32();
98        let height = bounds.size.height.as_f32() - AXIS_GAP;
99
100        // X scale
101        let x = ScalePoint::new(self.data.iter().map(|v| x_fn(v)).collect(), vec![0., width]);
102
103        // Y scale
104        let domain = self
105            .data
106            .iter()
107            .flat_map(|v| self.y.iter().map(|y_fn| y_fn(v)))
108            .chain(Some(Y::zero()))
109            .collect::<Vec<_>>();
110        let y = ScaleLinear::new(domain, vec![height, 10.]);
111
112        // Draw X axis
113        let data_len = self.data.len();
114        let x_label = self.data.iter().enumerate().filter_map(|(i, d)| {
115            if (i + 1) % self.tick_margin == 0 {
116                x.tick(&x_fn(d)).map(|x_tick| {
117                    let align = match i {
118                        0 => {
119                            if data_len == 1 {
120                                TextAlign::Center
121                            } else {
122                                TextAlign::Left
123                            }
124                        }
125                        i if i == data_len - 1 => TextAlign::Right,
126                        _ => TextAlign::Center,
127                    };
128                    AxisText::new(x_fn(d).into(), x_tick, cx.theme().muted_foreground).align(align)
129                })
130            } else {
131                None
132            }
133        });
134
135        Axis::new()
136            .x(height)
137            .x_label(x_label)
138            .stroke(cx.theme().border)
139            .paint(&bounds, window, cx);
140
141        // Draw grid
142        Grid::new()
143            .y((0..=3).map(|i| height * i as f32 / 4.0).collect())
144            .stroke(cx.theme().border)
145            .dash_array(&[px(4.), px(2.)])
146            .paint(&bounds, window);
147
148        // Draw area
149        for (i, y_fn) in self.y.iter().enumerate() {
150            let x = x.clone();
151            let y = y.clone();
152            let x_fn = x_fn.clone();
153            let y_fn = y_fn.clone();
154
155            let fill = *self
156                .fill
157                .get(i)
158                .unwrap_or(&cx.theme().chart_2.opacity(0.4).into());
159
160            let stroke = *self.stroke.get(i).unwrap_or(&cx.theme().chart_2);
161
162            Area::new()
163                .data(&self.data)
164                .x(move |d| x.tick(&x_fn(d)))
165                .y0(height)
166                .y1(move |d| y.tick(&y_fn(d)))
167                .stroke(stroke)
168                .stroke_style(self.stroke_style)
169                .fill(fill)
170                .paint(&bounds, window);
171        }
172    }
173}