novem/src/query.rs

use gpui::{AppContext, BorrowAppContext};
use std::{any::Any, collections::HashMap, marker::PhantomData};

pub trait QueryFn: Clone + 'static {
    type Data: 'static;
    type Error: 'static;
    type Context: Context;

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

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

struct Query {
    key: &'static str,
    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.raw, 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();

    let cloned_ent = ent.clone();
    let query_context = cx.global::<Store<F::Context>>().query_context.clone();
    cx.observe(&query_context, move |_, _, cx| {
        cx.update_global::<Store<F::Context>, _>(|store, cx| {
            store.invalidate_query(&cloned_ent, cx);
        });
    })
    .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.raw, 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();

    let cloned_ent = ent.clone();
    let query_context = cx.global::<Store<F::Context>>().query_context.clone();
    cx.observe(&query_context, move |_, _, cx| {
        cx.update_global::<Store<F::Context>, _>(|store, cx| {
            store.invalidate_query(&cloned_ent, cx);
        });
    })
    .detach();

    ent
}

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

pub fn read_query<'a, F, T>(
    query: &Entity<F>,
    cx: &'a gpui::Context<T>,
) -> QueryStatus<'a, F::Data, F::Error>
where
    F: QueryFn,
    T: 'static,
{
    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()),
    }
}

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

pub(crate) trait Context: Clone {}

pub struct Store<C>
where
    C: Context,
{
    query_data: HashMap<String, gpui::Entity<Query>>,
    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_context: cx.new(|_| ctx),
            query_data: std::collections::HashMap::new(),
        }
    }

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

        Entity {
            raw,
            _marker: PhantomData,
        }
    }

    fn ensure_query_data<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>) -> Entity<Q>
    where
        T: 'static,
        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(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<Context = C>,
        T: 'static,
    {
        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| {
            let c = query_context.read_with(cx, |c, _| c.clone())?;
            let result = q.run(&c).await;

            entity.raw.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<E, F>(&self, entity: &Entity<F>, cx: &mut gpui::Context<E>)
    where
        E: 'static,
        F: QueryFn<Context = C>,
    {
        let entity = entity.raw.read(cx);
        if let Some(entity) = self.query_data.get(entity.key) {
            entity.update(cx, |query, cx| {
                if !matches!(query.data, QueryData::Loading) {
                    query.data = QueryData::Stale;
                    cx.notify()
                }
            })
        }
    }
}

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