1 files changed, 50 insertions, 18 deletions

diff --git a/lazy/src/component.rs b/lazy/src/component.rs
index 4f1df650..d8f21f8a 100644
--- a/lazy/src/component.rs
+++ b/lazy/src/component.rs
@@ -11,7 +11,7 @@ use iced_native::{
 };
 
 use ouroboros::self_referencing;
-use std::cell::{Ref, RefCell};
+use std::cell::RefCell;
 use std::marker::PhantomData;
 
 /// A reusable, custom widget that uses The Elm Architecture.
@@ -46,6 +46,16 @@ pub trait Component<Message, Renderer> {
     /// Produces the widgets of the [`Component`], which may trigger an [`Event`](Component::Event)
     /// on user interaction.
     fn view(&self, state: &Self::State) -> Element<'_, Self::Event, Renderer>;
+
+    /// Update the [`Component`] state based on the provided [`Operation`](widget::Operation)
+    ///
+    /// By default, it does nothing.
+    fn operate(
+        &self,
+        _state: &mut Self::State,
+        _operation: &mut dyn widget::Operation<Message>,
+    ) {
+    }
 }
 
 /// Turns an implementor of [`Component`] into an [`Element`] that can be
@@ -106,6 +116,26 @@ where
         );
     }
 
+    fn rebuild_element_with_operation(
+        &self,
+        state: &mut S,
+        operation: &mut dyn widget::Operation<Message>,
+    ) {
+        let heads = self.state.borrow_mut().take().unwrap().into_heads();
+
+        heads.component.operate(state, operation);
+
+        *self.state.borrow_mut() = Some(
+            StateBuilder {
+                component: heads.component,
+                message: PhantomData,
+                state: PhantomData,
+                element_builder: |component| Some(component.view(state)),
+            }
+            .build(),
+        );
+    }
+
     fn with_element<T>(
         &self,
         f: impl FnOnce(&Element<'_, Event, Renderer>) -> T,
@@ -234,8 +264,14 @@ where
         &self,
         tree: &mut Tree,
         layout: Layout<'_>,
+        renderer: &Renderer,
         operation: &mut dyn widget::Operation<Message>,
     ) {
+        self.rebuild_element_with_operation(
+            tree.state.downcast_mut(),
+            operation,
+        );
+
         struct MapOperation<'a, B> {
             operation: &'a mut dyn widget::Operation<B>,
         }
@@ -274,6 +310,7 @@ where
             element.as_widget().operate(
                 &mut tree.children[0],
                 layout,
+                renderer,
                 &mut MapOperation { operation },
             );
         });
@@ -322,25 +359,25 @@ where
     }
 
     fn overlay<'b>(
-        &'b self,
+        &'b mut self,
         tree: &'b mut Tree,
         layout: Layout<'_>,
         renderer: &Renderer,
     ) -> Option<overlay::Element<'b, Message, Renderer>> {
         let overlay = OverlayBuilder {
             instance: self,
-            instance_ref_builder: |instance| instance.state.borrow(),
             tree,
             types: PhantomData,
             overlay_builder: |instance, tree| {
-                instance
-                    .as_ref()
-                    .unwrap()
-                    .borrow_element()
-                    .as_ref()
-                    .unwrap()
-                    .as_widget()
-                    .overlay(&mut tree.children[0], layout, renderer)
+                instance.state.get_mut().as_mut().unwrap().with_element_mut(
+                    move |element| {
+                        element.as_mut().unwrap().as_widget_mut().overlay(
+                            &mut tree.children[0],
+                            layout,
+                            renderer,
+                        )
+                    },
+                )
             },
         }
         .build();
@@ -362,15 +399,11 @@ where
 
 #[self_referencing]
 struct Overlay<'a, 'b, Message, Renderer, Event, S> {
-    instance: &'a Instance<'b, Message, Renderer, Event, S>,
+    instance: &'a mut Instance<'b, Message, Renderer, Event, S>,
     tree: &'a mut Tree,
     types: PhantomData<(Message, Event, S)>,
 
-    #[borrows(instance)]
-    #[covariant]
-    instance_ref: Ref<'this, Option<State<'a, Message, Renderer, Event, S>>>,
-
-    #[borrows(instance_ref, mut tree)]
+    #[borrows(mut instance, mut tree)]
     #[covariant]
     overlay: Option<overlay::Element<'this, Event, Renderer>>,
 }
@@ -514,7 +547,6 @@ where
             self.overlay = Some(
                 OverlayBuilder {
                     instance: overlay.instance,
-                    instance_ref_builder: |instance| instance.state.borrow(),
                     tree: overlay.tree,
                     types: PhantomData,
                     overlay_builder: |_, _| None,