1use std::{hash::Hash, rc::Rc};
2
3use gpui_kit::{
4 AnyElement, App, Bounds, ElementId, Hsla, IntoElement, PathBuilder, Pixels, Point,
5 SharedString, Window, fill, point, px,
6};
7
8use crate::{
9 ActiveTheme,
10 plot::{
11 Grid, Plot, PlotAxis, origin_point,
12 scale::{PlotValue, Scale, ScaleBand, ScaleLinear},
13 tooltip::{CrossLine, Tooltip, TooltipState},
14 },
15};
16
17use super::{
18 AXIS_GAP, MAX_BAND_WIDTH, TooltipContent, build_band_labels, caller_id, labeled_items,
19};
20
21pub struct CandlestickChart<T, X, Y>
22where
23 T: 'static,
24 X: Eq + Hash + Into<SharedString> + 'static,
25 Y: PlotValue,
26{
27 data: Vec<T>,
28 x: Option<Rc<dyn Fn(&T) -> X>>,
29 open: Option<Rc<dyn Fn(&T) -> Y>>,
30 high: Option<Rc<dyn Fn(&T) -> Y>>,
31 low: Option<Rc<dyn Fn(&T) -> Y>>,
32 close: Option<Rc<dyn Fn(&T) -> Y>>,
33 tick_margin: usize,
34 body_width_ratio: f32,
35 max_band_width: Pixels,
36 x_axis: bool,
37 grid: bool,
38 bullish: Option<Hsla>,
39 bearish: Option<Hsla>,
40 id: ElementId,
41 interactive: bool,
42 tooltip_content: TooltipContent<T>,
43}
44
45impl<T, X, Y> CandlestickChart<T, X, Y>
46where
47 X: Eq + Hash + Into<SharedString> + 'static,
48 Y: PlotValue,
49{
50 #[track_caller]
51 pub fn new<I>(data: I) -> Self
52 where
53 I: IntoIterator<Item = T>,
54 {
55 Self {
56 data: data.into_iter().collect(),
57 x: None,
58 open: None,
59 high: None,
60 low: None,
61 close: None,
62 tick_margin: 1,
63 body_width_ratio: 0.8,
64 max_band_width: px(MAX_BAND_WIDTH),
65 x_axis: true,
66 grid: true,
67 bullish: None,
68 bearish: None,
69 id: caller_id(),
70 interactive: true,
71 tooltip_content: TooltipContent::default(),
72 }
73 }
74
75 pub fn id(mut self, id: impl Into<ElementId>) -> Self {
82 self.id = id.into();
83 self
84 }
85
86 pub fn interactive(mut self, interactive: bool) -> Self {
95 self.interactive = interactive;
96 self
97 }
98
99 pub fn tooltip_title(mut self, title: impl Fn(&T) -> SharedString + 'static) -> Self {
101 self.tooltip_content.set_title(title);
102 self
103 }
104
105 pub fn tooltip_value(
110 mut self,
111 value: impl Fn(&T, usize, f64) -> SharedString + 'static,
112 ) -> Self {
113 self.tooltip_content.set_value(value);
114 self
115 }
116
117 pub fn tooltip_value_color<H>(mut self, color: impl Fn(&T, usize, f64) -> H + 'static) -> Self
123 where
124 H: Into<Hsla>,
125 {
126 self.tooltip_content.set_value_color(color);
127 self
128 }
129
130 pub fn tooltip_content<E>(
137 mut self,
138 content: impl Fn(&T, &mut Window, &mut App) -> E + 'static,
139 ) -> Self
140 where
141 E: IntoElement,
142 {
143 self.tooltip_content.set_content(content);
144 self
145 }
146
147 pub fn x(mut self, x: impl Fn(&T) -> X + 'static) -> Self {
148 self.x = Some(Rc::new(x));
149 self
150 }
151
152 pub fn open(mut self, open: impl Fn(&T) -> Y + 'static) -> Self {
153 self.open = Some(Rc::new(open));
154 self
155 }
156
157 pub fn high(mut self, high: impl Fn(&T) -> Y + 'static) -> Self {
158 self.high = Some(Rc::new(high));
159 self
160 }
161
162 pub fn low(mut self, low: impl Fn(&T) -> Y + 'static) -> Self {
163 self.low = Some(Rc::new(low));
164 self
165 }
166
167 pub fn close(mut self, close: impl Fn(&T) -> Y + 'static) -> Self {
168 self.close = Some(Rc::new(close));
169 self
170 }
171
172 pub fn tick_margin(mut self, tick_margin: usize) -> Self {
173 self.tick_margin = tick_margin;
174 self
175 }
176
177 pub fn body_width_ratio(mut self, ratio: f32) -> Self {
178 self.body_width_ratio = ratio;
179 self
180 }
181
182 pub fn max_band_width(mut self, width: impl Into<Pixels>) -> Self {
187 self.max_band_width = width.into();
188 self
189 }
190
191 pub fn x_axis(mut self, x_axis: bool) -> Self {
195 self.x_axis = x_axis;
196 self
197 }
198
199 pub fn grid(mut self, grid: bool) -> Self {
200 self.grid = grid;
201 self
202 }
203
204 pub fn bullish(mut self, color: impl Into<Hsla>) -> Self {
209 self.bullish = Some(color.into());
210 self
211 }
212
213 pub fn bearish(mut self, color: impl Into<Hsla>) -> Self {
217 self.bearish = Some(color.into());
218 self
219 }
220
221 fn candle_colors(&self, cx: &App) -> (Hsla, Hsla) {
223 (
224 self.bullish.unwrap_or(cx.macos().success),
225 self.bearish.unwrap_or(cx.macos().danger),
226 )
227 }
228
229 fn x_scale(&self, bounds: Bounds<Pixels>) -> Option<ScaleBand<X>> {
232 let x_fn = self.x.as_ref()?;
233 Some(
234 ScaleBand::new(
235 self.data.iter().map(|v| x_fn(v)),
236 [0., bounds.size.width.as_f32()],
237 )
238 .max_band_width(self.max_band_width.as_f32())
239 .padding_inner(0.4)
240 .padding_outer(0.2),
241 )
242 }
243
244 fn plot_height(&self, bounds: Bounds<Pixels>) -> f32 {
246 bounds.size.height.as_f32() - if self.x_axis { AXIS_GAP } else { 0. }
247 }
248}
249
250crate::impl_into_plot!(impl<T, X, Y> CandlestickChart<T, X, Y> where X: Eq + Hash + Into<SharedString> + 'static, Y: PlotValue);
251
252impl<T, X, Y> Plot for CandlestickChart<T, X, Y>
253where
254 X: Eq + Hash + Into<SharedString> + 'static,
255 Y: PlotValue,
256{
257 fn paint(&mut self, bounds: Bounds<Pixels>, window: &mut Window, cx: &mut App) {
258 let (Some(x_fn), Some(open_fn), Some(high_fn), Some(low_fn), Some(close_fn)) = (
259 self.x.as_ref(),
260 self.open.as_ref(),
261 self.high.as_ref(),
262 self.low.as_ref(),
263 self.close.as_ref(),
264 ) else {
265 return;
266 };
267
268 let height = self.plot_height(bounds);
269
270 let Some(x) = self.x_scale(bounds) else {
272 return;
273 };
274 let band_width = x.band_width();
275
276 let all_values = self
278 .data
279 .iter()
280 .flat_map(|d| [high_fn(d), low_fn(d), open_fn(d), close_fn(d)]);
281 let y = ScaleLinear::new(all_values, [height, 10.]);
282
283 let mut axis = PlotAxis::new().stroke(cx.macos().border);
285 if self.x_axis {
286 let labels = build_band_labels(
287 &self.data,
288 x_fn.as_ref(),
289 &x,
290 band_width,
291 &labeled_items(self.data.len(), None, self.tick_margin),
292 cx.macos().secondary,
293 );
294 axis = axis.x(height).x_label(labels);
295 }
296 axis.paint(&bounds, window, cx);
297
298 if self.grid {
300 Grid::new()
301 .y((0..=3).map(|i| height * i as f32 / 4.0))
302 .stroke(cx.macos().divider())
303 .paint(&bounds, window);
304 }
305
306 let (bullish, bearish) = self.candle_colors(cx);
308 let origin = bounds.origin;
309 let x_fn = x_fn.clone();
310 let open_fn = open_fn.clone();
311 let high_fn = high_fn.clone();
312 let low_fn = low_fn.clone();
313 let close_fn = close_fn.clone();
314
315 for d in &self.data {
316 let x_tick = x.tick(&x_fn(d));
317 let Some(x_tick) = x_tick else {
318 continue;
319 };
320
321 let open = open_fn(d);
323 let high = high_fn(d);
324 let low = low_fn(d);
325 let close = close_fn(d);
326
327 let open_y = y.tick(&open);
329 let high_y = y.tick(&high);
330 let low_y = y.tick(&low);
331 let close_y = y.tick(&close);
332
333 let (Some(open_y), Some(high_y), Some(low_y), Some(close_y)) =
334 (open_y, high_y, low_y, close_y)
335 else {
336 continue;
337 };
338
339 let is_bullish = close > open;
341 let color = if is_bullish { bullish } else { bearish };
342
343 let center_x = x_tick + band_width / 2.;
345 let body_width = band_width * self.body_width_ratio;
346 let body_left = center_x - body_width / 2.;
347 let body_right = center_x + body_width / 2.;
348
349 let mut wick_builder = PathBuilder::stroke(px(1.));
351 wick_builder.move_to(origin_point(px(center_x), px(high_y), origin));
352 wick_builder.line_to(origin_point(px(center_x), px(low_y), origin));
353
354 if let Ok(path) = wick_builder.build() {
355 window.paint_path(path, color);
356 }
357
358 let (top, bottom) = if is_bullish {
362 (close_y, open_y)
363 } else {
364 (open_y, close_y)
365 };
366
367 let body_bounds = Bounds::from_corners(
368 origin_point(px(body_left), px(top), origin),
369 origin_point(px(body_right), px(bottom), origin),
370 );
371
372 window.paint_quad(fill(body_bounds, color));
373 }
374 }
375
376 fn id(&self) -> Option<ElementId> {
377 self.interactive.then(|| self.id.clone())
378 }
379
380 fn tooltip_state(
381 &self,
382 position: Point<Pixels>,
383 bounds: Bounds<Pixels>,
384 _cx: &App,
385 ) -> Option<TooltipState> {
386 let x_fn = self.x.as_ref()?;
387 let x = self.x_scale(bounds)?;
388
389 if position.y.as_f32() > self.plot_height(bounds) {
391 return None;
392 }
393
394 let index = x.nearest_index(position.x.as_f32());
395 let d = self.data.get(index)?;
396 let center = x.tick(&x_fn(d))? + x.band_width() / 2.;
397
398 Some(TooltipState::new(
399 index,
400 point(px(center), position.y),
401 vec![],
402 ))
403 }
404
405 fn tooltip(
406 &self,
407 state: &TooltipState,
408 cursor: Point<Pixels>,
409 bounds: Bounds<Pixels>,
410 window: &mut Window,
411 cx: &mut App,
412 ) -> Option<AnyElement> {
413 let (x_fn, open_fn, high_fn, low_fn, close_fn) = (
414 self.x.as_ref()?,
415 self.open.as_ref()?,
416 self.high.as_ref()?,
417 self.low.as_ref()?,
418 self.close.as_ref()?,
419 );
420 let d = self.data.get(state.index)?;
421 let (open, close) = (open_fn(d), close_fn(d));
422 let (bullish, bearish) = self.candle_colors(cx);
423 let color = if close > open { bullish } else { bearish };
424
425 let band_width = self.x_scale(bounds)?.band_width();
429 let cross_line = CrossLine::new(state.cross_line)
430 .span(0., self.plot_height(bounds))
431 .band(px(band_width));
432
433 let tooltip = Tooltip::new(cursor, bounds.size)
434 .gap(px(8.))
435 .cross_line(cross_line);
436 let tooltip = self.tooltip_content.apply(
437 tooltip,
438 d,
439 || Some(x_fn(d).into()),
440 || {
441 [
442 ("Open", open),
443 ("High", high_fn(d)),
444 ("Low", low_fn(d)),
445 ("Close", close),
446 ]
447 .into_iter()
448 .map(|(label, value)| Some((color, label.to_string().into(), value.to_f64()?)))
449 .collect::<Option<Vec<_>>>()
450 },
451 window,
452 cx,
453 )?;
454
455 Some(tooltip.into_any_element())
456 }
457}