From fa3bd4280e7c34354db0a15ef1baba9fbbeb87f1 Mon Sep 17 00:00:00 2001
From: Héctor Ramón Jiménez <hector0193@gmail.com>
Date: Mon, 7 Mar 2022 16:33:25 +0700
Subject: Initialize `lazy::pure::Component` view properly

... during `children` and `diff` in the `Widget` implementation
---
 lazy/src/pure/component.rs | 47 +++++++++++++++++++++++++---------------------
 1 file changed, 26 insertions(+), 21 deletions(-)

(limited to 'lazy/src/pure/component.rs')

diff --git a/lazy/src/pure/component.rs b/lazy/src/pure/component.rs
index c79e34f8..ad82c55e 100644
--- a/lazy/src/pure/component.rs
+++ b/lazy/src/pure/component.rs
@@ -101,27 +101,6 @@ where
         &self,
         f: impl FnOnce(&mut Element<'_, Event, Renderer>) -> T,
     ) -> T {
-        if self
-            .state
-            .borrow()
-            .as_ref()
-            .unwrap()
-            .borrow_element()
-            .is_none()
-        {
-            let heads = self.state.borrow_mut().take().unwrap().into_heads();
-
-            *self.state.borrow_mut() = Some(
-                StateBuilder {
-                    component: heads.component,
-                    message: PhantomData,
-                    state: PhantomData,
-                    element_builder: |state| Some(state.view(&S::default())),
-                }
-                .build(),
-            );
-        }
-
         self.state
             .borrow_mut()
             .as_mut()
@@ -145,10 +124,36 @@ where
     }
 
     fn children(&self) -> Vec<Tree> {
+        let heads = self.state.borrow_mut().take().unwrap().into_heads();
+
+        *self.state.borrow_mut() = Some(
+            StateBuilder {
+                component: heads.component,
+                message: PhantomData,
+                state: PhantomData,
+                element_builder: |state| Some(state.view(&S::default())),
+            }
+            .build(),
+        );
+
         self.with_element(|element| vec![Tree::new(element)])
     }
 
     fn diff(&self, tree: &mut Tree) {
+        let heads = self.state.borrow_mut().take().unwrap().into_heads();
+
+        *self.state.borrow_mut() = Some(
+            StateBuilder {
+                component: heads.component,
+                message: PhantomData,
+                state: PhantomData,
+                element_builder: |state| {
+                    Some(state.view(tree.state.downcast_ref()))
+                },
+            }
+            .build(),
+        );
+
         self.with_element(|element| {
             tree.diff_children(std::slice::from_ref(&element))
         })
-- 
cgit