gpui: Add use state APIs (#34741)

This PR adds a component level state API to GPUI, as well as a few
utilities for simplified interactions with entities

Release Notes:

- N/A

crates/eval/src/example.rs

@@ -422,6 +422,13 @@ impl AppContext for ExampleContext {
         self.app.update_entity(handle, update)
     }
 
+    fn as_mut<'a, T>(&'a mut self, handle: &Entity<T>) -> Self::Result<gpui::GpuiBorrow<'a, T>>
+    where
+        T: 'static,
+    {
+        self.app.as_mut(handle)
+    }
+
     fn read_entity<T, R>(
         &self,
         handle: &Entity<T>,

crates/gpui/src/app.rs

@@ -448,15 +448,23 @@ impl App {
     }
 
     pub(crate) fn update<R>(&mut self, update: impl FnOnce(&mut Self) -> R) -> R {
-        self.pending_updates += 1;
+        self.start_update();
         let result = update(self);
+        self.finish_update();
+        result
+    }
+
+    pub(crate) fn start_update(&mut self) {
+        self.pending_updates += 1;
+    }
+
+    pub(crate) fn finish_update(&mut self) {
         if !self.flushing_effects && self.pending_updates == 1 {
             self.flushing_effects = true;
             self.flush_effects();
             self.flushing_effects = false;
         }
         self.pending_updates -= 1;
-        result
     }
 
     /// Arrange a callback to be invoked when the given entity calls `notify` on its respective context.
@@ -868,7 +876,6 @@ impl App {
         loop {
             self.release_dropped_entities();
             self.release_dropped_focus_handles();
-
             if let Some(effect) = self.pending_effects.pop_front() {
                 match effect {
                     Effect::Notify { emitter } => {
@@ -1819,6 +1826,13 @@ impl AppContext for App {
         })
     }
 
+    fn as_mut<'a, T>(&'a mut self, handle: &Entity<T>) -> GpuiBorrow<'a, T>
+    where
+        T: 'static,
+    {
+        GpuiBorrow::new(handle.clone(), self)
+    }
+
     fn read_entity<T, R>(
         &self,
         handle: &Entity<T>,
@@ -2015,3 +2029,79 @@ impl HttpClient for NullHttpClient {
         type_name::<Self>()
     }
 }
+
+/// A mutable reference to an entity owned by GPUI
+pub struct GpuiBorrow<'a, T> {
+    inner: Option<Lease<T>>,
+    app: &'a mut App,
+}
+
+impl<'a, T: 'static> GpuiBorrow<'a, T> {
+    fn new(inner: Entity<T>, app: &'a mut App) -> Self {
+        app.start_update();
+        let lease = app.entities.lease(&inner);
+        Self {
+            inner: Some(lease),
+            app,
+        }
+    }
+}
+
+impl<'a, T: 'static> std::borrow::Borrow<T> for GpuiBorrow<'a, T> {
+    fn borrow(&self) -> &T {
+        self.inner.as_ref().unwrap().borrow()
+    }
+}
+
+impl<'a, T: 'static> std::borrow::BorrowMut<T> for GpuiBorrow<'a, T> {
+    fn borrow_mut(&mut self) -> &mut T {
+        self.inner.as_mut().unwrap().borrow_mut()
+    }
+}
+
+impl<'a, T> Drop for GpuiBorrow<'a, T> {
+    fn drop(&mut self) {
+        let lease = self.inner.take().unwrap();
+        self.app.notify(lease.id);
+        self.app.entities.end_lease(lease);
+        self.app.finish_update();
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::{cell::RefCell, rc::Rc};
+
+    use crate::{AppContext, TestAppContext};
+
+    #[test]
+    fn test_gpui_borrow() {
+        let cx = TestAppContext::single();
+        let observation_count = Rc::new(RefCell::new(0));
+
+        let state = cx.update(|cx| {
+            let state = cx.new(|_| false);
+            cx.observe(&state, {
+                let observation_count = observation_count.clone();
+                move |_, _| {
+                    let mut count = observation_count.borrow_mut();
+                    *count += 1;
+                }
+            })
+            .detach();
+
+            state
+        });
+
+        cx.update(|cx| {
+            // Calling this like this so that we don't clobber the borrow_mut above
+            *std::borrow::BorrowMut::borrow_mut(&mut state.as_mut(cx)) = true;
+        });
+
+        cx.update(|cx| {
+            state.write(cx, false);
+        });
+
+        assert_eq!(*observation_count.borrow(), 2);
+    }
+}

crates/gpui/src/app/async_context.rs

@@ -3,7 +3,7 @@ use crate::{
     Entity, EventEmitter, Focusable, ForegroundExecutor, Global, PromptButton, PromptLevel, Render,
     Reservation, Result, Subscription, Task, VisualContext, Window, WindowHandle,
 };
-use anyhow::Context as _;
+use anyhow::{Context as _, anyhow};
 use derive_more::{Deref, DerefMut};
 use futures::channel::oneshot;
 use std::{future::Future, rc::Weak};
@@ -58,6 +58,15 @@ impl AppContext for AsyncApp {
         Ok(app.update_entity(handle, update))
     }
 
+    fn as_mut<'a, T>(&'a mut self, _handle: &Entity<T>) -> Self::Result<super::GpuiBorrow<'a, T>>
+    where
+        T: 'static,
+    {
+        Err(anyhow!(
+            "Cannot as_mut with an async context. Try calling update() first"
+        ))
+    }
+
     fn read_entity<T, R>(
         &self,
         handle: &Entity<T>,
@@ -364,6 +373,15 @@ impl AppContext for AsyncWindowContext {
             .update(self, |_, _, cx| cx.update_entity(handle, update))
     }
 
+    fn as_mut<'a, T>(&'a mut self, _: &Entity<T>) -> Self::Result<super::GpuiBorrow<'a, T>>
+    where
+        T: 'static,
+    {
+        Err(anyhow!(
+            "Cannot use as_mut() from an async context, call `update`"
+        ))
+    }
+
     fn read_entity<T, R>(
         &self,
         handle: &Entity<T>,

crates/gpui/src/app/context.rs

@@ -726,6 +726,13 @@ impl<T> AppContext for Context<'_, T> {
         self.app.update_entity(handle, update)
     }
 
