1use crate::aes::Aesthetic;
2use crate::coord::Coord;
3use crate::data::DataFrame;
4use crate::position::identity::PositionIdentity;
5use crate::position::Position;
6use crate::render::backend::{DrawBackend, LineStyle, Linetype};
7use crate::render::RenderError;
8use crate::scale::ScaleSet;
9use crate::stat::identity::StatIdentity;
10use crate::stat::Stat;
11use crate::theme::Theme;
12
13use super::{Geom, GeomParams};
14
15pub enum StepDirection {
17 Hv,
19 Vh,
21}
22
23pub struct GeomStep {
25 pub color: (u8, u8, u8),
26 pub width: f64,
27 pub alpha: f64,
28 pub direction: StepDirection,
29}
30
31impl Default for GeomStep {
32 fn default() -> Self {
33 GeomStep {
34 color: (0, 0, 0),
35 width: 1.5,
36 alpha: 1.0,
37 direction: StepDirection::Hv,
38 }
39 }
40}
41
42impl Geom for GeomStep {
43 fn draw(
44 &self,
45 data: &DataFrame,
46 coord: &dyn Coord,
47 scales: &ScaleSet,
48 _theme: &Theme,
49 backend: &mut dyn DrawBackend,
50 ) -> Result<(), RenderError> {
51 let x_col = data
52 .column("x")
53 .ok_or(RenderError::MissingAesthetic("x".into()))?;
54 let y_col = data
55 .column("y")
56 .ok_or(RenderError::MissingAesthetic("y".into()))?;
57 let color_col = data.column("color");
58
59 let plot_area = backend.plot_area();
60 let x_scale = scales.get(&Aesthetic::X);
61 let y_scale = scales.get(&Aesthetic::Y);
62
63 let mut sorted: Vec<usize> = (0..data.nrows()).collect();
65 sorted.sort_by(|&a, &b| {
66 let xa = x_col[a].as_f64().unwrap_or(0.0);
67 let xb = x_col[b].as_f64().unwrap_or(0.0);
68 xa.partial_cmp(&xb).unwrap_or(std::cmp::Ordering::Equal)
69 });
70
71 let raw: Vec<(f64, f64)> = sorted
73 .iter()
74 .map(|&i| {
75 let nx = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.0);
76 let ny = y_scale.map(|s| s.map(&y_col[i])).unwrap_or(0.0);
77 (nx, ny)
78 })
79 .collect();
80
81 let mut step_points: Vec<(f64, f64)> = Vec::new();
83 for (j, &(nx, ny)) in raw.iter().enumerate() {
84 if j > 0 {
85 let (prev_nx, prev_ny) = raw[j - 1];
86 match self.direction {
87 StepDirection::Hv => step_points.push((nx, prev_ny)),
88 StepDirection::Vh => step_points.push((prev_nx, ny)),
89 }
90 }
91 step_points.push((nx, ny));
92 }
93
94 let points: Vec<(f64, f64)> = step_points
95 .iter()
96 .map(|&(nx, ny)| {
97 let (px, py) = coord.transform((nx, ny), &plot_area);
98 let px = if px.is_finite() {
102 px
103 } else if nx < 0.0 {
104 plot_area.x
105 } else {
106 plot_area.x + plot_area.width
107 };
108 let py = if py.is_finite() {
109 py
110 } else if ny < 0.0 {
111 plot_area.y + plot_area.height
112 } else {
113 plot_area.y
114 };
115 (px, py)
116 })
117 .collect();
118
119 let line_color = color_col
120 .and_then(|cc| {
121 if cc.is_empty() {
122 None
123 } else {
124 scales.map_color(&Aesthetic::Color, &cc[0])
125 }
126 })
127 .unwrap_or(self.color);
128
129 if points.len() >= 2 {
130 backend.draw_line(
131 &points,
132 &LineStyle {
133 color: line_color,
134 alpha: self.alpha,
135 width: self.width,
136 linetype: Linetype::Solid,
137 },
138 )?;
139 }
140
141 Ok(())
142 }
143
144 fn required_aes(&self) -> Vec<Aesthetic> {
145 vec![Aesthetic::X, Aesthetic::Y]
146 }
147
148 fn default_stat(&self) -> Box<dyn Stat> {
149 Box::new(StatIdentity)
150 }
151
152 fn default_position(&self) -> Box<dyn Position> {
153 Box::new(PositionIdentity)
154 }
155
156 fn default_params(&self) -> GeomParams {
157 GeomParams::default()
158 }
159
160 fn name(&self) -> &str {
161 "step"
162 }
163
164 fn set_series_color(&mut self, color: (u8, u8, u8)) {
165 self.color = color;
166 }
167}