codecraft 0.2.0

//! A frame profiler: what the last few hundred frames were spent on.
//! CPU and GPU spans are charted in separate bands with separate scales.
use std::time::Instant;

use yakui::geometry::{Color as YakuiColor, Constraints, Rect, Vec2};
use yakui::paint::PaintRect;
use yakui::widgets::{List, Pad};
use yakui::{Alignment, CrossAxisAlignment};

use crate::ecs::Resource;
use crate::ui::color::{Color, yakui_color};
use crate::ui::widgets::{self, text_colored};

/// Which side of the machine a span was spent on.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum Lane {
    #[default]
    Cpu,
    Gpu,
}

impl Lane {
    pub const ALL: [Lane; 2] = [Lane::Cpu, Lane::Gpu];

    pub fn label(self) -> &'static str {
        match self {
            Lane::Cpu => "CPU",
            Lane::Gpu => "GPU",
        }
    }

    /// The family a lane's spans are coloured from: warm for CPU, cool for GPU.
    pub fn palette(self) -> &'static [Color] {
        match self {
            Lane::Cpu => &CPU_COLORS,
            Lane::Gpu => &GPU_COLORS,
        }
    }

    /// The colour for the `index`th distinct span in this lane.
    pub fn color(self, index: usize) -> Color {
        let palette = self.palette();
        palette[index % palette.len()]
    }
}

const CPU_COLORS: [Color; 5] = [
    Color::srgb(0.90, 0.55, 0.20),
    Color::srgb(0.85, 0.35, 0.25),
    Color::srgb(0.95, 0.72, 0.30),
    Color::srgb(0.75, 0.28, 0.35),
    Color::srgb(0.98, 0.85, 0.45),
];

const GPU_COLORS: [Color; 5] = [
    Color::srgb(0.25, 0.65, 0.85),
    Color::srgb(0.30, 0.80, 0.70),
    Color::srgb(0.45, 0.50, 0.90),
    Color::srgb(0.20, 0.45, 0.70),
    Color::srgb(0.55, 0.75, 0.95),
];

/// One measured span within a frame. Times are milliseconds from frame start.
#[derive(Clone, Debug, PartialEq)]
pub struct Span {
    pub name: String,
    pub lane: Lane,
    pub start: f64,
    pub end: f64,
}

impl Span {
    pub fn length(&self) -> f64 {
        (self.end - self.start).max(0.0)
    }
}

#[derive(Clone, Default, Debug)]
struct Frame {
    spans: Vec<Span>,
    total: f64,
}

#[derive(Clone, Debug)]
struct Stat {
    name: String,
    worst: f64,
    /// Legend position, kept stable so entries do not jump between frames.
    slot: usize,
}

struct Band {
    lane: Lane,
    frames: Vec<Frame>,
    cursor: usize,
    stats: Vec<Stat>,
    scale: f64,
}

impl Band {
    fn new(lane: Lane, frames: usize) -> Self {
        Self {
            lane,
            frames: vec![Frame::default(); frames.max(1)],
            cursor: 0,
            stats: Vec::new(),
            scale: FRAME_BUDGET,
        }
    }

    fn push(&mut self, spans: Vec<Span>) {
        let total = spans.iter().map(|span| span.end).fold(0.0f64, f64::max);
        let at = self.cursor;
        self.frames[at] = Frame { spans, total };
        self.cursor = (self.cursor + 1) % self.frames.len();

        // Never below the frame budget, so a quiet window still shows the 60 fps line and a spike reads as one.
        let peak = self
            .frames
            .iter()
            .map(|frame| frame.total)
            .fold(0.0, f64::max);
        self.scale = (peak * 1.15).max(FRAME_BUDGET);

        self.rebuild_stats();
    }

    fn current(&self) -> &Frame {
        &self.frames[(self.cursor + self.frames.len() - 1) % self.frames.len()]
    }

    fn total(&self) -> f64 {
        self.current().total
    }