+    fn as_mut<'a, E>(&'a mut self, handle: &Entity<E>) -> Self::Result<super::GpuiBorrow<'a, E>>
+    where
+        E: 'static,
+    {
+        self.app.as_mut(handle)
+    }
+
     fn read_entity<U, R>(
         &self,
         handle: &Entity<U>,

crates/gpui/src/app/entity_map.rs

@@ -1,4 +1,4 @@
-use crate::{App, AppContext, VisualContext, Window, seal::Sealed};
+use crate::{App, AppContext, GpuiBorrow, VisualContext, Window, seal::Sealed};
 use anyhow::{Context as _, Result};
 use collections::FxHashSet;
 use derive_more::{Deref, DerefMut};
@@ -105,7 +105,7 @@ impl EntityMap {
 
     /// Move an entity to the stack.
     #[track_caller]
-    pub fn lease<'a, T>(&mut self, pointer: &'a Entity<T>) -> Lease<'a, T> {
+    pub fn lease<T>(&mut self, pointer: &Entity<T>) -> Lease<T> {
         self.assert_valid_context(pointer);
         let mut accessed_entities = self.accessed_entities.borrow_mut();
         accessed_entities.insert(pointer.entity_id);
@@ -117,15 +117,14 @@ impl EntityMap {
         );
         Lease {
             entity,
-            pointer,
+            id: pointer.entity_id,
             entity_type: PhantomData,
         }
     }
 
     /// Returns an entity after moving it to the stack.
     pub fn end_lease<T>(&mut self, mut lease: Lease<T>) {
-        self.entities
+        self.entities.insert(lease.id, lease.entity.take().unwrap());
-            .insert(lease.pointer.entity_id, lease.entity.take().unwrap());
     }
 
     pub fn read<T: 'static>(&self, entity: &Entity<T>) -> &T {
@@ -187,13 +186,13 @@ fn double_lease_panic<T>(operation: &str) -> ! {
     )
 }
 
-pub(crate) struct Lease<'a, T> {
+pub(crate) struct Lease<T> {
     entity: Option<Box<dyn Any>>,
-    pub pointer: &'a Entity<T>,
+    pub id: EntityId,
     entity_type: PhantomData<T>,
 }
 
