iris/core/src/layout/widget/handle.rs

use std::{marker::Unsize, mem::MaybeUninit, ops::CoerceUnsized};

use crate::{
    layout::{EventManager, IdFnTag, IdTag, Ui, WIDGETS, Widget, WidgetData, WidgetLike},
    util::{
        Handle, MappedRefGuard, MappedRefGuardMut, Ref, RefGuard, RefGuardMut, RefMap, RefMapMut,
        RefMut, WeakHandle,
    },
};

pub(super) type SlotTy = u32;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct WidgetId {
    pub(super) i: SlotTy,
    pub(super) genr: SlotTy,
}

/// A strong handle to a widget. Cannot be cloned, only one may exist
/// Use .weak() to obtain weak references to this
pub struct WidgetHandle<W: ?Sized = dyn Widget>(WidgetId, Handle<WidgetData<W>>);
/// A weak handle to a widget. Implements Copy so you can easily pass into closures
pub struct WidgetRef<W: ?Sized = dyn Widget>(WidgetId, *const W);

pub type WidgetUpdate = WidgetId;

impl<W: Widget> WidgetHandle<W> {
    pub(super) fn new(id: WidgetId, widget: W) -> Self {
        let mut label = std::any::type_name::<W>().to_string();
        if let (Some(first), Some(last)) = (label.find(":"), label.rfind(":")) {
            label = label.split_at(first).0.to_string() + "::" + label.split_at(last + 1).1;
        }
        Self(
            id,
            WidgetData {
                id,
                widget,
                send: None,
                label,
                event_managers: Default::default(),
            }
            .into(),
        )
    }
}

impl<W: Widget + ?Sized + Unsize<dyn Widget>> WidgetHandle<W> {
    pub(super) fn weak_inner(&self) -> WeakHandle<WidgetData<W>> {
        self.1.weak()
    }

    pub fn any(self) -> WidgetHandle<dyn Widget> {
        self
    }

    /// DO NOT CALL OTHER THAN TEMP IN PAINTER
    pub(crate) fn as_any(&self) -> WidgetHandle<dyn Widget> {
        WidgetHandle(self.0, self.1.clone())
    }
}

impl<W: ?Sized> WidgetHandle<W> {
    pub fn id(&self) -> WidgetId {
        self.0
    }

    pub fn get(&self) -> RefMap<'_, W> {
        Ref::map(self.1.get(), |i| &mut i.widget)
    }

    pub fn get_mut(&self) -> RefMapMut<'_, W> {
        let inner = self.1.get_mut();
        if let Some(send) = &inner.send {
            let _ = send.send(self.0);
        }
        RefMut::map(inner, |i| &mut i.widget)
    }

    pub fn data(&self) -> Ref<'_, WidgetData<W>> {
        self.1.get()
    }

    pub(crate) fn data_mut(&self) -> RefMut<'_, WidgetData<W>> {
        self.1.get_mut()
    }

    pub(crate) fn get_mut_quiet(&self) -> RefMapMut<'_, W> {
        RefMut::map(self.1.get_mut(), |i| &mut i.widget)
    }

    pub fn get_label(&self) -> RefMap<'_, String> {
        Ref::map(self.1.get(), |i| &mut i.label)
    }

    pub fn set_label(&self, label: impl Into<String>) {
        self.1.get_mut().label = label.into();
    }

    /// TODO: THIS SHOULD ALWAYS BE 1, remove this probably (weak count might be nice though)
    pub fn refs(&self) -> usize {
        self.1.refs()
    }

    pub fn weak(&self) -> WidgetRef<W> {
        WidgetRef::new(self.0)
    }
}

pub type WRef<'a, W> = MappedRefGuard<'a, WidgetData<W>, W>;
pub type WRefMut<'a, W> = MappedRefGuardMut<'a, WidgetData<W>, W>;

impl<W: Widget> WidgetRef<W> {
    fn handle(&self) -> Handle<WidgetData<W>> {
        WIDGETS.get_type(self.0).unwrap()
    }

    pub fn get<'a>(&self) -> WRef<'a, W> {
        RefGuard::map(self.handle().get_take(), |i| &mut i.widget)
    }

