novem/src/query.rs

use gpui::{AppContext, BorrowAppContext};
use std::any::Any;

pub trait QueryFn<C>: Clone + 'static
where
    C: Context,
{
    type Data: 'static;
    type Error: 'static;

    fn key(&self) -> &'static str;

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

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

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

    cx.observe(&ent, |_, _, cx| {
        cx.notify();
    })
    .detach();

    let query_context = cx.global::<Store<C>>().query_context.clone();

    cx.observe(&query_context, move |_, _, cx| {
        cx.update_global::<Store<C>, _>(|store, cx| {
            store.invalidate_query(&query_fn, cx);
            store.ensure_query_data(&query_fn, cx);
        })
    })
    .detach();
}

pub fn read_query<'a, F, T, C>(
    query_fn: F,
    cx: &'a gpui::Context<T>,
) -> QueryStatus<'a, F::Data, F::Error>
where
    C: Context + 'static,
    F: QueryFn<C>,
    T: 'static,
{
    let store = cx.global::<Store<C>>();
    let Some(ent) = store.query_data.get(query_fn.key()) else {
        return QueryStatus::Loading;
    };
    let QueryState { data } = ent.read(cx);
    match 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()),
    }
}

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

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

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

pub(crate) type Entity = gpui::Entity<QueryState>;

pub(crate) trait Context: Clone {}

pub struct Store<C>
where
    C: Context,
{
    query_data: std::collections::HashMap<String, Entity>,
    query_context: gpui::Entity<C>,
}

impl<C> Store<C>
where
    C: Context + 'static,
{
    pub fn new(ctx: C, cx: &mut gpui::App) -> Self {
        Self {
            query_data: std::collections::HashMap::new(),
            query_context: cx.new(|_| ctx),
        }
    }

    fn entity_for<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>) -> Entity
    where
        Q: QueryFn<C>,
        T: 'static,
    {
        self.query_data
            .entry(query.key().into())
            .or_insert_with(|| {
                cx.new(|_| QueryState {
                    data: QueryData::Pending,
                })
            })
            .clone()
    }

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

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

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

        entity
    }

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

        entity.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| {
            let c = query_context.read_with(cx, |c, _| c.clone())?;
            let result = q.run(&c).await;

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

            anyhow::Ok(())
        })
    }

    fn invalidate_query<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>)
    where
        Q: QueryFn<C>,
    {
        if let Some(entity) = self.query_data.get(query.key()) {
            entity.update(cx, |query, cx| {
                if !matches!(query.data, QueryData::Loading) {
                    query.data = QueryData::Stale;
                    cx.notify()
                }
            })
        }
    }

    pub(crate) fn read_query<'a, Q, E>(
        &self,
        query_fn: Q,
        cx: &'a gpui::Context<E>,
    ) -> QueryStatus<'a, Q::Data, Q::Error>
    where
        Q: QueryFn<C>,
    {
        let Some(ent) = self.query_data.get(query_fn.key()) else {
            return QueryStatus::Loading;
        };
        let QueryState { data } = ent.read(cx);
        match data {
            QueryData::Loading | QueryData::Pending | QueryData::Stale => QueryStatus::Loading,
            QueryData::Some(data) => QueryStatus::Loaded(data.downcast_ref::<Q::Data>().unwrap()),
            QueryData::Err(error) => QueryStatus::Err(error.downcast_ref::<Q::Error>().unwrap()),
        }
    }
}

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