fix: improve component recycling algorithm

packages/iocraft/Cargo.toml

@@ -16,8 +16,6 @@ unicode-width = "0.1.13"
 textwrap = "0.16.1"
 generational-box = "0.5.6"
 any_key = "0.1.1"
-uuid = { version = "1.10.0", features = ["v4"] }
-indexmap = "2.5.0"
 
 [dev-dependencies]
 indoc = "2"

packages/iocraft/src/component.rs

@@ -2,6 +2,7 @@ use crate::{
     context::ContextStack,
     element::{ElementKey, ElementType},
     hook::{AnyHook, Hook, Hooks},
+    multimap::RemoveOnlyMultimap,
     props::{AnyProps, Props},
     render::{ComponentDrawer, ComponentUpdater, UpdateContext},
 };
@@ -12,7 +13,6 @@ use core::{
     task::{Context, Poll},
 };
 use futures::future::poll_fn;
-use indexmap::IndexMap;
 use taffy::NodeId;
 
 pub(crate) struct ComponentHelper<C: Component> {
@@ -189,7 +189,7 @@ impl InstantiatedComponent {
         if self.has_transparent_layout {
             // If the component has a transparent layout, provide the first child's layout to the
             // hooks and component.
-            if let Some((_, child)) = self.children.components.iter().next().as_ref() {
+            if let Some(child) = self.children.components.iter().next().as_ref() {
                 drawer.for_child_node_layout(child.node_id, |drawer| {
                     self.hooks.pre_component_draw(drawer);
                     self.component.draw(drawer);
@@ -206,7 +206,7 @@ impl InstantiatedComponent {
         self.children.draw(drawer);
 
         if self.has_transparent_layout {
-            if let Some((_, child)) = self.children.components.iter().next().as_ref() {
+            if let Some(child) = self.children.components.iter().next().as_ref() {
                 drawer.for_child_node_layout(child.node_id, |drawer| {
                     self.hooks.post_component_draw(drawer);
                 });
@@ -237,12 +237,12 @@ impl InstantiatedComponent {
 
 #[derive(Default)]
 pub(crate) struct Components {
-    pub components: IndexMap<ElementKey, InstantiatedComponent>,
+    pub components: RemoveOnlyMultimap<ElementKey, InstantiatedComponent>,
 }
 
 impl Components {
     pub fn draw(&mut self, drawer: &mut ComponentDrawer<'_>) {
-        for (_, component) in self.components.iter_mut() {
+        for component in self.components.iter_mut() {
             if component.has_transparent_layout {
                 component.draw(drawer);
             } else {
@@ -255,7 +255,7 @@ impl Components {
 
     pub fn poll_change(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
         let mut is_ready = false;
-        for component in self.components.values_mut() {
+        for component in self.components.iter_mut() {
             if Pin::new(&mut *component).poll_change(cx).is_ready() {
                 is_ready = true;
             }

packages/iocraft/src/lib.rs

@@ -99,6 +99,7 @@ mod context;
 mod element;
 mod handler;
 mod hook;
+mod multimap;
 mod props;
 mod render;
 mod style;

packages/iocraft/src/multimap.rs

@@ -0,0 +1,73 @@
+use std::{
+    collections::{HashMap, VecDeque},
+    hash::Hash,
+};
+
+/// This is a specialized map implementation which is optimized for tracking components across updates:
+///
+/// During updates, components are created and appended to the `AppendOnlyMultimap`. Once the
+/// update is complete, the map is converted to a `RemoveOnlyMultimap`, which can be iterated in
+/// insertion order. During the next update, components are removed from the map based on their key
+/// in order to be recycled. If multiple elements have duplicate keys, they're recycled in the same
+/// order they were first inserted.
+///
+/// All operations have O(1) time complexity.
+pub(crate) struct AppendOnlyMultimap<K, V> {
+    items: Vec<Option<V>>,
+    m: HashMap<K, VecDeque<usize>>,
+}
+
+impl<K, V> Default for AppendOnlyMultimap<K, V> {
+    fn default() -> Self {
+        Self {
+            items: Vec::new(),
+            m: HashMap::new(),
+        }
+    }
+}
+
+impl<K: Hash + Eq, V> AppendOnlyMultimap<K, V> {
+    pub fn push_back(&mut self, key: K, value: V) {
+        let index = self.items.len();
+        self.items.push(Some(value));
+        self.m.entry(key).or_default().push_back(index);
+    }
+}
+
+pub(crate) struct RemoveOnlyMultimap<K, V> {
+    items: Vec<Option<V>>,
+    m: HashMap<K, VecDeque<usize>>,
+}
+
+impl<K, V> Default for RemoveOnlyMultimap<K, V> {
+    fn default() -> Self {
+        Self {
+            items: Vec::new(),
+            m: HashMap::new(),
+        }
+    }
+}
+
+impl<K: Hash + Eq, V> RemoveOnlyMultimap<K, V> {
+    pub fn pop_front(&mut self, key: &K) -> Option<V> {
+        let index = self.m.get_mut(key)?.pop_front()?;
+        self.items[index].take()
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &V> {
+        self.items.iter().filter_map(|item| item.as_ref())
+    }
+
+    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut V> {
+        self.items.iter_mut().filter_map(|item| item.as_mut())
+    }
+}
+
+impl<K, V> From<AppendOnlyMultimap<K, V>> for RemoveOnlyMultimap<K, V> {
+    fn from(multimap: AppendOnlyMultimap<K, V>) -> Self {
+        Self {
+            items: multimap.items,
+            m: multimap.m,
+        }
+    }
+}

packages/iocraft/src/render.rs

@@ -2,14 +2,14 @@ use crate::{
     canvas::{Canvas, CanvasSubviewMut},
     component::{ComponentHelperExt, Components, InstantiatedComponent},
     context::{Context, ContextStack, SystemContext},
-    element::{ElementExt, ElementKey},
+    element::ElementExt,
+    multimap::AppendOnlyMultimap,
     props::AnyProps,
     terminal::{MockTerminalConfig, MockTerminalOutputStream, Terminal, TerminalEvents},
 };
 use core::{
     any::Any,
     cell::{Ref, RefMut},
-    mem,
     pin::Pin,
     task::{self, Poll},
 };
@@ -18,10 +18,8 @@ use futures::{
     future::{select, FutureExt, LocalBoxFuture},
     stream::{Stream, StreamExt},
 };
-use indexmap::IndexMap;
 use std::io;
 use taffy::{AvailableSpace, Layout, NodeId, Point, Size, Style, TaffyTree};
-use uuid::Uuid;
 
 pub(crate) struct UpdateContext<'a> {
     terminal: Option<&'a mut Terminal>,
@@ -140,7 +138,7 @@ impl<'a, 'b, 'c> ComponentUpdater<'a, 'b, 'c> {
     {
         self.component_context_stack
             .with_context(context, |component_context_stack| {
-                let mut used_components = IndexMap::with_capacity(self.children.components.len());
+                let mut used_components = AppendOnlyMultimap::default();
 
                 let mut direct_child_node_ids = Vec::new();
                 let child_node_ids = if self.transparent_layout {
@@ -153,7 +151,7 @@ impl<'a, 'b, 'c> ComponentUpdater<'a, 'b, 'c> {
 
                 for mut child in children {
                     let mut component: InstantiatedComponent =
-                        match self.children.components.swap_remove(child.key()) {
+                        match self.children.components.pop_front(child.key()) {
                             Some(component)
                                 if component.component().type_id()
                                     == child.helper().component_type_id() =>
@@ -182,25 +180,21 @@ impl<'a, 'b, 'c> ComponentUpdater<'a, 'b, 'c> {
                         child.props_mut(),
                     );
 
-                    let mut child_key = child.key().clone();
-                    while used_components.contains_key(&child_key) {
-                        child_key = ElementKey::new(Uuid::new_v4().as_u128());
-                    }
-                    used_components.insert(child_key, component);
+                    used_components.push_back(child.key().clone(), component);
                 }
 
                 self.context
                     .layout_engine
                     .set_children(self.node_id, &direct_child_node_ids)
                     .expect("we should be able to set the children");
 
-                for (_, component) in self.children.components.drain(..) {
+                for component in self.children.components.iter() {
                     self.context
                         .layout_engine
                         .remove(component.node_id())
                         .expect("we should be able to remove the node");
                 }
-                mem::swap(&mut self.children.components, &mut used_components);
+                self.children.components = used_components.into();
             });
     }
 }
@@ -507,9 +501,7 @@ mod tests {
             Box(flex_direction: FlexDirection::Column) {
                 Text(content: format!("tick: {}", tick))
                 MyInnerComponent(label: "a")
-                // without a key, these next elements may not be re-used across renders
                 #((0..2).map(|i| element! { MyInnerComponent(label: format!("b{}", i)) }))
-                // with a key, these next elements will definitely be re-used across renders
                 #((0..2).map(|i| element! { MyInnerComponent(key: i, label: format!("c{}", i)) }))
             }
         }
@@ -524,7 +516,7 @@ mod tests {
         let actual = canvases.iter().map(|c| c.to_string()).collect::<Vec<_>>();
         let expected = vec![
             "tick: 0\nrender count (a): 1\nrender count (b0): 1\nrender count (b1): 1\nrender count (c0): 1\nrender count (c1): 1\n",
-            "tick: 1\nrender count (a): 2\nrender count (b0): 2\nrender count (b1): 1\nrender count (c0): 2\nrender count (c1): 2\n",
+            "tick: 1\nrender count (a): 2\nrender count (b0): 2\nrender count (b1): 2\nrender count (c0): 2\nrender count (c1): 2\n",
         ];
         assert_eq!(actual, expected);
     }
@@ -559,4 +551,64 @@ mod tests {
         .to_string();
         assert_eq!(actual, "+--------+\n+--------+\n",);
     }
+
+    #[derive(Default, Props)]
+    struct AsyncTickerProps {
+        ticks: Option<State<i32>>,
+    }
+
+    #[component]
+    fn AsyncTicker<'a>(
+        props: &mut AsyncTickerProps,
+        mut hooks: Hooks,
+    ) -> impl Into<AnyElement<'a>> {
+        let mut ticks = props.ticks.unwrap();
+        hooks.use_future(async move {
+            ticks += 1;
+        });
+        element!(Box)
+    }
+
+    #[component]
+    fn AsyncTickerContainer(mut hooks: Hooks) -> impl Into<AnyElement<'static>> {
+        let mut system = hooks.use_context_mut::<SystemContext>();
+        let child_ticks = hooks.use_state(|| 0);
+        let mut tick = hooks.use_state(|| 0);
+
+        hooks.use_future(async move {
+            tick += 1;
+        });
+
+        if tick == 5 {
+            // make sure our children have all ticked exactly 10 times
+            assert_eq!(child_ticks, 10);
+            system.exit();
+        } else {
+            // do a few more render passes
+            tick += 1;
+        }
+
+        element! {
+            Box {
+                #((0..10).map(|_| {
+                    element! {
+                        AsyncTicker(ticks: child_ticks)
+                    }
+                }))
+            }
+        }
+    }
+
+    // This is a regression test for an issue where elements added via iterator without keys would
+    // be re-created on every render instead of being recycled.
+    #[apply(test!)]
+    async fn test_async_ticker_container() {
+        let canvases: Vec<_> = mock_terminal_render_loop(
+            element!(AsyncTickerContainer),
+            MockTerminalConfig::default(),
+        )
+        .collect()
+        .await;
+        assert!(canvases.len() > 0);
+    }
 }