iris/src/layout/id.rs

use std::{
    any::TypeId,
    marker::PhantomData,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc::Sender,
    },
};

use crate::{
    layout::{FnTag, Ui, Widget, WidgetLike, WidgetTag},
    util::{Id, RefCounter},
};

pub struct AnyWidget;

/// An identifier for a widget that can index a UI to get the associated widget.
/// It should always remain valid; it keeps a ref count and removes the widget from the UI if all
/// references are dropped.
///
/// W does not need to implement widget so that AnyWidget is valid;
/// Instead, add generic bounds on methods that take an ID if they need specific data.
///
/// TODO: ergonomic clones when they get put in rust-analyzer & don't cause ICEs?
#[repr(C)]
pub struct WidgetId<W = AnyWidget> {
    pub(super) ty: TypeId,
    pub(super) id: Id,
    counter: RefCounter,
    send: Sender<Id>,
    is_static: Arc<AtomicBool>,
    _pd: PhantomData<W>,
}

/// A WidgetId for a static widget that cannot be removed from a Ui.
/// Useful because ergonomic clones don't exist yet so you can easily use these in closures.
/// Do not use this if you want the widget to be freeable.
///
/// This is currently not perfectly efficient and just creates new WidgetIds every time it's used,
/// but they don't send drop messages to Ui.
/// Ideally I'd have an enum or something that lets you use either, but that doesn't seem worth it
/// right now; it's good enough and relatively cheap.
#[repr(C)]
pub struct StaticWidgetId<W = AnyWidget> {
    pub(super) ty: TypeId,
    pub(super) id: Id,
    _pd: PhantomData<W>,
}

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

impl<W> Clone for WidgetId<W> {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            ty: self.ty,
            counter: self.counter.clone(),
            send: self.send.clone(),
            is_static: self.is_static.clone(),
            _pd: PhantomData,
        }
    }
}

impl<W> WidgetId<W> {
    pub(super) fn new(id: Id, ty: TypeId, send: Sender<Id>, is_static: bool) -> Self {
        Self {
            ty,
            id,
            counter: RefCounter::new(),
            send,
            is_static: Arc::new(is_static.into()),
            _pd: PhantomData,
        }
    }

    pub fn any(self) -> WidgetId<AnyWidget> {
        self.cast_type()
    }

    pub fn as_any(&self) -> &WidgetId<AnyWidget> {
        // SAFETY: self is repr(C) and generic only used for phantom data
        unsafe { std::mem::transmute(self) }
    }

    pub fn key(&self) -> Id {
        self.id
    }

    pub(super) fn cast_type<W2>(self) -> WidgetId<W2> {
        // SAFETY: self is repr(C) and generic only used for phantom data
        unsafe { std::mem::transmute(self) }
    }

    pub fn refs(&self) -> u32 {
        self.counter.refs()
    }

    pub fn into_static(self) -> StaticWidgetId<W> {
        self.is_static.store(true, Ordering::Release);
        StaticWidgetId {
            ty: self.ty,
            id: self.id,
            _pd: PhantomData,
        }
    }
}

impl WidgetId {
    pub fn set_static<W>(&mut self, other: StaticWidgetId<W>) {
        let send = self.send.clone();
        drop(std::mem::replace(
            self,
            Self::new(other.id, self.ty, send, true),
        ));
    }
}

impl<W> Drop for WidgetId<W> {
    fn drop(&mut self) {
        if self.counter.drop() && !self.is_static.load(Ordering::Acquire) {
            let _ = self.send.send(self.id);
        }
    }
}

pub struct IdTag;
pub struct IdFnTag;

pub trait WidgetIdFn<W>: FnOnce(&mut Ui) -> WidgetId<W> {}
impl<W, F: FnOnce(&mut Ui) -> WidgetId<W>> WidgetIdFn<W> for F {}

/// TODO: does this ever make sense to use? it allows for invalid ids
pub trait Idable<Tag> {
    type Widget: Widget;
    fn set(self, ui: &mut Ui, id: &WidgetId<Self::Widget>);
    fn id(self, id: &WidgetId<Self::Widget>) -> impl WidgetIdFn<Self::Widget>
    where
        Self: Sized,
    {
        let id = id.clone();
        move |ui| {
            self.set(ui, &id);
            id
        }
    }
    fn id_static(self, id: StaticWidgetId<Self::Widget>) -> impl WidgetIdFn<Self::Widget>
    where
        Self: Sized,
    {
        move |ui| {
            let id = id.id(&ui.send);
            self.set(ui, &id);
            id
        }
    }
}

impl<W: Widget> Idable<WidgetTag> for W {
    type Widget = W;

    fn set(self, ui: &mut Ui, id: &WidgetId<Self::Widget>) {
        ui.set(id, self);
    }
}

impl<F: FnOnce(&mut Ui) -> W, W: Widget> Idable<FnTag> for F {
    type Widget = W;

    fn set(self, ui: &mut Ui, id: &WidgetId<Self::Widget>) {
        let w = self(ui);
        ui.set(id, w);
    }
}

impl<W: 'static> WidgetLike<IdTag> for WidgetId<W> {
    type Widget = W;
    fn add(self, _: &mut Ui) -> WidgetId<W> {
        self
    }
}

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

impl<W> StaticWidgetId<W> {
    pub fn to_id(&self, send: &Sender<Id>) -> WidgetId<W> {
        WidgetId::new(self.id, self.ty, send.clone(), true)
    }
    pub fn any(self) -> StaticWidgetId<AnyWidget> {
        // SAFETY: self is repr(C)
        unsafe { std::mem::transmute(self) }
    }
}

impl<W: 'static> WidgetLike<IdTag> for StaticWidgetId<W> {
    type Widget = W;
    fn add(self, ui: &mut Ui) -> WidgetId<W> {
        self.id(&ui.send)
    }
}

impl<W> Clone for StaticWidgetId<W> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<W> Copy for StaticWidgetId<W> {}

pub trait WidgetIdLike<W> {
    fn id(self, send: &Sender<Id>) -> WidgetId<W>;
}

impl<W> WidgetIdLike<W> for &WidgetId<W> {
    fn id(self, _: &Sender<Id>) -> WidgetId<W> {
        self.clone()
    }
}

impl<W> WidgetIdLike<W> for StaticWidgetId<W> {
    fn id(self, send: &Sender<Id>) -> WidgetId<W> {
        self.to_id(send)
    }
}

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

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

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