novem/src/query.rs

use gpui::BorrowAppContext;
use std::{any::Any, borrow::Cow, collections::HashMap, marker::PhantomData, ops::Deref};

pub(crate) trait QueryFn: Clone + 'static {
    type Data: 'static;
    type Error: std::fmt::Debug + 'static;
    type Context: Context;

    fn key(&self) -> Key;

    async fn run(&self, c: &Self::Context) -> Result<Self::Data, Self::Error>;
}

pub(crate) trait QueryAppContext: gpui::AppContext {
    fn ready<T>(value: T) -> Self::Result<T>;

    fn map_result<T, U>(result: Self::Result<T>, f: impl FnOnce(T) -> U) -> Self::Result<U>;
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub(crate) struct Key(Cow<'static, str>);

pub(crate) struct Query {
    data: QueryData,
}

enum QueryData {
    Pending,
    Loading,
    Stale,
    Some(Box<dyn Any>),
    Err(Box<dyn Any>),
}

pub enum QueryStatus<'a, Data, Error> {
    Loading,
    Loaded(&'a Data),
    Err(&'a Error),
}

#[derive(Clone)]
pub(crate) struct Entity<F>
where
    F: QueryFn,
{
    raw: gpui::Entity<Query>,
    _marker: PhantomData<fn() -> F>,
}

pub fn use_query<F, T>(query_fn: F, cx: &mut gpui::Context<T>) -> Entity<F>
where
    F: QueryFn,
    T: 'static,
    Store<F::Context>: gpui::Global,
{
    let ent = cx.update_global::<Store<F::Context>, _>(|store, cx| {
        let ent = store.entity_for(&query_fn, cx);
        store.ensure_query_data(&query_fn, cx);
        ent
    });

    cx.observe(&ent, move |_, ent, cx| {
        let query = ent.read(cx);
        if matches!(query.data, QueryData::Stale) {
            cx.update_global::<Store<F::Context>, _>(|store, cx| {
                store.ensure_query_data(&query_fn, cx);
            });
        }
        cx.notify();
    })
    .detach();

    ent
}

pub fn use_lazy_query<F, T>(query_fn: F, cx: &mut gpui::Context<T>) -> Entity<F>
where
    F: QueryFn,
    T: 'static,
    Store<F::Context>: gpui::Global,
{
    let ent = cx.update_global::<Store<F::Context>, _>(|store, cx| store.entity_for(&query_fn, cx));

    cx.observe(&ent, move |_, ent, cx| {
        let query = ent.read(cx);
        if matches!(query.data, QueryData::Stale) {
            cx.update_global::<Store<F::Context>, _>(|store, cx| {
                store.ensure_query_data(&query_fn, cx);
            });
        }
        cx.notify();
    })
    .detach();

    ent
}

pub async fn fetch_query<F>(query_fn: F, cx: &mut gpui::AsyncApp) -> anyhow::Result<Entity<F>>
where
    F: QueryFn,
{
    let ent = cx.update(|cx| {
        cx.update_global::<Store<F::Context>, _>(|store, cx| store.ensure_query_data(&query_fn, cx))
    })?;

    enum WaitState {
        Cached,
        Waiting {
            rx: futures::channel::oneshot::Receiver<()>,
            sub: gpui::Subscription,
        },
    }

    loop {
        let wait_state = cx.update(|cx| {
            let is_done = ent.read_with(cx, |query, _| {
                matches!(query.data, QueryData::Some(_) | QueryData::Err(_))
            });
            if is_done {
                WaitState::Cached
            } else {
                let (tx, rx) = futures::channel::oneshot::channel();
                let ent = ent.clone();
                let mut tx = Some(tx);

                let sub = cx.observe(&ent, move |ent, cx| {
                    let is_done = ent.read_with(cx, |query, _| {
                        matches!(query.data, QueryData::Some(_) | QueryData::Err(_))
                    });
                    if is_done && let Some(tx) = tx.take() {
                        _ = tx.send(());
                    }
                });

                WaitState::Waiting { rx, sub }
            }
        })?;

        match wait_state {
        | WaitState::Cached => {
            return Ok(ent);
        }
        | WaitState::Waiting { rx, sub } => {
            _ = sub;
            _ = rx.await;
        }
        }
    }
}

impl<F> Entity<F>
where
    F: QueryFn,
    Store<F::Context>: gpui::Global,
{
    pub fn refetch<E>(&self, cx: &mut gpui::Context<E>)
    where
        E: 'static,
    {
        cx.update_global::<Store<F::Context>, _>(|store, cx| {
            store.invalidate_query(self, cx);
        });
    }
}

impl<F> Deref for Entity<F>
where
    F: QueryFn,
{
    type Target = gpui::Entity<Query>;

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

pub fn read_query<'a, F>(query: &Entity<F>, cx: &'a gpui::App) -> QueryStatus<'a, F::Data, F::Error>
where
    F: QueryFn,
{
    let state = query.raw.read(cx);

    match &state.data {
    | QueryData::Loading | QueryData::Pending | QueryData::Stale => QueryStatus::Loading,
    | QueryData::Some(data) => QueryStatus::Loaded(data.downcast_ref::<F::Data>().unwrap()),
    | QueryData::Err(error) => QueryStatus::Err(error.downcast_ref::<F::Error>().unwrap()),
    }
}

pub fn watch_query<E, F, H>(query: &Entity<F>, on_notify: H, cx: &mut gpui::Context<E>) -> gpui::Subscription
where
    E: 'static,
    F: QueryFn,
    H: Fn(&mut E, &Entity<F>, &mut gpui::Context<E>) + Clone + 'static,
{
    let observed_query = query.clone();
    let sub = cx.observe(query, {
        let on_notify = on_notify.clone();
        move |this, _, cx| on_notify(this, &observed_query, cx)
    });

    let initial_query = query.clone();
    cx.spawn({
        let on_notify = on_notify.clone();
        async move |weak, cx| {
            let _ = weak.update(cx, |this, cx| on_notify(this, &initial_query, cx));
        }
    })
    .detach();

    sub
}

// ================= Store ==================

pub(crate) trait Context: Clone {}

pub struct Store<C>
where
    C: Context,
{
    query_data: HashMap<Key, gpui::Entity<Query>>,
    query_context: C,
}

impl<C> Store<C>
where
    C: Context + 'static,
{
    pub fn new(ctx: C) -> Self {
        Self {
            query_context: ctx,
            query_data: std::collections::HashMap::new(),
        }
    }

    pub(crate) fn update_query_context(&mut self, f: impl FnOnce(&mut C)) {
        f(&mut self.query_context);
    }

    fn entity_for<Q, CX>(&mut self, query: &Q, cx: &mut CX) -> CX::Result<Entity<Q>>
    where
        Q: QueryFn<Context = C>,
        CX: QueryAppContext,
    {
        let query_key = query.key();

        if let Some(raw) = self.query_data.get(&query_key) {
            return CX::ready(Entity {
                raw: raw.clone(),
                _marker: PhantomData,
            });
        }

        CX::map_result(
            cx.new(|_| Query {
                data: QueryData::Pending,
            }),
            |raw| {
                self.query_data.insert(query_key, raw.clone());
                Entity {
                    raw,
                    _marker: PhantomData,
                }
            },
        )
    }

    fn ensure_query_data<Q>(&mut self, query: &Q, cx: &mut gpui::App) -> Entity<Q>
    where
        Q: QueryFn<Context = C>,
    {
        let entity = self.entity_for(query, cx);

        let should_execute = entity.raw.read_with(cx, |state, _| {
            matches!(state.data, QueryData::Pending | QueryData::Stale)
        });

        if should_execute {
            self.execute_query_detached(query, cx).detach();
        }

        entity
    }

    fn execute_query_detached<Q>(
        &mut self,
        query: &Q,
        cx: &mut gpui::App,
    ) -> gpui::Task<anyhow::Result<()>>
    where
        Q: QueryFn<Context = C>,
    {
        let entity = self.entity_for(query, cx);

        entity.raw.update(cx, |state, cx| {
            state.data = QueryData::Loading;
            cx.notify();
        });

        let q = query.clone();
        let query_context = self.query_context.clone();

        cx.spawn(async move |cx| {
            println!("[query] {}", q.key());

            let result = q.run(&query_context).await;

            entity.raw.update(cx, |state, cx| {
                state.data = match result {
                | Ok(data) => {
                    println!("[query] OK {}", q.key());
                    QueryData::Some(Box::new(data))
                }
                | Err(err) => {
                    println!("[query] ERR {:?}: {:?}", q.key(), err);
                    QueryData::Err(Box::new(err))
                }
                };
                cx.notify();
            })?;

            anyhow::Ok(())
        })
    }

    async fn execute_query<'a, F>(&mut self, query_fn: &F, cx: &'a mut gpui::AsyncApp) -> Entity<F>
    where
        F: QueryFn<Context = C>,
    {
        let entity = self.entity_for(query_fn, cx).unwrap();

        entity.update(cx, |query, cx| {
            query.data = QueryData::Loading;
            cx.notify();
        });

        let result = query_fn.run(&self.query_context).await;

        entity
            .raw
            .update(cx, |query, cx| {
                query.data = match result {
                | Ok(data) => QueryData::Some(Box::new(data)),
                | Err(err) => QueryData::Err(Box::new(err)),
                };
                cx.notify();
                true
            })
            .unwrap();

        entity
    }

    fn invalidate_query<E, F>(&self, entity: &Entity<F>, cx: &mut gpui::Context<E>)
    where
        E: 'static,
        F: QueryFn<Context = C>,
    {
        entity.update(cx, |query, cx| {
            if !matches!(query.data, QueryData::Loading) {
                query.data = QueryData::Stale;
                cx.notify()
            }
        });
    }
}

impl QueryAppContext for gpui::App {
    fn ready<T>(value: T) -> Self::Result<T> {
        value
    }

    fn map_result<T, U>(result: Self::Result<T>, f: impl FnOnce(T) -> U) -> Self::Result<U> {
        f(result)
    }
}

impl QueryAppContext for gpui::AsyncApp {
    fn ready<T>(value: T) -> Self::Result<T> {
        Ok(value)
    }

    fn map_result<T, U>(result: Self::Result<T>, f: impl FnOnce(T) -> U) -> Self::Result<U> {
        result.map(f)
    }
}

impl<'a, E> QueryAppContext for gpui::Context<'a, E> {
    fn ready<T>(value: T) -> Self::Result<T> {
        value
    }

    fn map_result<T, U>(result: Self::Result<T>, f: impl FnOnce(T) -> U) -> Self::Result<U> {
        f(result)
    }
}

impl AsRef<str> for Key {
    fn as_ref(&self) -> &str {
        self.0.as_ref()
    }
}

impl From<&'static str> for Key {
    fn from(key: &'static str) -> Self {
        Key(Cow::Borrowed(key))
    }
}

impl From<String> for Key {
    fn from(key: String) -> Self {
        Key(Cow::Owned(key))
    }
}

impl std::fmt::Display for Key {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl<C> gpui::Global for Store<C> where C: Context + 'static {}