iris/src/layout/ui.rs

use image::DynamicImage;

use crate::{
    layout::{
        ActiveSensors, PainterCtx, Sensor, SensorMap, StaticWidgetId, TextData, TextureHandle,
        Textures, Vec2, Widget, WidgetId, WidgetInstance, WidgetLike,
    },
    render::Primitives,
    util::{DynBorrower, HashMap, HashSet, Id, IdTracker},
};
use std::{
    any::{Any, TypeId},
    ops::{Index, IndexMut},
    sync::mpsc::{Receiver, Sender, channel},
};

pub struct Ui<Ctx> {
    root: Option<WidgetId>,
    pub(super) widgets: Widgets,
    updates: Vec<Id>,
    recv: Receiver<Id>,
    pub(super) send: Sender<Id>,
    size: Vec2,
    // TODO: make these non pub(crate)
    pub(crate) primitives: Primitives,
    pub(crate) textures: Textures,
    text: TextData,
    full_redraw: bool,

    pub(super) active: Active,
    pub(super) sensor_map: SensorMap<Ctx>,
}

#[derive(Default)]
pub struct Widgets {
    ids: IdTracker,
    map: HashMap<Id, WidgetData>,
}

pub struct WidgetData {
    pub widget: Box<dyn Widget>,
    pub label: String,
    pub sized_children: HashSet<Id>,
    /// dynamic borrow checking
    pub borrowed: bool,
    /// whether this widget has any sensors
    pub sensor: bool,
}

impl<Ctx> Ui<Ctx> {
    pub fn add<W: Widget, Tag>(&mut self, w: impl WidgetLike<Ctx, Tag, Widget = W>) -> WidgetId<W> {
        w.add(self)
    }

    pub fn add_static<W: Widget, Tag>(
        &mut self,
        w: impl WidgetLike<Ctx, Tag, Widget = W>,
    ) -> StaticWidgetId<W> {
        let id = w.add(self);
        id.into_static()
    }

    /// useful for debugging
    pub fn set_label<W>(&mut self, id: &WidgetId<W>, label: String) {
        self.widgets.data_mut(&id.id).unwrap().label = label;
    }

    pub fn add_widget<W: Widget>(&mut self, w: W) -> WidgetId<W> {
        self.push(w)
    }

    pub fn push<W: Widget>(&mut self, w: W) -> WidgetId<W> {
        let id = self.id();
        self.widgets.insert(id.id.duplicate(), w);
        id
    }

    pub fn set<W: Widget>(&mut self, id: &WidgetId<W>, w: W) {
        self.widgets.insert(id.id.duplicate(), w);
    }

    pub fn set_root<Tag>(&mut self, w: impl WidgetLike<Ctx, Tag>) {
        self.root = Some(w.add(self).any());
        self.full_redraw = true;
    }

    pub fn new() -> Self {
        Self::default()
    }

    pub fn get<W: Widget>(&self, id: &WidgetId<W>) -> Option<&W> {
        self.widgets.get(id)
    }

    pub fn get_mut<W: Widget>(&mut self, id: &WidgetId<W>) -> Option<&mut W> {
        self.widgets.get_mut(id)
    }

    pub fn id<W: Widget>(&mut self) -> WidgetId<W> {
        WidgetId::new(
            self.widgets.reserve(),
            TypeId::of::<W>(),
            self.send.clone(),
            false,
        )
    }

    pub fn id_static<W: Widget>(&mut self) -> StaticWidgetId<W> {
        let id = self.id();
        id.into_static()
    }

    pub fn add_texture(&mut self, image: DynamicImage) -> TextureHandle {
        self.textures.add(image)
    }

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

    pub fn redraw_all(&mut self) {
        self.active.clear();
        // free before bc nothing should exist
        self.free();
        let mut ctx = PainterCtx::new(
            &self.widgets,
            &mut self.primitives,
            &mut self.active,
            &mut self.textures,
            &mut self.text,
            self.size,
        );
        if let Some(root) = &self.root {
            ctx.draw(&root.id);
        }
    }

    pub fn update(&mut self) {
        if self.full_redraw {
            self.redraw_all();
            self.full_redraw = false;
        } else if !self.updates.is_empty() {
            self.redraw_updates();
        }
    }

    fn redraw_updates(&mut self) {
        let mut ctx = PainterCtx::new(
            &self.widgets,
            &mut self.primitives,
            &mut self.active,
            &mut self.textures,
            &mut self.text,
            self.size,
        );
        for id in self.updates.drain(..) {
            ctx.redraw(&id);
        }
        self.free();
    }

    /// free any resources that don't have references anymore
    fn free(&mut self) {
        for id in self.recv.try_iter() {
            self.sensor_map.remove(&id);
            self.widgets.delete(id);
        }
        self.textures.free();
    }

    pub fn needs_redraw(&self) -> bool {
        self.full_redraw || !self.updates.is_empty()
    }

    pub fn num_widgets(&self) -> usize {
        self.widgets.len()
    }

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

    pub fn add_sensor<W>(&mut self, id: &WidgetId<W>, f: Sensor<Ctx>) {
        self.sensor_map
            .entry(id.id.duplicate())
            .or_default()
            .sensors
            .push(f);
        self.widgets.data_mut(&id.id).unwrap().sensor = true;
    }
}

