How GPUI's Entity Model Solves Rust State Management (With a Ratatui Example)

Introduction

Gpui has been gaining a lot of popularity these last couple of months, and I am proud to say I have been a fan of the framework long before that. As a matter of fact, here is a tweet by me saying that people should talk about it more (when people didn't talk about it more):

In my opinion, two things are contributing to the popularity of the framework today:

• Performance: I think it is clear to everyone now that gpui is the best alternative to electron applications.
• AI: Many people hate using Rust because of the borrow checker, but with agents, they no longer have to deal with it as much.

What I really loved about the framework, however, is its approach to state management. When I first started building terminal user interfaces with Rust, I struggled with how I should approach state management, because with Rust (unlike JavaScript) the borrow checker makes things especially hard.

After seeing how gpui handled this, by studying the codebase and reading this gem of an article, I noticed that the approach adopted can be transferred to any other application which wants to manage state effectively.

In this blog post, I will share how I built an experimental state management crate for Rust, and I will use it to build an application with ratatui.

State Management Crate

Our state management crate creates a state. It provides components of the application with ways to read and update said state. In Rust, each variable has a unique owner. So, in order to avoid the complexity that comes with the borrow checker, the state is going to be owned by one App object.

Instead of multiple components holding references to the state owned by the App, the App gives out cheap handles to the components which want to read or update it. This is possible because the App stores its entities in a SlotMap data structure.

A SlotMap is an array which allows O(1) insertion and retrieval of elements by key. A key is returned after inserting a value, and the key will always refer to the value which was inserted. This makes it ideal for storing our application entities.

For our use case, the slotmap crate is used. With that dependency in place, we can start building our crate, beginning with the App struct that owns everything.

App

So, let us define the App struct as follows:

pub struct App {
  map: SlotMap<DefaultKey, Option<Box<dyn Any>>>
}

• DefaultKey is the type of the key returned after inserting values into the SlotMap.
• Values are Option<Box<dyn Any>>, since a state can be anything; wrapping it in an Option is useful for preventing reentrant updates.

Entity

Next, let's define what an Entity is. In gpui, Entity represents the handle that the App gives out. It basically contains the key of the value inserted into the slotmap, a ref_count to count how many times the entity was cloned, and a PhantomData which specifies which type the Entity holds.

pub struct Entity<T> {
  entity_id: DefaultKey,
  ref_count: Rc<()>,
  _marker: PhantomData<T> 
}

With this, we can implement a method to register new entities.

impl App {
  fn new<T: 'static>(&mut self, build: impl FnOnce(&App) -> T) -> Entity<T> {
    let state = build(self);
    let entity_id = self.map.insert(Some(Box::new(state)));
    Entity {
      entity_id,
      ref_count: Rc::new(()),
      _marker: PhantomData
    }
  }
}

As you can see, the method accepts a closure which is handed a reference to the App. This is because some entities might need to be initialised with data from other entities in App.

Let us now implement a method to read from a state.

impl App {
    pub fn read<T: 'static>(&self, entity: &Entity<T>) -> &T {
        let value = self.map.get(entity.entity_id).expect("stale handle");
        match value {
            Some(value) => value.downcast_ref::<T>().unwrap(),
            None => {
                panic!("cannot read a value which is currently being updated")
            }
        }
    }
}

SlotMap::get returns None when the key isn't present in the map. This happens if the entity was removed (and the slot's generation bumped), or if the handle is simply invalid. We expect here rather than propagate the error, because reaching this state should never happen in a correctly-behaving program. If it does, it signals a bug in our architecture, not a recoverable condition. Since this is an experimental crate, we're choosing to fail loudly rather than handle every edge case gracefully.

The next step is to implement an update method in order to update entities.

impl App {
    pub fn update<T: 'static, R>(
        &mut self,
        entity: Entity<T>,
        update: impl FnOnce(&mut T, &mut App) -> R,
    ) -> R {
        let value = self.map.get_mut(entity.entity_id).expect("stale handle");
        let boxed = value.take().expect("entity already being updated");
        let mut state: T = *boxed.downcast::<T>().unwrap();
        let result: R = update(&mut state, self);
        *self.map.get_mut(entity.entity_id).unwrap() = Some(Box::new(state));
        result
    }
}

