feat: make hydration more robust

packages/web/src/dom.rs

@@ -22,11 +22,11 @@ pub struct WebsysDom {
     stack: Stack,
 
     /// A map from ElementID (index) to Node
-    nodes: NodeSlab,
+    pub(crate) nodes: NodeSlab,
 
     document: Document,
 
-    root: Element,
+    pub(crate) root: Element,
 
     sender_callback: Rc<dyn Fn(SchedulerMsg)>,
 
@@ -34,45 +34,20 @@ pub struct WebsysDom {
     // This is roughly a delegater
     // TODO: check how infero delegates its events - some are more performant
     listeners: FxHashMap<&'static str, ListenerEntry>,
-
-    // We need to make sure to add comments between text nodes
-    // We ensure that the text siblings are patched by preventing the browser from merging
-    // neighboring text nodes. Originally inspired by some of React's work from 2016.
-    //  -> https://reactjs.org/blog/2016/04/07/react-v15.html#major-changes
-    //  -> https://github.com/facebook/react/pull/5753
-    last_node_was_text: bool,
 }
 
 type ListenerEntry = (usize, Closure<dyn FnMut(&Event)>);
 
 impl WebsysDom {
-    pub fn new(root: Element, cfg: WebConfig, sender_callback: Rc<dyn Fn(SchedulerMsg)>) -> Self {
+    pub fn new(cfg: WebConfig, sender_callback: Rc<dyn Fn(SchedulerMsg)>) -> Self {
         let document = load_document();
 
         let nodes = NodeSlab::new(2000);
         let listeners = FxHashMap::default();
 
-        // re-hydrate the page - only supports one virtualdom per page
-        // hydration is the dubmest thing you've ever heard of
-        // just blast away the page and replace it completely.
-        if cfg.hydrate {
-            // // Load all the elements into the arena
-            // let node_list: NodeList = document.query_selector_all("dioxus-id").unwrap();
-            // let len = node_list.length() as usize;
-
-            // for x in 0..len {
-            //     let node: Node = node_list.get(x as u32).unwrap();
-            //     let el: &Element = node.dyn_ref::<Element>().unwrap();
-            //     let id: String = el.get_attribute("dioxus-id").unwrap();
-            //     let id = id.parse::<usize>().unwrap();
-            //     nodes[id] = Some(node);
-            // }
-
-            // Load all the event listeners into our listener register
-            // TODO
-        }
-
         let mut stack = Stack::with_capacity(10);
+
+        let root = load_document().get_element_by_id(&cfg.rootname).unwrap();
         let root_node = root.clone().dyn_into::<Node>().unwrap();
         stack.push(root_node);
 
@@ -83,15 +58,13 @@ impl WebsysDom {
             document,
             sender_callback,
             root,
-            last_node_was_text: false,
         }
     }
 
-    pub fn process_edits(&mut self, edits: &mut Vec<DomEdit>) {
+    pub fn apply_edits(&mut self, mut edits: Vec<DomEdit>) {
         for edit in edits.drain(..) {
             match edit {
                 DomEdit::PushRoot { root } => self.push(root),
-                DomEdit::PopRoot => self.pop(),
                 DomEdit::AppendChildren { many } => self.append_children(many),
                 DomEdit::ReplaceWith { m, root } => self.replace_with(m, root),
                 DomEdit::Remove { root } => self.remove(root),
@@ -137,11 +110,6 @@ impl WebsysDom {
         self.stack.push(real_node);
     }
 
-    // drop the node off the stack
-    fn pop(&mut self) {
-        self.stack.pop();
-    }
-
     fn append_children(&mut self, many: u32) {
         let root: Node = self
             .stack
@@ -150,13 +118,23 @@ impl WebsysDom {
             .unwrap()
             .clone();
 
+        // We need to make sure to add comments between text nodes
+        // We ensure that the text siblings are patched by preventing the browser from merging
+        // neighboring text nodes. Originally inspired by some of React's work from 2016.
+        //  -> https://reactjs.org/blog/2016/04/07/react-v15.html#major-changes
+        //  -> https://github.com/facebook/react/pull/5753
+        /*
+        todo: we need to track this for replacing/insert after/etc
+        */
+        let mut last_node_was_text = false;
+
         for child in self
             .stack
             .list
             .drain((self.stack.list.len() - many as usize)..)
         {
             if child.dyn_ref::<web_sys::Text>().is_some() {
-                if self.last_node_was_text {
+                if last_node_was_text {
                     let comment_node = self
                         .document
                         .create_comment("dioxus")
@@ -164,9 +142,9 @@ impl WebsysDom {
                         .unwrap();
                     root.append_child(&comment_node).unwrap();
                 }
-                self.last_node_was_text = true;
+                last_node_was_text = true;
             } else {
-                self.last_node_was_text = false;
+                last_node_was_text = false;
             }
             root.append_child(&child).unwrap();
         }

packages/web/src/lib.rs

@@ -55,7 +55,6 @@
 use std::rc::Rc;
 
 pub use crate::cfg::WebConfig;
-use crate::dom::load_document;
 use dioxus::SchedulerMsg;
 use dioxus::VirtualDom;
 pub use dioxus_core as dioxus;
@@ -66,6 +65,7 @@ mod cache;
 mod cfg;
 mod dom;
 mod nodeslab;
+mod rehydrate;
 mod ric_raf;
 
 /// Launch the VirtualDOM given a root component and a configuration.
@@ -146,24 +146,39 @@ pub async fn run_with_props<T: 'static + Send>(root: Component<T>, root_props: T
         wasm_bindgen::intern(s);
     }
 
-    let should_hydrate = cfg.hydrate;
-
-    let root_el = load_document().get_element_by_id(&cfg.rootname).unwrap();
-
     let tasks = dom.get_scheduler_channel();
 
     let sender_callback: Rc<dyn Fn(SchedulerMsg)> =
         Rc::new(move |event| tasks.unbounded_send(event).unwrap());
 
-    let mut websys_dom = dom::WebsysDom::new(root_el, cfg, sender_callback);
+    let should_hydrate = cfg.hydrate;
+
+    let mut websys_dom = dom::WebsysDom::new(cfg, sender_callback);
 
     log::trace!("rebuilding app");
-    let mut mutations = dom.rebuild();
 
-    // hydrating is simply running the dom for a single render. If the page is already written, then the corresponding
-    // ElementIds should already line up because the web_sys dom has already loaded elements with the DioxusID into memory
-    if !should_hydrate {
-        websys_dom.process_edits(&mut mutations.edits);
+    if should_hydrate {
+        // todo: we need to split rebuild and initialize into two phases
+        // it's a waste to produce edits just to get the vdom loaded
+        let _ = dom.rebuild();
+
+        if let Err(err) = websys_dom.rehydrate(&dom) {
+            log::error!(
+                "Rehydration failed {:?}. Rebuild DOM into element from scratch",
+                &err
+            );
+
+            websys_dom.root.set_text_content(None);
+
+            // errrrr we should split rebuild into two phases
+            // one that initializes things and one that produces edits
+            let edits = dom.rebuild();
+
+            websys_dom.apply_edits(edits.edits);
+        }
+    } else {
+        let edits = dom.rebuild();
+        websys_dom.apply_edits(edits.edits);
     }
 
     let work_loop = ric_raf::RafLoop::new();
@@ -185,9 +200,9 @@ pub async fn run_with_props<T: 'static + Send>(root: Component<T>, root_props: T
         // wait for the animation frame to fire so we can apply our changes
         work_loop.wait_for_raf().await;
 
-        for mut edit in mutations {
+        for edit in mutations {
             // actually apply our changes during the animation frame
-            websys_dom.process_edits(&mut edit.edits);
+            websys_dom.apply_edits(edit.edits);
         }
     }
 }

packages/web/src/rehydrate.rs

@@ -0,0 +1,157 @@
+use crate::dom::WebsysDom;
+use dioxus_core::{VNode, VirtualDom};
+use wasm_bindgen::JsCast;
+use web_sys::{Comment, Element, Node, Text};
+
+#[derive(Debug)]
+pub enum RehydrationError {
+    NodeTypeMismatch,
+    NodeNotFound,
+    VNodeNotInitialized,
+}
+use RehydrationError::*;
+
+impl WebsysDom {
+    // we're streaming in patches, but the nodes already exist
+    // so we're just going to write the correct IDs to the node and load them in
+    pub fn rehydrate(&mut self, dom: &VirtualDom) -> Result<(), RehydrationError> {
+        let root = self
+            .root
+            .clone()
+            .dyn_into::<Node>()
+            .map_err(|_| NodeTypeMismatch)?;
+
+        let root_scope = dom.base_scope();
+        let root_node = root_scope.root_node();
+
+        let mut nodes = vec![root];
+        let mut counter = vec![0];
+
+        let mut last_node_was_text = false;
+
+        // Recursively rehydrate the dom from the VirtualDom
+        self.rehydrate_single(
+            &mut nodes,
+            &mut counter,
+            dom,
+            root_node,
+            &mut last_node_was_text,
+        )
+    }
+
+    fn rehydrate_single(
+        &mut self,
+        nodes: &mut Vec<Node>,
+        place: &mut Vec<u32>,
+        dom: &VirtualDom,
+        node: &VNode,
+        last_node_was_text: &mut bool,
+    ) -> Result<(), RehydrationError> {
+        match node {
+            VNode::Text(t) => {
+                let node_id = t.id.get().ok_or(VNodeNotInitialized)?;
+
+                let cur_place = place.last_mut().unwrap();
+
+                // skip over the comment element
+                if *last_node_was_text {
+                    if cfg!(debug_assertions) {
+                        let node = nodes.last().unwrap().child_nodes().get(*cur_place).unwrap();
+                        let node_text = node.dyn_into::<Comment>().unwrap();
+                        assert_eq!(node_text.data(), "spacer");
+                    }
+                    *cur_place += 1;
+                }
+
+                let node = nodes
+                    .last()
+                    .unwrap()
+                    .child_nodes()
+                    .get(*cur_place)
+                    .ok_or(NodeNotFound)?;
+
+                let _text_el = node.dyn_ref::<Text>().ok_or(NodeTypeMismatch)?;
+
+                // in debug we make sure the text is the same
+                if cfg!(debug_assertions) {
+                    let contents = _text_el.node_value().unwrap();
+                    assert_eq!(t.text, contents);
+                }
+
+                *last_node_was_text = true;
+
+                self.nodes[node_id.0] = Some(node);
+
+                *cur_place += 1;
+            }
+
+            VNode::Element(vel) => {
+                let node_id = vel.id.get().ok_or(VNodeNotInitialized)?;
+
+                let cur_place = place.last_mut().unwrap();
+
+                let node = nodes.last().unwrap().child_nodes().get(*cur_place).unwrap();
+
+                use smallstr::SmallString;
+                use std::fmt::Write;
+
+                // 8 digits is enough, yes?
+                // 12 million nodes in one page?
+                let mut s: SmallString<[u8; 8]> = smallstr::SmallString::new();
+                write!(s, "{}", node_id).unwrap();
+
+                node.dyn_ref::<Element>()
+                    .unwrap()
+                    .set_attribute("dioxus-id", s.as_str())
+                    .unwrap();
+
+                self.nodes[node_id.0] = Some(node.clone());
+
+                *cur_place += 1;
+
+                nodes.push(node.clone());
+
+                place.push(0);
+
+                // we cant have the last node be text
+                let mut last_node_was_text = false;
+                for child in vel.children {
+                    self.rehydrate_single(nodes, place, dom, &child, &mut last_node_was_text)?;
+                }
+
+                place.pop();
+                nodes.pop();
+
+                if cfg!(debug_assertions) {
+                    let el = node.dyn_ref::<Element>().unwrap();
+                    let name = el.tag_name().to_lowercase();
+                    assert_eq!(name, vel.tag);
+                }
+            }
+
+            VNode::Placeholder(el) => {
+                let node_id = el.id.get().ok_or(VNodeNotInitialized)?;
+
+                let cur_place = place.last_mut().unwrap();
+                let node = nodes.last().unwrap().child_nodes().get(*cur_place).unwrap();
+
+                self.nodes[node_id.0] = Some(node);
+
+                *cur_place += 1;
+            }
+
+            VNode::Fragment(el) => {
+                for el in el.children {
+                    self.rehydrate_single(nodes, place, dom, &el, last_node_was_text)?;
+                }
+            }
+
+            VNode::Component(el) => {
+                let scope = dom.get_scope(el.scope.get().unwrap()).unwrap();
+                let node = scope.root_node();
+                self.rehydrate_single(nodes, place, dom, node, last_node_was_text)?;
+            }
+        }
+        Ok(())
+    }
+}