tachys on mounted hook

on_mounted.rs

use {
    tachys::{
        html::attribute::{self, Attribute},
        prelude::AddAnyAttr,
        view::{IntoRender, Mountable, Render, RenderHtml},
    },
    tracing::{instrument, warn},
};

trait StaticCallbackOnce: FnOnce() + Send + 'static {}
impl<T> StaticCallbackOnce for T where T: FnOnce() + Send + 'static {}

struct Callback(Option<Box<dyn StaticCallbackOnce>>);

#[extension_traits::extension(pub trait TachysHookMountEventsExt)]
impl<T> T
where
    Self: Sized,
    T: IntoRender,
{
    fn on_mount<M, U>(self, on_mounted: M, on_unmounted: U) -> OnMounted<<T as IntoRender>::Output>
    where
        M: StaticCallbackOnce,
        U: StaticCallbackOnce,
    {
        OnMounted {
            inner: self.into_render(),
            mount_callbacks: MountCallbacks {
                on_mounted: Callback::new(on_mounted),
                on_unmounted: Callback::new(on_unmounted),
            },
        }
    }
}

pub enum MountEvent {
    Mounted,
    Unmounted,
}

impl Callback {
    const fn empty() -> Self {
        Self(None)
    }
    fn new<F>(callback: F) -> Self
    where
        F: StaticCallbackOnce,
    {
        Self(Some(Box::new(callback)))
    }

    #[instrument(skip(self))]
    pub fn call(&mut self) {
        match self.0.take() {
            Some(some) => (some)(),
            None => warn!("this callback has already been executed"),
        }
    }
}

struct MountCallbacks {
    on_mounted: Callback,
    on_unmounted: Callback,
}

impl MountCallbacks {
    const fn empty() -> Self {
        Self {
            on_mounted: Callback::empty(),
            on_unmounted: Callback::empty(),
        }
    }
}

pub struct OnMounted<Inner> {
    inner: Inner,
    mount_callbacks: MountCallbacks,
}

pub struct OnMountedState<Inner> {
    inner: Inner,
    mount_callbacks: MountCallbacks,
}

impl<Inner> Render for OnMounted<Inner>
where
    Inner: Render,
    OnMountedState<Inner::State>: Mountable,
{
    type State = OnMountedState<<Inner as Render>::State>;

    fn build(self) -> Self::State {
        let Self { inner, mount_callbacks } = self;
        OnMountedState {
            inner: inner.build(),
            mount_callbacks,
        }
    }

    fn rebuild(self, state: &mut Self::State) {
        state.mount_callbacks = self.mount_callbacks;
        self.inner.rebuild(&mut state.inner)
    }
}

impl<Inner> AddAnyAttr for OnMounted<Inner>
where
    Inner: AddAnyAttr,
{
    type Output<SomeNewAttr: Attribute> = OnMounted<Inner::Output<SomeNewAttr>>;

    fn add_any_attr<NewAttr: Attribute>(self, attr: NewAttr) -> Self::Output<NewAttr> {
        OnMounted {
            inner: self.inner.add_any_attr(attr),
            mount_callbacks: self.mount_callbacks,
        }
    }
}

impl<Inner> RenderHtml for OnMounted<Inner>
where
    Inner: RenderHtml + AddAnyAttr + Send,
{
    type AsyncOutput = Inner::AsyncOutput;

    type Owned = OnMounted<Inner::Owned>;

    const MIN_LENGTH: usize = Inner::MIN_LENGTH;
    const EXISTS: bool = Inner::EXISTS;

    fn dry_resolve(&mut self) {
        self.inner.dry_resolve();
    }

    async fn resolve(self) -> Self::AsyncOutput {
        self.inner.resolve().await
    }

    fn to_html_with_buf(
        self,
        buf: &mut String,
        position: &mut tachys::view::Position,
        escape: bool,
        mark_branches: bool,
        extra_attrs: Vec<attribute::any_attribute::AnyAttribute>,
    ) {
        self.inner
            .to_html_with_buf(buf, position, escape, mark_branches, extra_attrs);
    }

    fn hydrate<const FROM_SERVER: bool>(self, cursor: &tachys::hydration::Cursor, position: &tachys::view::PositionState) -> Self::State {
        OnMountedState {
            inner: self.inner.hydrate::<FROM_SERVER>(cursor, position),
            mount_callbacks: self.mount_callbacks,
        }
    }

    fn into_owned(self) -> Self::Owned {
        OnMounted {
            inner: self.inner.into_owned(),
            mount_callbacks: self.mount_callbacks,
        }
    }
}

impl<Inner> Mountable for OnMountedState<Inner>
where
    Inner: Mountable,
{
    fn unmount(&mut self) {
        self.mount_callbacks.on_unmounted.call();
        self.inner.unmount();
    }

    fn mount(&mut self, parent: &tachys::renderer::types::Element, marker: Option<&tachys::renderer::types::Node>) {
        self.mount_callbacks.on_mounted.call();
        self.inner.mount(parent, marker);
    }

    fn insert_before_this(&self, child: &mut dyn Mountable) -> bool {
        self.inner.insert_before_this(child)
    }

    fn elements(&self) -> Vec<tachys::renderer::types::Element> {
        self.inner.elements()
    }
}