    pub fn get_mut<'a>(&self) -> WRefMut<'a, W> {
        let inner = self.handle().get_take_mut();

        if let Some(send) = &inner.send {
            let _ = send.send(self.0);
        }
        RefGuardMut::map(inner, |i| &mut i.widget)
    }
}

impl<W: ?Sized + Widget> WidgetRef<W> {
    pub fn data<'a>(&self) -> Option<RefGuard<'a, WidgetData<dyn Widget>>> {
        Some(WIDGETS.get(self.0)?.get_take())
    }

    pub(crate) fn data_mut(&self) -> Option<RefGuardMut<'_, WidgetData<dyn Widget>>> {
        Some(WIDGETS.get(self.0)?.get_take_mut())
    }
}

impl<W: Widget + ?Sized + Unsize<dyn Widget>> WidgetRef<W> {
    pub fn any(self) -> WidgetRef<dyn Widget> {
        WidgetRef::new(self.0)
    }
}

impl<W: ?Sized> WidgetRef<W> {
    fn new(id: WidgetId) -> Self {
        Self(id, unsafe { MaybeUninit::zeroed().assume_init() })
    }

    pub fn id(&self) -> WidgetId {
        self.0
    }
}

impl WidgetId {
    /// THIS SHOULD ONLY BE CALLED FOR TEMP STUFF DURING PAINTING
    pub(crate) fn strong(self) -> Option<WidgetHandle> {
        Some(WidgetHandle(self, WIDGETS.get(self)?))
    }

    pub fn weak(self) -> WidgetRef {
        WidgetRef::new(self)
    }

    pub(crate) fn map_event_managers(&self, f: impl Fn(Handle<dyn EventManager>)) {
        let Some(data) = self.weak().data() else {
            return;
        };
        for m in data.event_managers() {
            f(m)
        }
    }
}

pub trait WidgetIdFn<W: ?Sized = dyn Widget>: FnOnce(&mut Ui) -> WidgetHandle<W> {}
impl<W: ?Sized, F: FnOnce(&mut Ui) -> WidgetHandle<W>> WidgetIdFn<W> for F {}

pub trait WidgetRet: FnOnce(&mut Ui) -> WidgetHandle {}
impl<F: FnOnce(&mut Ui) -> WidgetHandle> WidgetRet for F {}

impl<W: Widget + ?Sized + Unsize<dyn Widget> + 'static> WidgetLike<IdTag> for WidgetHandle<W> {
    type Widget = W;
    fn add(self, _: &mut Ui) -> WidgetHandle<W> {
        self
    }
}

impl<W: Widget + ?Sized + Unsize<dyn Widget> + 'static, F: FnOnce(&mut Ui) -> WidgetHandle<W>>
    WidgetLike<IdFnTag> for F
{
    type Widget = W;
    fn add(self, ui: &mut Ui) -> WidgetHandle<W> {
        self(ui)
    }
}

pub trait IdLike<W> {
    fn id(&self) -> WidgetId;
}

impl<W> IdLike<W> for WidgetHandle<W> {
    fn id(&self) -> WidgetId {
        self.id()
    }
}

impl<W> IdLike<W> for WidgetRef<W> {
    fn id(&self) -> WidgetId {
        self.id()
    }
}

impl<T: ?Sized + Unsize<U>, U: ?Sized> CoerceUnsized<WidgetHandle<U>> for WidgetHandle<T> {}
impl<T: ?Sized + Unsize<U>, U: ?Sized> CoerceUnsized<WidgetRef<U>> for WidgetRef<T> {}

impl<W> PartialEq for WidgetHandle<W> {
    fn eq(&self, other: &Self) -> bool {
        self.id() == other.id()
    }
}

impl<W> std::fmt::Debug for WidgetHandle<W> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.id().fmt(f)
    }
}

impl<W: ?Sized> Clone for WidgetRef<W> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<W: ?Sized> Copy for WidgetRef<W> {}

impl<W: ?Sized> PartialEq for WidgetRef<W> {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}
unsafe impl<W: ?Sized> Send for WidgetRef<W> {}
unsafe impl<W: ?Sized> Sync for WidgetRef<W> {}