wip

src/query.rs

@@ -1,16 +1,27 @@
 use gpui::{AppContext, BorrowAppContext};
-use std::any::Any;
+use std::{any::Any, collections::HashMap, marker::PhantomData};
 
-pub trait QueryFn<C>: Clone + 'static
-where
-    C: Context,
-{
+pub trait QueryFn: Clone + 'static {
     type Data: 'static;
     type Error: 'static;
+    type Context: Context;
 
     fn key(&self) -> &'static str;
 
-    async fn run(&self, c: &C) -> Result<Self::Data, Self::Error>;
+    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> {
@@ -19,46 +30,77 @@ pub enum QueryStatus<'a, Data, Error> {
     Err(&'a Error),
 }
 
-pub fn use_query<F, T, C>(query_fn: F, cx: &mut gpui::Context<T>)
+#[derive(Clone)]
+pub(crate) struct Entity<F>
 where
-    C: Context + 'static,
-    F: QueryFn<C>,
-    T: 'static,
-    Store<C>: gpui::Global,
+    F: QueryFn,
 {
-    let ent = cx.update_global::<Store<C>, _>(|store, cx| store.ensure_query_data(&query_fn, cx));
+    raw: gpui::Entity<Query>,
+    _marker: PhantomData<fn() -> F>,
+}
 
-    cx.observe(&ent, |_, _, cx| {
+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| store.ensure_query_data(&query_fn, cx));
+
+    cx.observe(&ent.raw, |_, _, cx| {
         cx.notify();
     })
     .detach();
 
-    let query_context = cx.global::<Store<C>>().query_context.clone();
-
+    let query_context = cx.global::<Store<F::Context>>().query_context.clone();
     cx.observe(&query_context, move |_, _, cx| {
-        cx.update_global::<Store<C>, _>(|store, cx| {
+        cx.update_global::<Store<F::Context>, _>(|store, cx| {
             store.invalidate_query(&query_fn, cx);
             store.ensure_query_data(&query_fn, cx);
         })
     })
     .detach();
+
+    ent
 }
 
-pub fn read_query<'a, F, T, C>(
-    query_fn: F,
+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, |_, ent, cx| {
+        cx.notify();
+    })
+    .detach();
+
+    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(&query_fn, cx);
+            store.ensure_query_data(&query_fn, cx);
+        })
+    })
+    .detach();
+
+    ent
+}
+
+pub fn read_query<'a, F, T>(
+    query: &Entity<F>,
     cx: &'a gpui::Context<T>,
 ) -> QueryStatus<'a, F::Data, F::Error>
 where
-    C: Context + 'static,
-    F: QueryFn<C>,
+    F: QueryFn,
     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 {
+    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()),
@@ -67,27 +109,13 @@ where
 
 // ================= 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_data: HashMap<String, gpui::Entity<Query>>,
     query_context: gpui::Entity<C>,
 }
 
@@ -102,29 +130,36 @@ where
         }
     }
 
-    fn entity_for<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>) -> Entity
+    fn entity_for<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>) -> Entity<Q>
     where
-        Q: QueryFn<C>,
+        Q: QueryFn<Context = C>,
         T: 'static,
     {
-        self.query_data
+        let raw = self
+            .query_data
             .entry(query.key().into())
             .or_insert_with(|| {
-                cx.new(|_| QueryState {
+                cx.new(|_| Query {
+                    key: query.key(),
                     data: QueryData::Pending,
                 })
             })
-            .clone()
+            .clone();
+
+        Entity {
+            raw,
+            _marker: PhantomData,
+        }
     }
 
-    fn ensure_query_data<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>) -> Entity
+    fn ensure_query_data<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>) -> Entity<Q>
     where
         T: 'static,
-        Q: QueryFn<C>,
+        Q: QueryFn<Context = C>,
     {
         let entity = self.entity_for(query, cx);
 
-        let should_execute = entity.read_with(cx, |state, _| {
+        let should_execute = entity.raw.read_with(cx, |state, _| {
             matches!(state.data, QueryData::Pending | QueryData::Stale)
         });
 
@@ -141,12 +176,12 @@ where
         cx: &mut gpui::Context<T>,
     ) -> gpui::Task<anyhow::Result<()>>
     where
-        Q: QueryFn<C>,
+        Q: QueryFn<Context = C>,
         T: 'static,
     {
         let entity = self.entity_for(query, cx);
 
-        entity.update(cx, |state, cx| {
+        entity.raw.update(cx, |state, cx| {
             state.data = QueryData::Loading;
             cx.notify();
         });
@@ -158,7 +193,7 @@ where
             let c = query_context.read_with(cx, |c, _| c.clone())?;
             let result = q.run(&c).await;
 
-            entity.update(cx, |state, cx| {
+            entity.raw.update(cx, |state, cx| {
                 state.data = match result {
                     Ok(data) => QueryData::Some(Box::new(data)),
                     Err(err) => QueryData::Err(Box::new(err)),
@@ -172,7 +207,8 @@ where
 
     fn invalidate_query<Q, T>(&mut self, query: &Q, cx: &mut gpui::Context<T>)
     where
-        Q: QueryFn<C>,
+        Q: QueryFn<Context = C>,
+        T: 'static,
     {
         if let Some(entity) = self.query_data.get(query.key()) {
             entity.update(cx, |query, cx| {
@@ -183,25 +219,6 @@ where
             })
         }
     }
-
-    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 {}