novem/src/query.rs

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

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

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

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

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

pub fn use_query<F, T>(query_fn: F, cx: &mut gpui::Context<T>)
where
    F: QueryFn,
    T: 'static,
{
    let ent = cx.update_global::<app::Global, _>(|global, cx| {
        let entity = global.query_store.ensure_query(&query_fn, cx);

        if entity.read_with(cx, |state, _| matches!(state.data, QueryData::Pending)) {
            global.query_store.execute_query(query_fn, entity.clone(), cx).detach();
        }

        entity
    });

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

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

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

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

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

pub struct Store {
    query_data: std::collections::HashMap<String, Entity>,
}

impl Store {
    pub fn new() -> Self {
        Self {
            query_data: std::collections::HashMap::new(),
        }
    }

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

    fn execute_query<Q, T>(
        &mut self,
        query: Q,
        entity: Entity,
        cx: &mut gpui::Context<T>,
    ) -> gpui::Task<anyhow::Result<()>>
    where
        Q: QueryFn,
        T: 'static,
    {
        entity.update(cx, |state, cx| {
            state.data = QueryData::Loading;
            cx.notify();
        });

        cx.spawn(async move |_, cx| {
            let result = query.run().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(())
        })
    }

    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  {
        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 => QueryStatus::Loading,
            QueryData::Some(data) => QueryStatus::Loaded(data.downcast_ref::<Q::Data>().unwrap()),
            QueryData::Err(error) => QueryStatus::Err(error.downcast_ref::<Q::Error>().unwrap())
        }
    }
}