-impl<T: 'static> core::ops::Deref for Lease<'_, T> {
+impl<T: 'static> core::ops::Deref for Lease<T> {
     type Target = T;
 
     fn deref(&self) -> &Self::Target {
@@ -201,13 +200,13 @@ impl<T: 'static> core::ops::Deref for Lease<'_, T> {
     }
 }
 
-impl<T: 'static> core::ops::DerefMut for Lease<'_, T> {
+impl<T: 'static> core::ops::DerefMut for Lease<T> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.entity.as_mut().unwrap().downcast_mut().unwrap()
     }
 }
 
-impl<T> Drop for Lease<'_, T> {
+impl<T> Drop for Lease<T> {
     fn drop(&mut self) {
         if self.entity.is_some() && !panicking() {
             panic!("Leases must be ended with EntityMap::end_lease")
@@ -437,6 +436,19 @@ impl<T: 'static> Entity<T> {
         cx.update_entity(self, update)
     }
 
+    /// Updates the entity referenced by this handle with the given function.
+    pub fn as_mut<'a, C: AppContext>(&self, cx: &'a mut C) -> C::Result<GpuiBorrow<'a, T>> {
+        cx.as_mut(self)
+    }
+
+    /// Updates the entity referenced by this handle with the given function.
+    pub fn write<C: AppContext>(&self, cx: &mut C, value: T) -> C::Result<()> {
+        self.update(cx, |entity, cx| {
+            *entity = value;
+            cx.notify();
+        })
+    }
+
     /// Updates the entity referenced by this handle with the given function if
     /// the referenced entity still exists, within a visual context that has a window.
     /// Returns an error if the entity has been released.

crates/gpui/src/app/test_context.rs

@@ -9,6 +9,7 @@ use crate::{
 };
 use anyhow::{anyhow, bail};
 use futures::{Stream, StreamExt, channel::oneshot};
+use rand::{SeedableRng, rngs::StdRng};
 use std::{cell::RefCell, future::Future, ops::Deref, rc::Rc, sync::Arc, time::Duration};
 
 /// A TestAppContext is provided to tests created with `#[gpui::test]`, it provides
@@ -63,6 +64,13 @@ impl AppContext for TestAppContext {
         app.update_entity(handle, update)
     }
 
+    fn as_mut<'a, T>(&'a mut self, _: &Entity<T>) -> Self::Result<super::GpuiBorrow<'a, T>>
+    where
+        T: 'static,
+    {
+        panic!("Cannot use as_mut with a test app context. Try calling update() first")
+    }
+
     fn read_entity<T, R>(
         &self,
         handle: &Entity<T>,
@@ -134,6 +142,12 @@ impl TestAppContext {
         }
     }
 
+    /// Create a single TestAppContext, for non-multi-client tests
+    pub fn single() -> Self {
+        let dispatcher = TestDispatcher::new(StdRng::from_entropy());
+        Self::build(dispatcher, None)
+    }
+
     /// The name of the test function that created this `TestAppContext`
     pub fn test_function_name(&self) -> Option<&'static str> {
         self.fn_name
@@ -914,6 +928,13 @@ impl AppContext for VisualTestContext {
         self.cx.update_entity(handle, update)
     }
 