The method receives a closure which is yielded a mutable reference to the entity it wants to update. The expect after take panics in case of reentrancy: if the entity is already being updated, we cannot call update on it again. This is why we chose to use Option as the value type in the SlotMap . It gives us a way to signal that this value is not available right now.

Dropping stale entities

When there are no more handles pointing to a state in the map, the state is marked as stale and should be removed.

To achieve this, we add a pending field to our App struct to keep track of the stale entities:

pub struct App {
  map: SlotMap<DefaultKey, Option<Box<dyn Any>>>,
  pending: Rc<RefCell<Vec<DefaultKey>>>
}

When an entity handle is dropped, it checks if it is the only handle remaining to the state. If no other handle remains, it pushes its entity_id onto the pending queue.

Entity should have access to the pending queue. We give it a Weak reference rather than an Rc, so that holding an Entity never keeps App's queue (and thus App itself) artificially alive:

pub struct Entity<T> {
  entity_id: DefaultKey,
  ref_count: Rc<()>,
  pending: Weak<RefCell<Vec<DefaultKey>>>,
  _marker: PhantomData<T>
}

With this new field in place, we need to update new_entity to populate it whenever an entity is created:

pub fn new_entity<T: 'static>(&mut self, build: impl FnOnce(&App) -> T) -> Entity<T> {
        let entity = build(self);
        let entity_id = self.map.insert(Some(Box::new(entity)));
        Entity {
            entity_id,
            ref_count: Rc::new(()),
            pending: Rc::downgrade(&self.pending),
            _marker: PhantomData,
        }
    }

Then we implement Drop for Entity, so that when the last handle to a state goes out of scope, its slot is queued for removal:

impl<T> Drop for Entity<T> {
    fn drop(&mut self) {
        if Rc::strong_count(&self.ref_count) == 1
            && let Some(queue) = self.pending.upgrade()
        {
            queue.borrow_mut().push(self.entity_id);
        }
    }
}

Finally, we implement a method called flush_entities which drains the queue and actually removes those stale slots from the map:

impl App {
  pub fn flush_entities(&mut self) {
        let ids: Vec<DefaultKey> = self.pending.borrow_mut().drain(..).collect();
        for id in ids {
            self.map.remove(id);
        }
  }
}

Usage

Let's test our framework by creating a simple Ratatui application. In the majority of applications, one component updates a state, and another one reads from that state.

In our test application, we will demonstrate this by creating two widgets. One will act as an input and update an input_buffer. The second widget will just read the input_buffer and reflect its changes (In real world applications, this second widget might use the input to filter search results for example).

We will start by creating a Root struct. The Root struct holds a reference to an App struct:

struct Root<'a> {
  app: &'a mut App,
  widgets: Vec<Box<dyn Component>>
}

The Root also holds a list of widgets, defined using the Component trait. We define the Component trait as follows:

trait Component {
    fn render_with_cx(&mut self, cx: &mut App, frame: &mut Frame, area: Rect);
    fn handle_key(&mut self, cx: &mut App, key: KeyEvent);
}

Each component uses render_with_cx to draw on the frame while still having access to the context. Components also define a handle_key method in order to handle events.

Let us create our two widgets with their implementations of Component.

struct InputWidget {
    input_buffer: Entity<String>,
}

struct EchoWidget {
    input_buffer: Entity<String>,
}

Each widget owns an Entity<String>. The input widget updates input_buffer and the echo widget reads it and displays it.

impl InputWidget {
    fn init(input_buffer: Entity<String>) -> Self {
        Self { input_buffer }
    }
}

impl EchoWidget {
    fn init(input_buffer: Entity<String>) -> Self {
        Self { input_buffer }
    }
}

impl Component for InputWidget {
    fn render_with_cx(&mut self, cx: &mut App, frame: &mut Frame, area: Rect) {
        let block = Block::bordered()
            .title("Input")
            .padding(Padding::uniform(1));
        let inner = block.inner(area);
        block.render(area, frame.buffer_mut());

        let input_buffer = cx.read(&self.input_buffer);
        let text_width = input_buffer.chars().count() as u16;

        frame.set_cursor_position(Position::new(inner.x + text_width, inner.y));
        Paragraph::new(&input_buffer[..])
            .wrap(Wrap { trim: false })
            .render(inner, frame.buffer_mut());
    }

