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

use std::{
    marker::Unsize,
    mem::MaybeUninit,
    ops::{CoerceUnsized, Deref, DerefMut},
    sync::mpsc::Sender,
};

use crate::{
    layout::{IdFnTag, IdTag, Ui, WIDGETS, Widget, WidgetLike},
    util::{Handle, Ref, 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(crate) struct WidgetData<W: ?Sized> {
    pub id: WidgetId,
    pub send: Option<Sender<WidgetUpdate>>,
    pub label: String,
    pub widget: W,
}

pub enum WidgetUpdate {
    Drop(WidgetId),
    Mutate(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,
            }
            .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(WidgetUpdate::Mutate(self.0));
        }
        RefMut::map(inner, |i| &mut i.widget)
    }

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

    pub 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(self.0, null_ptr())
    }
}

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> {
        let handle = self.handle();
        WRef {
            val: unsafe {
                std::mem::transmute::<RefMap<'_, W>, RefMap<'_, W>>(Ref::map(handle.get(), |i| {
                    &mut i.widget
                }))
            },
            _handle: handle,
        }
    }

    pub fn get_mut<'a>(&self) -> WRefMut<'a, W> {
        let handle = self.handle();
        if let Some(send) = &handle.get().send {
            let _ = send.send(WidgetUpdate::Mutate(self.0));
        }
        WRefMut {
            val: unsafe {
                std::mem::transmute::<RefMapMut<'_, W>, RefMapMut<'_, W>>(RefMut::map(
                    handle.get_mut(),
                    |i| &mut i.widget,
                ))
            },
            _handle: handle,
        }
    }
}

pub struct WRef<'a, W> {
    _handle: Handle<WidgetData<W>>,
    val: RefMap<'a, W>,
}

pub struct WRefMut<'a, W> {
    _handle: Handle<WidgetData<W>>,
    val: RefMapMut<'a, W>,
}

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

fn null_ptr<T: ?Sized>() -> *const T {
    unsafe { MaybeUninit::zeroed().assume_init() }
}

impl<W: ?Sized> WidgetRef<W> {
    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)?))
    }
}

impl<W: ?Sized> Drop for WidgetData<W> {
    fn drop(&mut self) {
        if let Some(send) = &self.send {
            WIDGETS.remove(self.id);
            let _ = send.send(WidgetUpdate::Drop(self.id));
        }
    }
}

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 && self.1 == other.1
    }
}

impl<T> Deref for WRef<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.val
    }
}

impl<T> Deref for WRefMut<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.val
    }
}

impl<T> DerefMut for WRefMut<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.val
    }
}