    /// Ordered by worst over the window, not latest, so the legend does not reorder every frame.
    fn rebuild_stats(&mut self) {
        let mut stats: Vec<Stat> = Vec::new();
        for frame in &self.frames {
            for span in &frame.spans {
                match stats.iter_mut().find(|stat| stat.name == span.name) {
                    Some(stat) => stat.worst = stat.worst.max(span.length()),
                    None => stats.push(Stat {
                        name: span.name.clone(),
                        worst: span.length(),
                        slot: 0,
                    }),
                }
            }
        }
        stats.sort_by(|a, b| b.worst.total_cmp(&a.worst));
        for (index, stat) in stats.iter_mut().enumerate() {
            stat.slot = index;
        }
        self.stats = stats;
    }

    fn slot(&self, name: &str) -> usize {
        self.stats
            .iter()
            .find(|stat| stat.name == name)
            .map(|stat| stat.slot)
            .unwrap_or(0)
    }

    fn color(&self, name: &str) -> Color {
        self.lane.color(self.slot(name))
    }
}

/// 60 fps in milliseconds: the line drawn across the chart and the axis floor.
const FRAME_BUDGET: f64 = 1000.0 / 60.0;

const HISTORY: usize = 240;

/// A span being timed; it records itself when dropped.
pub struct Timing<'a> {
    profiler: &'a mut Profiler,
    name: &'static str,
    lane: Lane,
    start: Instant,
}

impl Drop for Timing<'_> {
    fn drop(&mut self) {
        let started = self.profiler.frame_start;
        let from = self.start.duration_since(started).as_secs_f64() * 1000.0;
        let to = self.start.elapsed().as_secs_f64() * 1000.0 + from;
        self.profiler.record(Span {
            name: self.name.to_string(),
            lane: self.lane,
            start: from,
            end: to,
        });
    }
}

/// The profiler: what this frame has cost so far, and the window behind it.
#[derive(Resource)]
pub struct Profiler {
    /// Whether it is running; off by default because timing costs a little.
    pub on: bool,
    bands: Vec<Band>,
    pending: Vec<Span>,
    frame_start: Instant,
}

impl Default for Profiler {
    fn default() -> Self {
        Self::new()
    }
}

impl Profiler {
    pub fn new() -> Self {
        Self {
            on: false,
            bands: Lane::ALL
                .into_iter()
                .map(|lane| Band::new(lane, HISTORY))
                .collect(),
            pending: Vec::new(),
            frame_start: Instant::now(),
        }
    }

    /// Starts a frame. Everything timed after this is measured from here.
    pub fn begin_frame(&mut self) {
        self.frame_start = Instant::now();
        self.pending.clear();
    }

    /// Files this frame's spans into the history.
    pub fn end_frame(&mut self) {
        if !self.on {
            self.pending.clear();
            return;
        }
        let spans = std::mem::take(&mut self.pending);
        for band in &mut self.bands {
            let lane = band.lane;
            band.push(
                spans
                    .iter()
                    .filter(|span| span.lane == lane)
                    .cloned()
                    .collect(),
            );
        }
    }

    /// How long after this frame started `at` was, in milliseconds.
    pub fn since_frame_start(&self, at: Instant) -> f64 {
        at.saturating_duration_since(self.frame_start).as_secs_f64() * 1000.0
    }

