iris/core/src/ui/render_state.rs

use crate::{
    ActiveData, Axis, DrawLayers, IdLike, MaskIdx, Painter, PixelRegion, SizeCtx, StrongWidget,
    UiRegion, UiRsc, UiVec2, WidgetId, Widgets,
    ui::cache::Cache,
    util::{HashMap, HashSet, Vec2, forget_ref},
};

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

    old_root: Option<WidgetId>,
    resized: bool,
    draw_started: HashSet<WidgetId>,
}

impl UiRenderState {
    pub fn new() -> Self {
        Self {
            active: Default::default(),
            layers: Default::default(),
            cache: Default::default(),
            output_size: Vec2::ZERO,
            old_root: None,
            resized: false,
            draw_started: 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.needs_full_redraw(root) {
            self.redraw_all(root, rsc);
            self.old_root = root.map(|r| r.id());
            self.resized = false;
        } else 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, 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>,
        mask: MaskIdx,
        old_children: Option<Vec<WidgetId>>,
        rsc: &mut dyn UiRsc,
    ) {
        let mut old_children = old_children.unwrap_or_default();
        if let Some(active) = self.active.get_mut(&id)
            && !rsc.widgets().needs_redraw.contains(&id)
        {
            // check to see if we can skip drawing first
            if active.region == region {
                return;
            } else if active.region.size() == region.size() {
                // TODO: epsilon?
                let from = active.region;
                self.mov(id, from, region);
                return;
            }
            // 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
        self.draw_started.insert(id);

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

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

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

        // add to active
        let active = ActiveData {
            id,
            region,
            parent,
            textures,
            primitives,
            children,
            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);
    }

    fn mov(&mut self, id: WidgetId, from: UiRegion, to: UiRegion) {
        let active = self.active.get_mut(&id).unwrap();
        for h in &active.primitives {
            let region = self.layers[h.layer].region_mut(h);
            *region = region.outside(&from).within(&to);
        }
        active.region = active.region.outside(&from).within(&to);
        // SAFETY: children cannot be recursive
        let children = unsafe { forget_ref(&active.children) };
        for child in children {
            self.mov(*child, from, to);
        }
    }

    /// 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> {
        self.cache.remove(id);
        let inst = self.remove(id, true, rsc);
        if let Some(inst) = &inst {
            for c in &inst.children {
                self.remove_rec(*c, rsc);
            }
        }
        inst
    }

    fn clear(&mut self, rsc: &mut dyn UiRsc) {
        for (_, active) in self.active.drain() {
            rsc.on_undraw(&active);
        }
        self.cache.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
    }

    // Scheduling and drawing must use the same full-redraw predicate.
    fn needs_full_redraw<'a>(&self, root: impl Into<Option<&'a StrongWidget>>) -> bool {
        self.root_changed(root) || self.resized
    }

    pub fn needs_redraw<'a>(
        &self,
        root: impl Into<Option<&'a StrongWidget>>,
        widgets: &Widgets,
    ) -> bool {
        self.needs_full_redraw(root) || 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(", "));
        }
    }

    pub fn window_region(&self, id: &impl IdLike) -> Option<PixelRegion> {
        let region = self.active.get(&id.id())?.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) {
        rsc.widgets_mut().needs_redraw.remove(&id);
        self.draw_started.remove(&id);
        // check if parent depends on the desired size of this, if so then redraw it first
        for axis in [Axis::X, Axis::Y] {
            if let Some(&(outer, old)) = self.cache.size.axis_dyn(axis).get(&id)
                && let Some(current) = self.active.get(&id)
                && let Some(pid) = current.parent
            {
                self.cache.size.axis_dyn(axis).remove(&id);
                let new = self.size_ctx(id, outer, rsc).len_axis(id, axis);
                self.cache.size.axis_dyn(axis).insert(id, (outer, new));
                if new != old {
                    self.redraw(pid, rsc);
                }
            }
        }

        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.mask,
            Some(active.children),
            rsc,
        );
    }

    pub(super) fn size_ctx<'b>(
        &'b mut self,
        source: WidgetId,
        outer: UiVec2,
        rsc: &'b mut dyn UiRsc,
    ) -> SizeCtx<'b> {
        let ui = rsc.ui_mut();
        SizeCtx {
            source,
            cache: &mut self.cache,
            text: &mut ui.text,
            widgets: &ui.widgets,
            outer,
            output_size: self.output_size,
            id: source,
        }
    }
}

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