    fn handle_key(&mut self, cx: &mut App, key: KeyEvent) {
        match key.code {
            KeyCode::Char(c) => cx.update(self.input_buffer.clone(), |input_buffer, _cx| {
                input_buffer.push(c);
            }),
            KeyCode::Backspace => cx.update(self.input_buffer.clone(), |input_buffer, _cx| {
                if input_buffer.len() > 0 {
                    input_buffer.remove(input_buffer.len() - 1);
                }
            }),
            _ => {}
        }
    }
}

In render_with_cx we render a block and a paragraph. To learn more about widgets and how to render them, you can check the ratatui documentation. The interesting part is that the paragraph is initialised with the input_buffer contents, thus the call to cx.read(&self.input_buffer).

In handle_key we listen to key events and then update our input_buffer state with cx.update().

The EchoWidget component only reads the input buffer entity and displays it:

impl Component for EchoWidget {
    fn render_with_cx(&mut self, cx: &mut App, frame: &mut Frame, area: Rect) {
        let block = Block::bordered()
            .title("Echo")
            .padding(Padding::uniform(1));
        let inner = block.inner(area);
        block.render(area, frame.buffer_mut());

        let input_buffer = cx.read(&self.input_buffer);
        Paragraph::new(&input_buffer[..])
            .wrap(Wrap { trim: false })
            .render(inner, frame.buffer_mut());
    }

    // Echo widget is read-only: it never receives input focus or key events.
    fn handle_key(&mut self, _cx: &mut App, _key: KeyEvent) {}
}

Entity creation happens when initialising the Root struct by calling cx.new_entity:

impl<'a> Root<'a> {
    fn init(cx: &'a mut App) -> Root<'a> {
        let input_buffer = cx.new_entity(|_| "Enter Text".into());
        Root {
            app: cx,
            widgets: vec![
                Box::new(InputWidget::init(input_buffer.clone())),
                Box::new(EchoWidget::init(input_buffer.clone())),
            ],
        }
    }
}

We also define a run method for Root which does the drawing by calling the render_with_cx method on each component, and listens for events by calling handle_key_event on each component in the root.

impl<'a> Root<'a> {
    fn run(&mut self, terminal: &mut DefaultTerminal) -> anyhow::Result<()> {
        loop {
            self.app.flush_entities();
            let _ = terminal.draw(|frame| {
                let [top, bottom] =
                    Layout::vertical([Constraint::Percentage(50), Constraint::Percentage(50)])
                        .areas(frame.area());
                self.widgets[0].render_with_cx(self.app, frame, top);
                self.widgets[1].render_with_cx(self.app, frame, bottom);
            });

            if let Event::Key(key_event) = event::read()? {
                if !key_event.is_press() {
                    continue;
                }
                let is_quit = key_event.code == KeyCode::Esc
                    || (key_event.code == KeyCode::Char('c')
                        && key_event.modifiers.contains(KeyModifiers::CONTROL));
                if is_quit {
                    break Ok(());
                }
                // Only the input widget (top, index 0) is interactive.
                self.widgets[0].handle_key(self.app, key_event);
            }
        }
    }
}

Finally, we define our main function:

fn main() {
    let mut app: App = App::init();
    let mut root = Root::init(&mut app);
    let _ = ratatui::run(|terminal| root.run(terminal));
}

Conclusion

Using our state management crate inspired by gpui, we were able to effectively manage state without dealing with the complexities that the borrow checker brings. Rather than wrapping each piece of shared state in an Rc<RefCell<T>>, as is common in Rust UI code, our state lives in a single SlotMap owned by the App. Components simply hold cheap Entity<T> handles into that map, instead of sharing mutable references directly.

The state management crate we developed here is experimental. It could be enhanced with other features, mainly:

• Only redraw when the UI is marked as dirty, instead of each frame.
• Components subscribing to/observing a certain state. (We didn't have to deal with this here, since we redraw the entire UI each tick.)

Code is available on GitHub:

• State Management Crate
• Test Application