iris/core/src/ui/render_state.rs

use crate::{
    ActiveData, Axis, DrawLayers, IdLike, MaskIdx, MoveIdx, Moves, OnResize, Painter, PixelRegion,
    Size, StrongWidget, UiRegion, UiRsc, WidgetId, Widgets,
    util::{HashMap, HashSet, Vec2},
};

const AXES: [Axis; 2] = [Axis::X, Axis::Y];

pub struct UiRenderState {
    pub active: HashMap<WidgetId, ActiveData>,
    pub layers: DrawLayers,
    pub(super) output_size: Vec2,

    old_root: Option<WidgetId>,
    resized: bool,
    draw_started: HashSet<WidgetId>,
    /// A widget's move slot, which outlives any one `ActiveData`: a redraw
    /// replaces that while its children go on pointing at the slot.
    slots: HashMap<WidgetId, MoveIdx>,
    pub moves: Moves,
}

impl UiRenderState {
    pub fn new() -> Self {
        Self {
            active: Default::default(),
            layers: Default::default(),
            output_size: Vec2::ZERO,
            old_root: None,
            resized: false,
            draw_started: Default::default(),
            slots: Default::default(),
            moves: Default::default(),
        }
    }

    pub fn resize(&mut self, size: impl Into<Vec2>) {
        self.output_size = size.into();
        self.resized = true;
    }

    pub fn output_size(&self) -> Vec2 {
        self.output_size
    }