+    fn as_mut<'a, T>(&'a mut self, handle: &Entity<T>) -> Self::Result<super::GpuiBorrow<'a, T>>
+    where
+        T: 'static,
+    {
+        self.cx.as_mut(handle)
+    }
+
     fn read_entity<T, R>(
         &self,
         handle: &Entity<T>,

crates/gpui/src/element.rs

@@ -39,7 +39,7 @@ use crate::{
 use derive_more::{Deref, DerefMut};
 pub(crate) use smallvec::SmallVec;
 use std::{
-    any::Any,
+    any::{Any, type_name},
     fmt::{self, Debug, Display},
     mem, panic,
 };
@@ -220,14 +220,17 @@ impl<C: RenderOnce> Element for Component<C> {
         window: &mut Window,
         cx: &mut App,
     ) -> (LayoutId, Self::RequestLayoutState) {
+        window.with_global_id(ElementId::Name(type_name::<C>().into()), |_, window| {
             let mut element = self
                 .component
                 .take()
                 .unwrap()
                 .render(window, cx)
                 .into_any_element();
+
             let layout_id = element.request_layout(window, cx);
             (layout_id, element)
+        })
     }
 
     fn prepaint(
@@ -239,7 +242,9 @@ impl<C: RenderOnce> Element for Component<C> {
         window: &mut Window,
         cx: &mut App,
     ) {
+        window.with_global_id(ElementId::Name(type_name::<C>().into()), |_, window| {
             element.prepaint(window, cx);
+        })
     }
 
     fn paint(
@@ -252,7 +257,9 @@ impl<C: RenderOnce> Element for Component<C> {
         window: &mut Window,
         cx: &mut App,
     ) {
+        window.with_global_id(ElementId::Name(type_name::<C>().into()), |_, window| {
             element.paint(window, cx);
+        })
     }
 }
 

crates/gpui/src/gpui.rs

@@ -197,6 +197,11 @@ pub trait AppContext {
     where
         T: 'static;
 
+    /// Update a entity in the app context.
+    fn as_mut<'a, T>(&'a mut self, handle: &Entity<T>) -> Self::Result<GpuiBorrow<'a, T>>
+    where
+        T: 'static;
+
     /// Read a entity from the app context.
     fn read_entity<T, R>(
         &self,

crates/gpui/src/window.rs

@@ -2424,6 +2424,53 @@ impl Window {
         result
     }
 
+    /// Use a piece of state that exists as long this element is being rendered in consecutive frames.
+    pub fn use_keyed_state<S: 'static>(
+        &mut self,
+        key: impl Into<ElementId>,
+        cx: &mut App,
+        init: impl FnOnce(&mut Self, &mut App) -> S,
+    ) -> Entity<S> {
+        let current_view = self.current_view();
+        self.with_global_id(key.into(), |global_id, window| {
+            window.with_element_state(global_id, |state: Option<Entity<S>>, window| {
+                if let Some(state) = state {
+                    (state.clone(), state)
+                } else {
+                    let new_state = cx.new(|cx| init(window, cx));
+                    cx.observe(&new_state, move |_, cx| {
+                        cx.notify(current_view);
+                    })
+                    .detach();
+                    (new_state.clone(), new_state)
+                }
+            })
+        })
+    }
+
+    /// Immediately push an element ID onto the stack. Useful for simplifying IDs in lists
+    pub fn with_id<R>(&mut self, id: impl Into<ElementId>, f: impl FnOnce(&mut Self) -> R) -> R {
+        self.with_global_id(id.into(), |_, window| f(window))
+    }
+
+    /// Use a piece of state that exists as long this element is being rendered in consecutive frames, without needing to specify a key
+    ///
+    /// NOTE: This method uses the location of the caller to generate an ID for this state.
+    ///       If this is not sufficient to identify your state (e.g. you're rendering a list item),
+    ///       you can provide a custom ElementID using the `use_keyed_state` method.
+    #[track_caller]
+    pub fn use_state<S: 'static>(
+        &mut self,
+        cx: &mut App,
+        init: impl FnOnce(&mut Self, &mut App) -> S,
+    ) -> Entity<S> {
+        self.use_keyed_state(
+            ElementId::CodeLocation(*core::panic::Location::caller()),
+            cx,
+            init,
+        )
+    }
+
     /// Updates or initializes state for an element with the given id that lives across multiple
     /// frames. If an element with this ID existed in the rendered frame, its state will be passed
     /// to the given closure. The state returned by the closure will be stored so it can be referenced
@@ -4577,6 +4624,8 @@ pub enum ElementId {
     NamedInteger(SharedString, u64),
     /// A path.
     Path(Arc<std::path::Path>),
+    /// A code location.
+    CodeLocation(core::panic::Location<'static>),
 }
 
 impl ElementId {
@@ -4596,6 +4645,7 @@ impl Display for ElementId {
             ElementId::NamedInteger(s, i) => write!(f, "{}-{}", s, i)?,
             ElementId::Uuid(uuid) => write!(f, "{}", uuid)?,
             ElementId::Path(path) => write!(f, "{}", path.display())?,
+            ElementId::CodeLocation(location) => write!(f, "{}", location)?,
         }
 
         Ok(())

crates/gpui_macros/src/derive_app_context.rs

@@ -53,6 +53,16 @@ pub fn derive_app_context(input: TokenStream) -> TokenStream {
                 self.#app_variable.update_entity(handle, update)
             }
 
+            fn as_mut<'y, 'z, T>(
+                &'y mut self,
+                handle: &'z gpui::Entity<T>,
+            ) -> Self::Result<gpui::GpuiBorrow<'y, T>>
+            where
+                T: 'static,
+            {
+                self.#app_variable.as_mut(handle)
+            }
+
             fn read_entity<T, R>(
                 &self,
                 handle: &gpui::Entity<T>,