ai-app/iris/core/src/ui/render_state.rs

use crate::{
    ActiveData, IdLike, MaskIdx, MoveIdx, Painter, PixelRegion, PrimitiveLayers, RegionAlign,
    StrongWidget, UiRegion, UiRsc, UiVec2, WidgetId, Widgets,
    render::MoveOffset,
    util::{HashMap, HashSet, Id, Vec2},
};

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

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

    /// `Widget::draw` calls and `Primitives::region_mut` rewrites since the
    /// last `take_counters`. LAYOUT.md section 8's pass conditions are
    /// stated in terms of these two: an unchanged frame must cost 0 of
    /// each, and moving one widget must cost 0 draws and 0 rewrites
    /// regardless of how many primitives are in its subtree.
    draw_count: u64,
    region_mut_count: u64,
    mov_count: u64,
}

/// A move chain more than this deep would mean something else is wrong
/// (an accidental cycle) -- see `resolve_move` in shader.wgsl, which walks
/// the identical bound and must be kept in step with this constant.
pub const MOVE_CHAIN_LIMIT: usize = 16;

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(),
            draw_count: 0,
            region_mut_count: 0,
            mov_count: 0,
        }
    }

    /// Reads and zeroes the (draws, region_mut rewrites, move_offsets
    /// writes) counters -- call once per frame before `update()` to
    /// measure exactly that frame, per LAYOUT.md section 8.
    pub fn take_counters(&mut self) -> (u64, u64, u64) {
        (
            std::mem::take(&mut self.draw_count),
            std::mem::take(&mut self.region_mut_count),
            std::mem::take(&mut self.mov_count),
        )
    }

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

    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_redraw_all(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,
                MoveOffset::NONE_PARENT,
                MaskIdx::NONE,
                None,
                None,
                rsc,
            );
        }
    }

    /// The slot an *already-active* widget's `move_offsets` entry chains
    /// to, read back from `self.active`. Only valid where the parent is
    /// guaranteed to already be in `self.active` -- true for `redraw()`,
    /// which targets a widget that was fully drawn on some earlier update,
    /// but **not** for a widget being drawn as part of its own parent's
    /// `Widget::draw` call: that parent's `ActiveData` is not inserted
    /// until its `draw` returns (below), so a child drawn partway through
    /// it would always read back "no parent" here. `Painter::widget_at`
    /// avoids that trap by passing its own already-known `move_slot`
    /// straight through instead of asking `self.active` to look it up.
    fn move_parent_of(&self, parent: Option<WidgetId>) -> u32 {
        parent
            .and_then(|p| self.active.get(&p))
            .map(|p| p.move_slot.idx() as u32)
            .unwrap_or(MoveOffset::NONE_PARENT)
    }

    // 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_slot: u32,
        mask: MaskIdx,
        old_children: Option<Vec<WidgetId>>,
        old_move_slot: Option<MoveIdx>,
        rsc: &mut dyn UiRsc,
    ) {
        let mut old_children = old_children.unwrap_or_default();
        let mut old_move_slot = old_move_slot;
        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, rsc);
                return;
            } else if rsc
                .widgets()
                .get_dyn(id)
                .map(|w| w.is_size_independent())
                .unwrap_or(false)
            {
                // The offered region changed shape, but this widget's own
                // drawn output does not depend on it (a fixed-size leaf) --
                // rewrite its own primitives' regions in place (O(primitives
                // owned directly by this widget, which for a leaf is O(1))
                // instead of redrawing. See LAYOUT.md section 3.
                let from = active.region;
                for h in &active.primitives {
                    let r = self.layers[h.layer].region_mut(h);
                    *r = r.outside(&from).within(&region);
                    self.region_mut_count += 1;
                }
                active.region = 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;
            old_move_slot = Some(active.move_slot);
        }

        // draw widget
        self.draw_started.insert(id);

        let move_slot = match old_move_slot {
            // Reused across a real redraw of the same id: the fresh
            // geometry this draw is about to write is placed at its
            // correct absolute position by `region` itself, so any delta
            // accumulated before this redraw is now stale and would
            // double-offset it if left in place. The chain link (`parent`)
            // is untouched -- the logical parent has not changed.
            Some(slot) => {
                let entry = rsc.ui_mut().move_offsets.get_mut(slot);
                entry.delta = [0.0, 0.0];
                slot
            }
            None => {
                let slot = rsc
                    .ui_mut()
                    .move_offsets
                    .push(MoveOffset::new([0.0, 0.0], parent_move_slot));
                rsc.ui_mut().move_offsets.push_ref(slot);
                if parent_move_slot != MoveOffset::NONE_PARENT {
                    rsc.ui_mut()
                        .move_offsets
                        .push_ref(Id::preset(parent_move_slot));
                }
                slot
            }
        };

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

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

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

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

        // 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);
    }

    /// O(1): write the delta for this widget's own slot in
    /// `move_offsets`. No primitive is touched and there is no recursion --
    /// every descendant's primitive references this slot transitively
    /// through the parent chain the shader walks (`resolve_move`), so it
    /// picks the new delta up for free. See LAYOUT.md section 2.
    fn mov(&mut self, id: WidgetId, from: UiRegion, to: UiRegion, rsc: &mut dyn UiRsc) {
        let Some(active) = self.active.get_mut(&id) else {
            return;
        };
        let slot = active.move_slot;
        active.region = to;
        let from_px = from.top_left().to_abs(self.output_size);
        let to_px = to.top_left().to_abs(self.output_size);
        let delta = to_px - from_px;
        let entry = rsc.ui_mut().move_offsets.get_mut(slot);
        entry.delta[0] += delta.x;
        entry.delta[1] += delta.y;
        self.mov_count += 1;
    }

    /// Move an already-active widget to `to`. Used by `Painter::reposition`,
    /// for a parent that drew a child provisionally (at the whole region it
    /// was offered) and now knows where the child actually belongs.
    ///
    /// Unlike `mov` (called by `draw_inner`'s own dispatch, where the
    /// *offered* region really did move and `active.region` already tracks
    /// it), the child here was not offered a smaller region -- it was
    /// offered everything and chose, on its own, to occupy only
    /// `active.size` of it. By convention every widget in this crate that
    /// does that anchors its own content at the top-left of whatever it
    /// was given (`Rect`/`Image`/`Sized`/`MaxSize` -- see their `draw`
    /// bodies), so that is where this assumes the child was actually
    /// painted, not `active.region` itself (which is the *offered* box,
    /// usually bigger). A nested `Aligned` whose own child is not top-left
    /// anchored -- i.e. `Aligned` wrapping `Aligned` -- is the one shape
    /// this does not cover; none of iris's widgets or examples build that
    /// today. See LAYOUT.md's "Rejected, and why" / deviations for the
    /// full reasoning.
    ///
    /// The delta is overwritten, not accumulated like `mov`'s: `from` is
    /// recomputed fresh from `active.size`/`active.region` every call, so
    /// repeating the same `reposition` (e.g. an unrelated redraw elsewhere
    /// re-running this widget's parent without its own layout changing)
    /// must land on the same answer, not drift further each time.
    pub(super) fn reposition(&mut self, id: WidgetId, to: UiRegion, rsc: &mut dyn UiRsc) {
        let Some(active) = self.active.get(&id) else {
            return;
        };
        let from = active
            .size
            .to_uivec2()
            .align(RegionAlign::TOP_LEFT)
            .within(&active.region);
        let slot = active.move_slot;
        let from_px = from.top_left().to_abs(self.output_size);
        let to_px = to.top_left().to_abs(self.output_size);
        let delta = to_px - from_px;
        let entry = rsc.ui_mut().move_offsets.get_mut(slot);
        entry.delta = [delta.x, delta.y];
        self.mov_count += 1;
    }

    /// 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 {
                // Permanent removal: retire this widget's own move slot
                // (the self-ownership ref taken when it was allocated) and
                // the up-link ref it held on its parent's slot -- read from
                // the arena entry itself, not from `active.parent`, since
                // the parent's own `ActiveData` may already be gone by the
                // time a deep descendant is retired (see LAYOUT.md
                // section 2's lifecycle note).
                let parent_slot = rsc.ui_mut().move_offsets[active.move_slot.idx()].parent;
                rsc.ui_mut().move_offsets.remove(active.move_slot);
                if parent_slot != MoveOffset::NONE_PARENT {
                    rsc.ui_mut().move_offsets.remove(Id::preset(parent_slot));
                }
                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);
            }
        }
        inst
    }

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

    /// What `update` will redraw everything for. Named and shared with
    /// `needs_redraw` rather than written out twice, because the two must
    /// agree: `needs_redraw` is what asks for the frame that `update` would
    /// draw, so a condition in one and not the other is a frame nobody
    /// requests and a stale window. `resized` was missing from `needs_redraw`,
    /// which is latent on Wayland only because winit asks for a redraw after a
    /// resize by itself -- a resize changes neither the root nor any widget,
    /// so nothing else here would have asked.
    fn needs_redraw_all<'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_redraw_all(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), primitives) in self.layers.iter_depth() {
            let indent = " ".repeat(depth * 2);
            let len = primitives.instances().len();
            print!("{indent}{idx}: {len} primitives");
            if len >= 1 {
                print!(" ({})", primitives.instances()[0].binding);
            }
            println!();
        }
    }

    /// `active[id].region`, corrected by every `move_offsets` delta between
    /// `id` and the root -- the CPU-side twin of the vertex shader's chain
    /// walk, over the same arena, so the two cannot disagree about where a
    /// widget is. O(chain depth), not O(primitives). See LAYOUT.md
    /// section 2b.
    pub fn resolved_region(&self, id: &impl IdLike, rsc: &dyn UiRsc) -> Option<UiRegion> {
        let active = self.active.get(&id.id())?;
        let delta = self.resolve_move_chain(active.move_slot, rsc);
        Some(active.region.offset(UiVec2::abs(delta)))
    }

    /// The plain-Rust twin of `resolve_move` in shader.wgsl: sums the
    /// pixel delta along the parent chain starting at `slot`. Both walks
    /// share `MOVE_CHAIN_LIMIT` as their bound so the two cannot disagree
    /// about where the chain ends.
    fn resolve_move_chain(&self, mut slot: MoveIdx, rsc: &dyn UiRsc) -> Vec2 {
        let offsets = &rsc.ui().move_offsets;
        let mut delta = Vec2::ZERO;
        for i in 0..MOVE_CHAIN_LIMIT {
            let entry = &offsets[slot.idx()];
            delta.x += entry.delta[0];
            delta.y += entry.delta[1];
            if entry.parent == MoveOffset::NONE_PARENT {
                return delta;
            }
            slot = Id::preset(entry.parent);
            debug_assert!(
                i + 1 < MOVE_CHAIN_LIMIT,
                "move offset chain exceeded MOVE_CHAIN_LIMIT; a widget's `parent` link is \
                 probably cyclic"
            );
        }
        delta
    }

    pub fn window_region(&self, id: &impl IdLike, rsc: &dyn UiRsc) -> Option<PixelRegion> {
        let region = self.resolved_region(id, rsc)?;
        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);
        if self.draw_started.contains(&id) {
            return;
        }

        let Some(active) = self.remove(id, false, rsc) else {
            return;
        };
        let old_size = active.size;
        let parent = active.parent;
        // `old_move_slot` being `Some` below means the slot is reused in
        // place rather than freshly parented, so this is only reached for
        // logging/clarity's sake, never actually used to link a new slot.
        let parent_move_slot = self.move_parent_of(parent);

        self.draw_inner(
            active.layer,
            id,
            active.region,
            parent,
            parent_move_slot,
            active.mask,
            Some(active.children),
            Some(active.move_slot),
            rsc,
        );

        // If this widget's own reported size changed, its parent's layout
        // (which placed it using the old size) is now stale and needs to
        // relay out too. Checked after the real draw, not before it --
        // there is no query left that answers "what size would this be"
        // without actually drawing (LAYOUT.md section 5).
        if let Some(pid) = parent {
            let new_size = self.active.get(&id).map(|a| a.size);
            if new_size != Some(old_size) {
                self.redraw(pid, rsc);
            }
        }
    }
}

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