    /// Times a stretch of work on one side of the machine.
    ///
    /// ```no_run
    /// # use codecraft::ui::profiler::{Lane, Profiler};
    /// # fn demo(profiler: &mut Profiler) {
    /// let _timing = profiler.time("physics", Lane::Cpu);
    /// // ... the work ...
    /// # }
    /// ```
    pub fn time(&mut self, name: &'static str, lane: Lane) -> Timing<'_> {
        let start = Instant::now();
        Timing {
            profiler: self,
            name,
            lane,
            start,
        }
    }

    /// Files a span measured elsewhere, such as GPU query timings.
    pub fn record(&mut self, span: Span) {
        if self.on {
            self.pending.push(span);
        }
    }

    /// What each lane cost on the last filed frame.
    pub fn totals(&self) -> Vec<(Lane, f64)> {
        self.bands
            .iter()
            .map(|band| (band.lane, band.total()))
            .collect()
    }

    fn band(&self, lane: Lane) -> &Band {
        self.bands
            .iter()
            .find(|band| band.lane == lane)
            .unwrap_or(&self.bands[0])
    }

    /// Draws the chart in the top-right corner when the profiler is on.
    pub fn show(&self) {
        if !self.on {
            return;
        }
        widgets::corner(Alignment::TOP_RIGHT, || {
            widgets::panel_colored(YakuiColor::rgba(10, 13, 20, 224), || {
                yakui::pad(Pad::all(PADDING), || {
                    let mut bands = List::column();
                    bands.item_spacing = BAND_GAP;
                    bands.main_axis_size = yakui::MainAxisSize::Min;
                    bands.show(|| {
                        for lane in Lane::ALL {
                            let band = self.band(lane);
                            let mut column = List::column();
                            column.main_axis_size = yakui::MainAxisSize::Min;
                            column.show(|| {
                                text_colored(
                                    PX,
                                    format!("{} {:.2} MS", lane.label(), band.total()),
                                    yakui_color(lane.color(0)),
                                );
                                let mut row = List::row();
                                row.item_spacing = 18.0;
                                row.main_axis_size = yakui::MainAxisSize::Min;
                                row.cross_axis_alignment = CrossAxisAlignment::Start;
                                row.show(|| {
                                    chart(band);
                                    legend(band);
                                });
                            });
                        }
                    });
                });
            });
        });
    }
}

const WIDTH: f32 = 420.0;
const LEGEND: f32 = 210.0;
const BAND_HEIGHT: f32 = 88.0;
const BAND_GAP: f32 = 22.0;
const PADDING: f32 = 8.0;

const COLUMN: f32 = 2.0;
const COLUMN_GAP: f32 = 0.0;

const ROW: f32 = 13.0;
const SWATCH: f32 = 8.0;
const PX: f32 = 14.0;

/// One stacked column per frame, newest at the right, with the frame budget drawn across.
fn chart(band: &Band) {
    let stride = COLUMN + COLUMN_GAP;
    let columns = ((WIDTH / stride).floor().max(1.0) as usize).min(band.frames.len());
    let mut bars: Vec<(Rect, YakuiColor)> = vec![(
        Rect::from_pos_size(Vec2::ZERO, Vec2::new(WIDTH, BAND_HEIGHT)),
        YakuiColor::rgba(0, 0, 0, 90),
    )];
    // The 60 fps line; it goes off the top once the whole window has been slower than that.
    let budget = 1.0 - (FRAME_BUDGET / band.scale).clamp(0.0, 1.0) as f32;
    bars.push((
        Rect::from_pos_size(Vec2::new(0.0, BAND_HEIGHT * budget), Vec2::new(WIDTH, 1.0)),
        YakuiColor::rgba(140, 148, 168, 140),
    ));
    for column in 0..columns {
        // Walk back from the newest so the latest frame lands at the right-hand edge.
        let age = columns - 1 - column;
        let at = (band.cursor + band.frames.len() - 1 - age) % band.frames.len();
        let x = column as f32 * stride;
        for span in &band.frames[at].spans {
            let from = (span.start / band.scale).clamp(0.0, 1.0) as f32;
            let to = (span.end / band.scale).clamp(0.0, 1.0) as f32;
            let top = BAND_HEIGHT * (1.0 - to);
            let bottom = BAND_HEIGHT * (1.0 - from);
            if bottom - top < 0.5 {
                continue;
            }
            bars.push((
                Rect::from_pos_size(Vec2::new(x, top), Vec2::new(COLUMN, bottom - top)),
                yakui_color(band.color(&span.name)),
            ));
        }
    }

    yakui::constrained(
        Constraints::tight(Vec2::new(WIDTH, BAND_HEIGHT)),
        move || {
            yakui::canvas(move |ctx| {
                let origin = ctx.layout.get(ctx.dom.current()).unwrap().rect.pos();
                for (rect, color) in &bars {
                    let mut paint =
                        PaintRect::new(Rect::from_pos_size(origin + rect.pos(), rect.size()));
                    paint.color = *color;
                    paint.add(ctx.paint);
                }
            });
        },
    );
}

/// One row per span in the newest frame, at its stable slot, with a colour swatch.
fn legend(band: &Band) {
    let rows = (BAND_HEIGHT / ROW).floor().max(1.0) as usize;
    let mut spans: Vec<&Span> = band.current().spans.iter().collect();
    spans.sort_by_key(|span| band.slot(&span.name));
    spans.truncate(rows);

    yakui::constrained(
        Constraints {
            min: Vec2::new(LEGEND, BAND_HEIGHT),
            max: Vec2::new(LEGEND, BAND_HEIGHT),
        },
        || {
            let mut column = List::column();
            column.main_axis_size = yakui::MainAxisSize::Min;
            column.show(|| {
                for span in spans {
                    let color = yakui_color(band.color(&span.name));
                    let mut row = List::row();
                    row.item_spacing = 5.0;
                    row.main_axis_size = yakui::MainAxisSize::Min;
                    row.cross_axis_alignment = CrossAxisAlignment::Center;
                    row.show(|| {
                        yakui::colored_box(color, Vec2::splat(SWATCH));
                        text_colored(PX, format!("{:.2} {}", span.length(), span.name), color);
                    });
                }
            });
        },
    );
}

#[cfg(test)]
mod tests {
    use super::*;