    pub fn update<'a>(&mut self, root: impl Into<Option<&'a StrongWidget>>, rsc: &mut dyn UiRsc) {
        // safety mechanism for memory leaks; might wanna return a result instead so user can
        // decide whether to panic or not
        if !rsc.widgets().waiting.is_empty() {
            let widgets = rsc.widgets();
            let len = widgets.waiting.len();
            let all: Vec<_> = widgets
                .waiting
                .iter()
                .map(|&w| format!("'{}' ({w:?})", widgets.label(w)))
                .collect();
            panic!(
                "{len} widget(s) were never upgraded\n\
                this is likely a memory leak; consider upgrading to strong if you plan on using it later\n\
                weak widgets: {all:#?}"
            );
        }
        let root = root.into();
        if self.root_changed(root) {
            self.redraw_all(root, rsc);
            self.old_root = root.map(|r| r.id());
        } else if self.resized {
            // A region is a fraction of the output plus an offset, resolved
            // against the window in the shader, so a resize moves the whole
            // drawing on its own. Only a widget that read pixels can be wrong.
            for (&id, active) in &self.active {
                if active.reads_output {
                    rsc.widgets_mut().needs_redraw.insert(id);
                }
            }
        }
        self.resized = false;
        if rsc.widgets().has_updates() {
            self.redraw_updates(rsc);
        }
    }

    fn redraw_all(&mut self, root: Option<&StrongWidget>, rsc: &mut dyn UiRsc) {
        self.clear(rsc);
        // free all resources & cache
        if let Some(id) = root {
            self.draw_inner(
                0,
                id.id(),
                UiRegion::FULL,
                None,
                MoveIdx::NONE,
                false,
                MaskIdx::NONE,
                None,
                rsc,
            );
        }
    }

    // TODO: should prolly make a DrawInfo struct or smth for everything other than rsc
    #[allow(clippy::too_many_arguments)]
    pub(super) fn draw_inner(
        &mut self,
        layer: usize,
        id: WidgetId,
        region: UiRegion,
        parent: Option<WidgetId>,
        parent_move: MoveIdx,
        slotted: bool,
        mask: MaskIdx,
        old_children: Option<Vec<WidgetId>>,
        rsc: &mut dyn UiRsc,
    ) -> Size {
        let mut old_children = old_children.unwrap_or_default();
        if self.active.contains_key(&id) {
            if let Some(size) = self.try_reuse(id, region, parent_move, rsc) {
                return size;
            }
            // if not, then maintain resize and track old children to remove unneeded
            let active = self.remove(id, false, rsc).unwrap();
            old_children = active.children;
        }

        // draw widget
        let (move_idx, local) = match slotted {
            // Its box becomes its slot's, so it draws in the slot's own
            // coordinates and the box it was given is one entry to rewrite.
            true => (self.move_slot(id, parent_move, region), UiRegion::FULL),
            false => {
                self.drop_slot(id);
                (parent_move, region)
            }
        };
        rsc.widgets_mut().needs_redraw.remove(&id);
        self.draw_started.insert(id);

        let mut painter = Painter {
            state: self,
            region: local,
            mask,
            layer,
            id,
            textures: Vec::new(),
            primitives: Vec::new(),
            children: Vec::new(),
            size_deps: Vec::new(),
            reads_output: false,
            move_idx,
            rsc,
        };

        let mut widget = painter.rsc.widgets().get_dyn_dynamic(id);
        let size = widget.draw(&mut painter);
        drop(widget);

        let Painter {
            state: _,
            rsc: _,
            region: _,
            mask,
            textures,
            primitives,
            children,
            size_deps,
            reads_output,
            move_idx,
            layer,
            id,
        } = painter;

        debug_assert!(
            Self::hints_agree(id, size, rsc),
            "'{}' ({id:?}) drew a size its size_hint disagrees with",
            rsc.widgets().label(id)
        );

        // add to active
        let active = ActiveData {
            id,
            region,
            size,
            parent,
            textures,
            primitives,
            children,
            size_deps,
            reads_output,
            move_idx,
            parent_move,
            mask,
            layer,
        };

        // remove old children that weren't kept
        for c in &old_children {
            if !active.children.contains(c) {
                self.remove_rec(*c, rsc);
            }
        }

        rsc.on_draw(&active);
        self.active.insert(id, active);
        size
    }

    /// The slot a widget's box is held in, made on its first placed draw and
    /// kept until it stops being drawn -- a redraw replaces its `ActiveData`
    /// while descendants go on naming the slot.
    fn move_slot(&mut self, id: WidgetId, parent: MoveIdx, region: UiRegion) -> MoveIdx {
        if let Some(&idx) = self.slots.get(&id) {
            self.moves.set_parent(idx, parent);
            self.moves.set(idx, region);
            return idx;
        }
        let idx = self.moves.push(parent, region);
        self.slots.insert(id, idx);
        idx
    }

    /// Gives up a slot a widget no longer needs, because it is drawn somewhere
    /// that does not place it. Its descendants name it, so this is only
    /// reached where they are about to be drawn again.
    fn drop_slot(&mut self, id: WidgetId) {
        if let Some(idx) = self.slots.remove(&id) {
            self.moves.remove(idx);
        }
    }

    /// The drawing a widget already has, kept for a new box if the box has not
    /// changed in a way it depends on.
    fn try_reuse(
        &mut self,
        id: WidgetId,
        region: UiRegion,
        parent_move: MoveIdx,
        rsc: &mut dyn UiRsc,
    ) -> Option<Size> {
        if rsc.widgets().needs_redraw.contains(&id) {
            return None;
        }
        let active = self.active.get(&id)?;
        // Drawn somewhere else in the tree: its box is in coordinates it no
        // longer sits in, and its slot names the wrong parent.
        if active.parent_move != parent_move {
            return None;
        }
        let (size, old, slot) = (active.size, active.region, active.move_idx);
        if old == region {
            return Some(size);
        }
        // Only a placed widget can be given a different box without drawing
        // again: everything it drew is a fraction of its slot's box, so one
        // entry says where all of it went.
        if slot == parent_move {
            return None;
        }
        let mut changed = [false; 2];
        for (axis, c) in AXES.into_iter().zip(changed.iter_mut()) {
            *c = region.axis(axis).len() != old.axis(axis).len();
        }
        if changed.iter().any(|&c| c) {
            let widget = rsc.widgets().get_dyn(id)?;
            let redraws = AXES
                .into_iter()
                .zip(changed)
                .any(|(axis, c)| c && widget.on_resize(axis) != OnResize::Scale);
            if redraws {
                return None;
            }
        }
        // Anything under it that has to be drawn again is drawn by drawing
        // this, because whatever reads that widget's size sits in between and
        // has to lay out around whatever it comes to.
        if changed.iter().any(|&c| c) && self.redraws_under(id, changed, rsc) {
            return None;
        }
        self.moves.set(slot, region);
        self.active.get_mut(&id).unwrap().region = region;
        Some(size)
    }

    /// Whether anything under `id` would have to be drawn again for the box
    /// it is a fraction of changing length, `changed` saying which axes of
    /// that box did.
    ///
    /// A part of a box with no relative extent on an axis is a fixed length,
    /// held as offsets from that box's start, and composing anything into it
    /// leaves no relative extent either. So a widget whose own box did not
    /// change length has no descendant whose box did, and the walk stops
    /// there -- an 80-wide child of a widened row is not asked at all.
    fn redraws_under(&self, id: WidgetId, changed: [bool; 2], rsc: &dyn UiRsc) -> bool {
        let Some(active) = self.active.get(&id) else {
            return false;
        };
        active.children.iter().any(|&child| {
            let Some(data) = self.active.get(&child) else {
                return false;
            };
            let mut own = changed;
            for (axis, c) in AXES.into_iter().zip(own.iter_mut()) {
                *c &= data.region.axis(axis).len().rel != 0.0;
            }
            if !own.iter().any(|&c| c) {
                return false;
            }
            let redraws = match rsc.widgets().get_dyn(child) {
                Some(widget) => AXES
                    .into_iter()
                    .zip(own)
                    .any(|(axis, c)| c && widget.on_resize(axis) != OnResize::Scale),
                None => true,
            };
            redraws || self.redraws_under(child, own, rsc)
        })
    }

    fn hints_agree(id: WidgetId, size: Size, rsc: &dyn UiRsc) -> bool {
        let Some(widget) = rsc.widgets().get_dyn(id) else {
            return true;
        };
        AXES.into_iter().all(|axis| {
            widget
                .size_hint(axis)
                .is_none_or(|hint| hint == size.axis(axis))
        })
    }

    /// NOTE: instance textures are cleared and self.textures freed
    fn remove(&mut self, id: WidgetId, undraw: bool, rsc: &mut dyn UiRsc) -> Option<ActiveData> {
        let mut active = self.active.remove(&id);
        if let Some(active) = &mut active {
            for h in &active.primitives {
                let mask = self.layers.free(h);
                if mask != MaskIdx::NONE {
                    rsc.ui_mut().masks.remove(mask);
                }
            }
            active.textures.clear();
            rsc.ui_mut().textures.free();
            if undraw {
                rsc.on_undraw(active);
            }
        }
        active
    }

    fn remove_rec(&mut self, id: WidgetId, rsc: &mut dyn UiRsc) -> Option<ActiveData> {
        let inst = self.remove(id, true, rsc);
        if let Some(inst) = &inst {
            for c in &inst.children {
                self.remove_rec(*c, rsc);
            }
        }
        // After the descendants, whose slots name this one as their parent.
        if let Some(idx) = self.slots.remove(&id) {
            self.moves.remove(idx);
        }
        inst
    }

    fn clear(&mut self, rsc: &mut dyn UiRsc) {
        for (_, active) in self.active.drain() {
            rsc.on_undraw(&active);
        }
        self.slots.clear();
        self.moves.clear();
        self.layers.clear();
        rsc.widgets_mut().needs_redraw.clear();
        rsc.free();
    }

    pub fn redraw_updates(&mut self, rsc: &mut dyn UiRsc) {
        while let Some(&id) = rsc.widgets().needs_redraw.iter().next() {
            self.redraw(id, rsc);
        }
        rsc.free();
    }

    pub fn root_changed<'a>(&self, root: impl Into<Option<&'a StrongWidget>>) -> bool {
        root.into().map(|r| r.id()) != self.old_root
    }

    pub fn needs_redraw<'a>(
        &self,
        root: impl Into<Option<&'a StrongWidget>>,
        widgets: &Widgets,
    ) -> bool {
        self.root_changed(root) || self.resized || widgets.has_updates()
    }

    pub fn active_widgets(&self) -> usize {
        self.active.len()
    }

    pub fn debug(&self, widgets: &Widgets, label: &str) -> impl Iterator<Item = &ActiveData> {
        self.active.iter().filter_map(move |(&id, inst)| {
            let l = widgets.label(id);
            if l == label { Some(inst) } else { None }
        })
    }

    pub fn debug_layers(&self) {
        for ((idx, depth), draws) in self.layers.iter_depth() {
            let indent = " ".repeat(depth * 2);
            let counts: Vec<String> = draws
                .primitives()
                .iter()
                .map(|l| l.as_ref().map_or(0, |l| l.instances().len()).to_string())
                .collect();
            println!("{indent}{idx}: [{}]", counts.join(", "));
        }
    }

    /// Where a widget is on screen: its box composed through the boxes it
    /// sits within, which is the walk the vertex shader does.
    pub fn window_region(&self, id: &impl IdLike) -> Option<PixelRegion> {
        let active = self.active.get(&id.id())?;
        let region = self.moves.resolve(active.parent_move, active.region);
        Some(region.to_px(self.output_size))
    }

    /// redraws a widget that's currently active (drawn)
    pub fn redraw(&mut self, id: WidgetId, rsc: &mut dyn UiRsc) {
        self.draw_started.remove(&id);
        // Whoever read this widget's size may be a different size now, so the
        // highest reader is what draws. Everything between the two is marked
        // as well: their own boxes have not changed, so the mark is the only
        // thing stopping the draw reusing its way past this widget.
        if let Some(top) = self.mark_readers(id, rsc) {
            self.redraw(top, rsc);
            // Cleared by that draw if it reached here; if it did not, this is
            // no longer drawn and asking again would not end.
            rsc.widgets_mut().needs_redraw.remove(&id);
            return;
        }
        rsc.widgets_mut().needs_redraw.remove(&id);

        if self.draw_started.contains(&id) {
            return;
        }

        let Some(active) = self.remove(id, false, rsc) else {
            return;
        };

        self.draw_inner(
            active.layer,
            id,
            active.region,
            active.parent,
            active.parent_move,
            active.move_idx != active.parent_move,
            active.mask,
            Some(active.children),
            rsc,
        );
    }

    /// The furthest ancestor that read this widget's size, directly or through
    /// widgets that did the same, marking everything below it on the way.
    fn mark_readers(&self, id: WidgetId, rsc: &mut dyn UiRsc) -> Option<WidgetId> {
        let mut top = None;
        let mut at = id;
        while let Some(active) = self.active.get(&at)
            && let Some(parent) = active.parent
            && self
                .active
                .get(&parent)
                .is_some_and(|p| p.size_deps.contains(&at))
        {
            rsc.widgets_mut().needs_redraw.insert(at);
            top = Some(parent);
            at = parent;
        }
        top
    }
}

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