some bs

src/query.rs

@@ -2,13 +2,16 @@ use crate::app;
 use gpui::{AppContext, BorrowAppContext};
 use std::any::Any;
 
-pub trait QueryFn: Clone + 'static {
+pub trait QueryFn<C>: Clone + 'static
+where
+    C: Context + 'static,
+{
     type Data: 'static;
     type Error: 'static;
 
     fn key(&self) -> &'static str;
 
-    async fn run(&self) -> Result<Self::Data, Self::Error>;
+    async fn run(&self, c: &C) -> Result<Self::Data, Self::Error>;
 }
 
 pub enum QueryStatus<'a, Data, Error> {
@@ -17,25 +20,29 @@ pub enum QueryStatus<'a, Data, Error> {
     Err(&'a Error),
 }
 
-pub fn use_query<F, T>(query_fn: F, cx: &mut gpui::Context<T>)
+pub fn use_query<F, T, C>(query_fn: F, cx: &mut gpui::Context<T>)
 where
-    F: QueryFn,
+    C: Context + 'static,
+    F: QueryFn<C>,
     T: 'static,
+    Store<C>: gpui::Global,
 {
-    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
-    });
+    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();
 }
 
 // ================= Store ==================
@@ -47,26 +54,37 @@ pub(crate) struct QueryState {
 pub(crate) enum QueryData {
     Pending,
     Loading,
+    Stale,
     Some(Box<dyn Any>),
     Err(Box<dyn Any>),
 }
 
 pub(crate) type Entity = gpui::Entity<QueryState>;
 
-pub struct Store {
+pub(crate) trait Context {}
+
+pub struct Store<C>
+where
+    C: Context,
+{
     query_data: std::collections::HashMap<String, Entity>,
+    query_context: gpui::Entity<C>,
 }
 
-impl Store {
-    pub fn new() -> Self {
+impl<C> Store<C>
+where
+    C: Context + 'static,
+{
+    pub fn new<E>(ctx: C, cx: &mut gpui::Context<E>) -> Self {
         Self {
             query_data: std::collections::HashMap::new(),
+            query_context: cx.new(|_| ctx),
         }
     }
 
-    fn ensure_query<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
     where
-        Q: QueryFn,
+        Q: QueryFn<C>,
         T: 'static,
     {
         self.query_data
@@ -79,23 +97,43 @@ impl Store {
             .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,
-        entity: Entity,
+        query: &Q,
         cx: &mut gpui::Context<T>,
     ) -> gpui::Task<anyhow::Result<()>>
     where
-        Q: QueryFn,
+        Q: QueryFn<C>,
         T: 'static,
     {
+        let entity = self.entity_for(query, cx);
+
         entity.update(cx, |state, cx| {
             state.data = QueryData::Loading;
             cx.notify();
         });
 
         cx.spawn(async move |_, cx| {
-            let result = query.run().await;
+            let query_context = self.query_context.read(cx);
+            let result = query.run(query_context).await;
 
             entity.update(cx, |state, cx| {
                 state.data = match result {
@@ -109,15 +147,36 @@ impl Store {
         })
     }
 
-    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  {
+    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 => QueryStatus::Loading,
+            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())
+            QueryData::Err(error) => QueryStatus::Err(error.downcast_ref::<Q::Error>().unwrap()),
         }
     }
 }