    fn span(name: &str, lane: Lane, start: f64, end: f64) -> Span {
        Span {
            name: name.to_string(),
            lane,
            start,
            end,
        }
    }

    fn running() -> Profiler {
        let mut profiler = Profiler::new();
        profiler.on = true;
        profiler
    }

    #[test]
    fn a_profiler_that_is_off_records_nothing() {
        let mut profiler = Profiler::new();
        profiler.begin_frame();
        profiler.record(span("draw", Lane::Cpu, 0.0, 4.0));
        profiler.end_frame();

        assert_eq!(profiler.totals(), vec![(Lane::Cpu, 0.0), (Lane::Gpu, 0.0)]);
    }

    #[test]
    fn a_frame_is_filed_when_it_ends() {
        let mut profiler = running();
        profiler.begin_frame();
        profiler.record(span("update", Lane::Cpu, 0.0, 3.0));
        profiler.record(span("draw", Lane::Gpu, 0.0, 7.0));
        profiler.end_frame();

        assert_eq!(profiler.totals(), vec![(Lane::Cpu, 3.0), (Lane::Gpu, 7.0)]);
    }

    #[test]
    fn the_lanes_are_kept_apart_and_scaled_apart() {
        let mut profiler = running();
        profiler.begin_frame();
        profiler.record(span("update", Lane::Cpu, 0.0, 1.0));
        profiler.record(span("draw", Lane::Gpu, 0.0, 90.0));
        profiler.end_frame();

        let cpu = profiler.band(Lane::Cpu);
        let gpu = profiler.band(Lane::Gpu);
        assert_eq!(cpu.current().spans.len(), 1, "one span each");
        assert_eq!(gpu.current().spans.len(), 1);
        assert!(
            gpu.scale > cpu.scale * 4.0,
            "a ninety millisecond GPU frame must not set the CPU axis: \
             cpu {} against gpu {}",
            cpu.scale,
            gpu.scale,
        );
    }

    #[test]
    fn the_lanes_are_coloured_apart() {
        for slot in 0..8 {
            assert_ne!(
                Lane::Cpu.color(slot),
                Lane::Gpu.color(slot),
                "slot {slot} is the same colour on both sides",
            );
        }
        for cpu in CPU_COLORS {
            assert!(!GPU_COLORS.contains(&cpu), "{cpu:?} is in both families",);
        }
    }

