Skip to main content

matrix_rain/
widget.rs

1use ratatui::buffer::Buffer;
2use ratatui::layout::Rect;
3use ratatui::style::{Color, Modifier, Style};
4use ratatui::widgets::StatefulWidget;
5
6use crate::config::MatrixConfig;
7use crate::state::MatrixRainState;
8use crate::stream::Stream;
9use crate::theme::ColorRamp;
10
11const TRUECOLOR_SENTINEL: u16 = u16::MAX;
12
13pub struct MatrixRain<'a> {
14    config: &'a MatrixConfig,
15}
16
17impl<'a> MatrixRain<'a> {
18    pub fn new(config: &'a MatrixConfig) -> Self {
19        Self { config }
20    }
21}
22
23impl<'a> StatefulWidget for MatrixRain<'a> {
24    type State = MatrixRainState;
25
26    fn render(self, area: Rect, buf: &mut Buffer, state: &mut Self::State) {
27        state.advance(area, self.config);
28        if area.width == 0 || area.height == 0 {
29            return;
30        }
31
32        if state.color_count().is_none() {
33            state.set_color_count(detect_color_count());
34        }
35        let tier = Tier::from_count(state.color_count().unwrap_or(8));
36
37        let ramp = self.config.theme.ramp();
38        let head_white = self.config.head_white;
39        let bold_head = self.config.bold_head;
40        let background = self.config.background;
41
42        for (col, stream) in state.streams().iter().enumerate() {
43            if !stream.is_active() {
44                continue;
45            }
46            paint_stream(
47                stream, area, buf, &ramp, head_white, bold_head, background, tier, col as u16,
48            );
49        }
50    }
51}
52
53#[derive(Clone, Copy, Debug, PartialEq, Eq)]
54enum Tier {
55    Truecolor,
56    Color256,
57    Color16,
58}
59
60impl Tier {
61    fn from_count(count: u16) -> Self {
62        if count > 256 {
63            Tier::Truecolor
64        } else if count == 256 {
65            Tier::Color256
66        } else {
67            Tier::Color16
68        }
69    }
70}
71
72fn detect_color_count() -> u16 {
73    let truecolor = std::env::var("COLORTERM")
74        .map(|v| matches!(v.trim(), "truecolor" | "24bit"))
75        .unwrap_or(false);
76    if truecolor {
77        TRUECOLOR_SENTINEL
78    } else {
79        crossterm::style::available_color_count()
80    }
81}
82
83fn paint_stream(
84    stream: &Stream,
85    area: Rect,
86    buf: &mut Buffer,
87    ramp: &ColorRamp,
88    head_white: bool,
89    bold_head: bool,
90    background: Option<Color>,
91    tier: Tier,
92    col: u16,
93) {
94    let head_int = stream.head_row().floor() as i32;
95    let length = stream.length();
96    let glyphs = stream.glyphs();
97    let buf_area = buf.area;
98
99    for i in 0..length {
100        let screen_row_i = head_int - i as i32;
101        if screen_row_i < 0 || screen_row_i >= area.height as i32 {
102            continue;
103        }
104        let screen_row = screen_row_i as u16;
105
106        let Some(glyph) = glyphs.get(i as usize).copied() else {
107            continue;
108        };
109
110        let color = pick_color(ramp, head_white, i, length, tier);
111
112        if should_skip(i, length, color, ramp.fade, background) {
113            continue;
114        }
115
116        let Some(x) = area.x.checked_add(col) else {
117            continue;
118        };
119        let Some(y) = area.y.checked_add(screen_row) else {
120            continue;
121        };
122
123        let buf_max_x = buf_area.x.saturating_add(buf_area.width);
124        let buf_max_y = buf_area.y.saturating_add(buf_area.height);
125        if x < buf_area.x || x >= buf_max_x || y < buf_area.y || y >= buf_max_y {
126            continue;
127        }
128
129        let mut style = Style::default().fg(color);
130        if i == 0 && bold_head {
131            style = style.add_modifier(Modifier::BOLD);
132        }
133
134        let cell = buf.get_mut(x, y);
135        cell.set_char(glyph);
136        cell.set_style(style);
137    }
138}
139
140fn pick_color(ramp: &ColorRamp, head_white: bool, i: u16, length: u16, tier: Tier) -> Color {
141    if i == 0 {
142        return if head_white { ramp.head } else { ramp.bright };
143    }
144    let denom = length.saturating_sub(1).max(1);
145    let t = (i as f32) / (denom as f32);
146
147    match tier {
148        Tier::Truecolor => interpolate_smooth(ramp, t),
149        Tier::Color256 => pick_nearest_stop(ramp, t),
150        Tier::Color16 => pick_named_zone(ramp, t),
151    }
152}
153
154fn pick_nearest_stop(ramp: &ColorRamp, t: f32) -> Color {
155    let stops = [ramp.head, ramp.bright, ramp.mid, ramp.dim, ramp.fade];
156    let idx = ((t * 4.0).round() as usize).min(4);
157    stops[idx]
158}
159
160fn interpolate_smooth(ramp: &ColorRamp, t: f32) -> Color {
161    let stops = [ramp.head, ramp.bright, ramp.mid, ramp.dim, ramp.fade];
162    let scaled = (t.clamp(0.0, 1.0)) * 4.0;
163    let lo = (scaled.floor() as usize).min(4);
164    let hi = (lo + 1).min(4);
165    let local = scaled - lo as f32;
166    let (lr, lg, lb) = to_rgb(stops[lo]);
167    let (hr, hg, hb) = to_rgb(stops[hi]);
168    let r = ((1.0 - local) * lr as f32 + local * hr as f32).round() as u8;
169    let g = ((1.0 - local) * lg as f32 + local * hg as f32).round() as u8;
170    let b = ((1.0 - local) * lb as f32 + local * hb as f32).round() as u8;
171    Color::Rgb(r, g, b)
172}
173
174fn pick_named_zone(ramp: &ColorRamp, t: f32) -> Color {
175    let stop = if t < 0.34 {
176        ramp.bright
177    } else if t < 0.67 {
178        ramp.mid
179    } else {
180        ramp.fade
181    };
182    nearest_named(stop)
183}
184
185fn should_skip(i: u16, length: u16, color: Color, fade: Color, background: Option<Color>) -> bool {
186    if let Some(bg) = background {
187        return color == bg;
188    }
189    if i == 0 {
190        return false;
191    }
192    if color == fade {
193        return true;
194    }
195    let denom = length.saturating_sub(1).max(1);
196    let t = (i as f32) / (denom as f32);
197    t >= 0.875
198}
199
200fn to_rgb(c: Color) -> (u8, u8, u8) {
201    match c {
202        Color::Rgb(r, g, b) => (r, g, b),
203        Color::Black => (0, 0, 0),
204        Color::Red => (128, 0, 0),
205        Color::Green => (0, 128, 0),
206        Color::Yellow => (128, 128, 0),
207        Color::Blue => (0, 0, 128),
208        Color::Magenta => (128, 0, 128),
209        Color::Cyan => (0, 128, 128),
210        Color::Gray => (192, 192, 192),
211        Color::DarkGray => (128, 128, 128),
212        Color::LightRed => (255, 0, 0),
213        Color::LightGreen => (0, 255, 0),
214        Color::LightYellow => (255, 255, 0),
215        Color::LightBlue => (0, 0, 255),
216        Color::LightMagenta => (255, 0, 255),
217        Color::LightCyan => (0, 255, 255),
218        Color::White => (255, 255, 255),
219        Color::Indexed(_) | Color::Reset => (255, 255, 255),
220    }
221}
222
223const NAMED_PALETTE: &[(Color, (u8, u8, u8))] = &[
224    (Color::Black, (0, 0, 0)),
225    (Color::Red, (128, 0, 0)),
226    (Color::Green, (0, 128, 0)),
227    (Color::Yellow, (128, 128, 0)),
228    (Color::Blue, (0, 0, 128)),
229    (Color::Magenta, (128, 0, 128)),
230    (Color::Cyan, (0, 128, 128)),
231    (Color::Gray, (192, 192, 192)),
232    (Color::DarkGray, (128, 128, 128)),
233    (Color::LightRed, (255, 0, 0)),
234    (Color::LightGreen, (0, 255, 0)),
235    (Color::LightYellow, (255, 255, 0)),
236    (Color::LightBlue, (0, 0, 255)),
237    (Color::LightMagenta, (255, 0, 255)),
238    (Color::LightCyan, (0, 255, 255)),
239    (Color::White, (255, 255, 255)),
240];
241
242fn nearest_named(target: Color) -> Color {
243    let (tr, tg, tb) = to_rgb(target);
244    let mut best = NAMED_PALETTE[0].0;
245    let mut best_dist = u32::MAX;
246    for &(named, (nr, ng, nb)) in NAMED_PALETTE {
247        let dr = (tr as i32 - nr as i32).unsigned_abs();
248        let dg = (tg as i32 - ng as i32).unsigned_abs();
249        let db = (tb as i32 - nb as i32).unsigned_abs();
250        let dist = dr * dr + dg * dg + db * db;
251        if dist < best_dist {
252            best_dist = dist;
253            best = named;
254        }
255    }
256    best
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    fn fully_active_config(seed_density: f32) -> MatrixConfig {
264        MatrixConfig {
265            density: seed_density,
266            ..MatrixConfig::default()
267        }
268    }
269
270    fn classic_ramp() -> ColorRamp {
271        ColorRamp {
272            head: Color::Rgb(0xFF, 0xFF, 0xFF),
273            bright: Color::Rgb(0xCC, 0xFF, 0xCC),
274            mid: Color::Rgb(0x00, 0xFF, 0x00),
275            dim: Color::Rgb(0x00, 0x99, 0x00),
276            fade: Color::Rgb(0x00, 0x33, 0x00),
277        }
278    }
279
280    #[test]
281    fn render_with_zero_width_area_is_noop() {
282        let cfg = MatrixConfig::default();
283        let mut state = MatrixRainState::with_seed(0);
284        let mut buf = Buffer::empty(Rect::new(0, 0, 10, 10));
285        MatrixRain::new(&cfg).render(Rect::new(0, 0, 0, 10), &mut buf, &mut state);
286    }
287
288    #[test]
289    fn render_with_zero_height_area_is_noop() {
290        let cfg = MatrixConfig::default();
291        let mut state = MatrixRainState::with_seed(0);
292        let mut buf = Buffer::empty(Rect::new(0, 0, 10, 10));
293        MatrixRain::new(&cfg).render(Rect::new(0, 0, 10, 0), &mut buf, &mut state);
294    }
295
296    #[test]
297    fn does_not_paint_outside_widget_area() {
298        let cfg = fully_active_config(1.0);
299        let mut state = MatrixRainState::with_seed(42);
300        let buf_area = Rect::new(0, 0, 20, 20);
301        let mut buf = Buffer::empty(buf_area);
302        for y in 0..20 {
303            for x in 0..20 {
304                buf.get_mut(x, y).set_char('#');
305            }
306        }
307        let widget_area = Rect::new(5, 5, 10, 10);
308        for _ in 0..50 {
309            MatrixRain::new(&cfg).render(widget_area, &mut buf, &mut state);
310            state.tick();
311        }
312        for y in 0..20 {
313            for x in 0..20 {
314                let inside = (5..15).contains(&x) && (5..15).contains(&y);
315                if !inside {
316                    assert_eq!(
317                        buf.get(x, y).symbol(),
318                        "#",
319                        "cell ({x},{y}) outside widget area was modified"
320                    );
321                }
322            }
323        }
324    }
325
326    #[test]
327    fn paints_at_least_some_cells_with_high_density() {
328        let cfg = fully_active_config(1.0);
329        let mut state = MatrixRainState::with_seed(42);
330        let mut buf = Buffer::empty(Rect::new(0, 0, 20, 20));
331        let widget_area = Rect::new(0, 0, 20, 20);
332        for _ in 0..120 {
333            MatrixRain::new(&cfg).render(widget_area, &mut buf, &mut state);
334            state.tick();
335        }
336        let mut painted = 0;
337        for y in 0..20 {
338            for x in 0..20 {
339                let sym = buf.get(x, y).symbol();
340                if !sym.is_empty() && sym != " " {
341                    painted += 1;
342                }
343            }
344        }
345        assert!(painted > 0, "expected some cells to be painted");
346    }
347
348    #[test]
349    fn honors_non_zero_origin() {
350        let cfg = fully_active_config(1.0);
351        let mut state = MatrixRainState::with_seed(42);
352        let mut buf = Buffer::empty(Rect::new(0, 0, 30, 30));
353        let widget_area = Rect::new(7, 11, 8, 8);
354        for _ in 0..120 {
355            MatrixRain::new(&cfg).render(widget_area, &mut buf, &mut state);
356            state.tick();
357        }
358        let mut painted_inside = 0;
359        let mut painted_outside = 0;
360        for y in 0..30 {
361            for x in 0..30 {
362                let sym = buf.get(x, y).symbol();
363                if !sym.is_empty() && sym != " " {
364                    let inside = (7..15).contains(&x) && (11..19).contains(&y);
365                    if inside {
366                        painted_inside += 1;
367                    } else {
368                        painted_outside += 1;
369                    }
370                }
371            }
372        }
373        assert!(painted_inside > 0, "no cells painted inside offset area");
374        assert_eq!(painted_outside, 0, "cells painted outside offset area");
375    }
376
377    #[test]
378    fn resize_between_renders_does_not_panic() {
379        let cfg = MatrixConfig::default();
380        let mut state = MatrixRainState::with_seed(42);
381        let sizes = [(20u16, 20u16), (5, 30), (40, 5), (1, 1), (0, 10), (15, 15)];
382        for (w, h) in sizes {
383            let mut buf = Buffer::empty(Rect::new(0, 0, w.max(1), h.max(1)));
384            MatrixRain::new(&cfg).render(Rect::new(0, 0, w, h), &mut buf, &mut state);
385        }
386    }
387
388    #[test]
389    fn tier_from_count_buckets() {
390        assert_eq!(Tier::from_count(8), Tier::Color16);
391        assert_eq!(Tier::from_count(15), Tier::Color16);
392        assert_eq!(Tier::from_count(16), Tier::Color16);
393        assert_eq!(Tier::from_count(255), Tier::Color16);
394        assert_eq!(Tier::from_count(256), Tier::Color256);
395        assert_eq!(Tier::from_count(257), Tier::Truecolor);
396        assert_eq!(Tier::from_count(u16::MAX), Tier::Truecolor);
397    }
398
399    #[test]
400    fn nearest_stop_endpoints() {
401        let r = classic_ramp();
402        assert_eq!(pick_nearest_stop(&r, 0.0), r.head);
403        assert_eq!(pick_nearest_stop(&r, 1.0), r.fade);
404        assert_eq!(pick_nearest_stop(&r, 0.5), r.mid);
405    }
406
407    #[test]
408    fn smooth_interpolation_endpoints_match_stops() {
409        let r = classic_ramp();
410        assert_eq!(interpolate_smooth(&r, 0.0), r.head);
411        assert_eq!(interpolate_smooth(&r, 1.0), r.fade);
412        assert_eq!(interpolate_smooth(&r, 0.25), r.bright);
413        assert_eq!(interpolate_smooth(&r, 0.5), r.mid);
414        assert_eq!(interpolate_smooth(&r, 0.75), r.dim);
415    }
416
417    #[test]
418    fn smooth_interpolation_midpoint_is_between_stops() {
419        let r = classic_ramp();
420        // t=0.125 sits between head (white) and bright (pale green).
421        match interpolate_smooth(&r, 0.125) {
422            Color::Rgb(rr, gg, bb) => {
423                // Should be between head (255,255,255) and bright (204,255,204).
424                assert!(rr > 204 && rr < 255, "r out of range: {rr}");
425                assert_eq!(gg, 255);
426                assert!(bb > 204 && bb < 255, "b out of range: {bb}");
427            }
428            _ => panic!("expected Rgb"),
429        }
430    }
431
432    #[test]
433    fn named_zone_collapses_to_named_colors() {
434        let r = classic_ramp();
435        // bright zone (early trail): 0xCCFFCC is closest to LightGreen (0,255,0)... actually it's pale green.
436        let early = pick_named_zone(&r, 0.1);
437        let mid = pick_named_zone(&r, 0.5);
438        let late = pick_named_zone(&r, 0.9);
439        // All should be one of the 16 named variants (no Rgb).
440        for c in [early, mid, late] {
441            assert!(
442                !matches!(c, Color::Rgb(..) | Color::Indexed(..)),
443                "Color16 path returned non-named color: {c:?}"
444            );
445        }
446    }
447
448    #[test]
449    fn nearest_named_white_for_white_input() {
450        assert_eq!(nearest_named(Color::Rgb(0xFF, 0xFF, 0xFF)), Color::White);
451        assert_eq!(nearest_named(Color::Rgb(0x00, 0x00, 0x00)), Color::Black);
452        assert_eq!(nearest_named(Color::Rgb(0x00, 0xFF, 0x00)), Color::LightGreen);
453    }
454
455    #[test]
456    fn pick_color_head_respects_head_white() {
457        let r = classic_ramp();
458        for tier in [Tier::Truecolor, Tier::Color256, Tier::Color16] {
459            assert_eq!(pick_color(&r, true, 0, 10, tier), r.head);
460            assert_eq!(pick_color(&r, false, 0, 10, tier), r.bright);
461        }
462    }
463
464    #[test]
465    fn skip_when_color_matches_background() {
466        let r = classic_ramp();
467        assert!(should_skip(3, 10, Color::Black, r.fade, Some(Color::Black)));
468        assert!(!should_skip(3, 10, Color::Green, r.fade, Some(Color::Black)));
469    }
470
471    #[test]
472    fn skip_fade_zone_when_background_none() {
473        let r = classic_ramp();
474        // Tail cell (i=length-1, t=1.0) should be skipped when background is None.
475        assert!(should_skip(9, 10, r.fade, r.fade, None));
476        // Head cell never skipped.
477        assert!(!should_skip(0, 10, r.head, r.fade, None));
478        // Middle cell not skipped.
479        assert!(!should_skip(4, 10, r.mid, r.fade, None));
480    }
481
482    #[test]
483    fn detection_caches_into_state_after_first_render() {
484        let cfg = MatrixConfig::default();
485        let mut state = MatrixRainState::with_seed(0);
486        assert!(state.color_count().is_none());
487        let mut buf = Buffer::empty(Rect::new(0, 0, 5, 5));
488        MatrixRain::new(&cfg).render(Rect::new(0, 0, 5, 5), &mut buf, &mut state);
489        assert!(state.color_count().is_some());
490    }
491
492    #[test]
493    fn detection_does_not_overwrite_pre_set_count() {
494        let cfg = MatrixConfig::default();
495        let mut state = MatrixRainState::with_seed(0);
496        state.set_color_count(42);
497        let mut buf = Buffer::empty(Rect::new(0, 0, 5, 5));
498        MatrixRain::new(&cfg).render(Rect::new(0, 0, 5, 5), &mut buf, &mut state);
499        assert_eq!(state.color_count(), Some(42));
500    }
501
502    #[test]
503    fn renders_under_each_tier_without_panic() {
504        let cfg = fully_active_config(1.0);
505        for forced in [16u16, 256, TRUECOLOR_SENTINEL] {
506            let mut state = MatrixRainState::with_seed(0xBEEF);
507            state.set_color_count(forced);
508            let mut buf = Buffer::empty(Rect::new(0, 0, 20, 10));
509            for _ in 0..30 {
510                MatrixRain::new(&cfg).render(Rect::new(0, 0, 20, 10), &mut buf, &mut state);
511                state.tick();
512            }
513        }
514    }
515}