impl<W: Widget, Ctx> Index<&WidgetId<W>> for Ui<Ctx> {
    type Output = W;

    fn index(&self, id: &WidgetId<W>) -> &Self::Output {
        self.get(id).unwrap()
    }
}

impl<W: Widget, Ctx> IndexMut<&WidgetId<W>> for Ui<Ctx> {
    fn index_mut(&mut self, id: &WidgetId<W>) -> &mut Self::Output {
        self.updates.push(id.id.duplicate());
        self.get_mut(id).unwrap()
    }
}

impl<W: Widget, Ctx> Index<StaticWidgetId<W>> for Ui<Ctx> {
    type Output = W;

    fn index(&self, id: StaticWidgetId<W>) -> &Self::Output {
        self.widgets.get_static(&id).unwrap()
    }
}

impl<W: Widget, Ctx> IndexMut<StaticWidgetId<W>> for Ui<Ctx> {
    fn index_mut(&mut self, id: StaticWidgetId<W>) -> &mut Self::Output {
        self.updates.push(id.id.id());
        self.widgets.get_static_mut(&id).unwrap()
    }
}

impl Widgets {
    pub fn new() -> Self {
        Self {
            ids: IdTracker::default(),
            map: HashMap::new(),
        }
    }

    pub fn get_dyn(&self, id: &Id) -> Option<&dyn Widget> {
        Some(self.map.get(id)?.widget.as_ref())
    }

    pub fn get_dyn_mut(&mut self, id: &Id) -> Option<&mut dyn Widget> {
        Some(self.map.get_mut(id)?.widget.as_mut())
    }

    /// get_dyn but dynamic borrow checking of widgets
    /// lets you do recursive (tree) operations, like the painter does
    pub fn get_dyn_dynamic(&self, id: &Id) -> WidgetWrapper<'_> {
        // must guarantee no other mutable references to this widget exist
        // done through the borrow variable
        let data: &mut WidgetData = unsafe { std::mem::transmute(self.map.get(id)) };
        if data.borrowed {
            panic!("tried to mutably borrow the same widget twice");
        }
        WidgetWrapper::new(data.widget.as_mut(), &mut data.borrowed)
    }

    pub fn get_static<W: Widget>(&self, id: &StaticWidgetId<W>) -> Option<&W> {
        self.get_dyn(&id.id.id())?.as_any().downcast_ref()
    }

    pub fn get_static_mut<W: Widget>(&mut self, id: &StaticWidgetId<W>) -> Option<&mut W> {
        self.get_dyn_mut(&id.id.id())?.as_any_mut().downcast_mut()
    }

    pub fn get<W: Widget>(&self, id: &WidgetId<W>) -> Option<&W> {
        self.get_dyn(&id.id)?.as_any().downcast_ref()
    }

    pub fn get_mut<W: Widget>(&mut self, id: &WidgetId<W>) -> Option<&mut W> {
        self.get_dyn_mut(&id.id)?.as_any_mut().downcast_mut()
    }

    pub fn insert<W: Widget>(&mut self, id: Id, widget: W) {
        self.insert_any(id, Box::new(widget), std::any::type_name::<W>().to_string());
    }

    pub fn data(&self, id: &Id) -> Option<&WidgetData> {
        self.map.get(id)
    }

    pub fn data_mut(&mut self, id: &Id) -> Option<&mut WidgetData> {
        self.map.get_mut(id)
    }

    pub fn insert_any(&mut self, id: Id, widget: Box<dyn Widget>, label: String) {
        self.map.insert(
            id,
            WidgetData {
                widget,
                label,
                sized_children: HashSet::new(),
                borrowed: false,
                sensor: false,
            },
        );
    }

    pub fn delete(&mut self, id: Id) {
        self.map.remove(&id);
        self.ids.free(id);
    }

    pub fn reserve(&mut self) -> Id {
        self.ids.next()
    }

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

    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }
}

pub type WidgetWrapper<'a> = DynBorrower<'a, dyn Widget>;

impl dyn Widget {
    pub fn as_any(&self) -> &dyn Any {
        self
    }

    pub fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}

impl<Ctx> Default for Ui<Ctx> {
    fn default() -> Self {
        let (send, recv) = channel();
        Self {
            root: Default::default(),
            widgets: Widgets::new(),
            updates: Default::default(),
            primitives: Default::default(),
            textures: Textures::new(),
            text: TextData::default(),
            full_redraw: false,
            active: Default::default(),
            sensor_map: Default::default(),
            send,
            recv,
            size: Vec2::ZERO,
        }
    }
}

pub type ActiveWidgets = HashMap<Id, WidgetInstance>;

#[derive(Default)]
pub struct Active {
    pub sensors: ActiveSensors,
    pub widgets: ActiveWidgets,
}

impl Active {
    pub fn clear(&mut self) {
        self.sensors.clear();
        self.widgets.clear();
    }

    pub fn remove(&mut self, id: &Id) -> Option<WidgetInstance> {
        let instance = self.widgets.remove(id);
        self.sensors.remove(id);
        instance
    }

    pub fn add(&mut self, id: &Id, instance: WidgetInstance, widgets: &Widgets) {
        if widgets.data(id).is_some_and(|w| w.sensor) {
            self.sensors.insert(id.duplicate(), instance.region);
        }
        self.widgets.insert(id.duplicate(), instance);
    }
}