    #[test]
    fn the_axis_never_hides_the_frame_budget() {
        let mut profiler = running();
        profiler.begin_frame();
        profiler.record(span("tiny", Lane::Cpu, 0.0, 0.01));
        profiler.end_frame();

        assert!(profiler.band(Lane::Cpu).scale >= FRAME_BUDGET);
    }

    #[test]
    fn a_spike_raises_the_axis() {
        let mut profiler = running();
        profiler.begin_frame();
        profiler.record(span("stall", Lane::Cpu, 0.0, 100.0));
        profiler.end_frame();

        assert!(profiler.band(Lane::Cpu).scale > 100.0);
    }

    #[test]
    fn the_legend_orders_by_the_worst_it_has_seen() {
        let mut profiler = running();
        profiler.begin_frame();
        profiler.record(span("fast", Lane::Cpu, 0.0, 1.0));
        profiler.record(span("slow", Lane::Cpu, 1.0, 9.0));
        profiler.end_frame();
        profiler.begin_frame();
        profiler.record(span("fast", Lane::Cpu, 0.0, 2.0));
        profiler.record(span("slow", Lane::Cpu, 2.0, 2.5));
        profiler.end_frame();

        let band = profiler.band(Lane::Cpu);
        assert_eq!(band.slot("slow"), 0);
        assert_eq!(band.slot("fast"), 1);
    }

    #[test]
    fn a_span_keeps_its_colour() {
        let mut profiler = running();
        for _ in 0..3 {
            profiler.begin_frame();
            profiler.record(span("draw", Lane::Cpu, 0.0, 5.0));
            profiler.record(span("update", Lane::Cpu, 5.0, 6.0));
            profiler.end_frame();
        }
        let band = profiler.band(Lane::Cpu);
        assert_eq!(band.color("draw"), Lane::Cpu.color(0));
        assert_eq!(band.color("update"), Lane::Cpu.color(1));
    }

    #[test]
    fn the_history_wraps() {
        let mut profiler = running();
        for i in 0..HISTORY + 10 {
            profiler.begin_frame();
            profiler.record(span("draw", Lane::Cpu, 0.0, i as f64 % 5.0 + 1.0));
            profiler.end_frame();
        }
        assert_eq!(profiler.band(Lane::Cpu).frames.len(), HISTORY);
    }

    #[test]
    fn a_timing_files_itself_when_it_is_dropped() {
        let mut profiler = running();
        profiler.begin_frame();
        {
            let _timing = profiler.time("work", Lane::Cpu);
            std::thread::sleep(std::time::Duration::from_millis(2));
        }
        profiler.end_frame();

        let band = profiler.band(Lane::Cpu);
        assert_eq!(band.current().spans.len(), 1);
        assert!(
            band.total() >= 1.5,
            "it should have measured the sleep: {}",
            band.total(),
        );
    }

    #[test]
    fn it_draws_only_while_it_is_on() {
        use crate::ui::state::Ui;

        let mut ui = Ui::new();
        ui.yakui.set_surface_size(Vec2::new(1280.0, 720.0));
        ui.yakui
            .set_unscaled_viewport(Rect::from_pos_size(Vec2::ZERO, Vec2::new(1280.0, 720.0)));
        let calls = |ui: &mut Ui, profiler: &Profiler| {
            ui.yakui.start();
            profiler.show();
            ui.yakui.finish();
            ui.yakui
                .paint()
                .layers()
                .iter()
                .map(|layer| layer.calls.len())
                .sum::<usize>()
        };

        let mut profiler = running();
        for _ in 0..20 {
            profiler.begin_frame();
            profiler.record(span("huge", Lane::Cpu, 0.0, 4000.0));
            profiler.record(span("draw", Lane::Gpu, 0.0, 8.0));
            profiler.end_frame();
        }
        assert!(calls(&mut ui, &profiler) > 0, "on: something is drawn");

        profiler.on = false;
        assert_eq!(calls(&mut ui, &profiler), 0, "off: nothing at all");
    }
}