From c1b36b44d485a60e9195b8d70d314603929b99ae Mon Sep 17 00:00:00 2001
From: mc1098 <m.cripps1@uni.brighton.ac.uk>
Date: Sat, 25 Sep 2021 10:40:27 +0100
Subject: [PATCH] Add safe first_node fn

`first_node` function can fail and in some situations, when the VNode
has not been mounted yet, so this adds a safe version that returns an
`Option<Node>` and refactors the previous version to use the "unchecked"
prefix as it can and should panic if the first node cannot be found.
---
 packages/yew/src/virtual_dom/vcomp.rs |  2 +-
 packages/yew/src/virtual_dom/vnode.rs | 25 ++++++++++++++++++++++---
 packages/yew/src/virtual_dom/vtag.rs  |  2 +-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/packages/yew/src/virtual_dom/vcomp.rs b/packages/yew/src/virtual_dom/vcomp.rs
index ac340bcbef5..ae33995e4b7 100644
--- a/packages/yew/src/virtual_dom/vcomp.rs
+++ b/packages/yew/src/virtual_dom/vcomp.rs
@@ -391,7 +391,7 @@ mod tests {
         let test_node: Node = document().create_text_node("test").into();
         let test_node_ref = NodeRef::new(test_node);
         let check_node_ref = |vnode: VNode| {
-            assert_eq!(vnode.first_node(), test_node_ref.get().unwrap());
+            assert_eq!(vnode.unchecked_first_node(), test_node_ref.get().unwrap());
         };
 
         let props = Props {
diff --git a/packages/yew/src/virtual_dom/vnode.rs b/packages/yew/src/virtual_dom/vnode.rs
index 49bbd7ca378..518b86a2760 100644
--- a/packages/yew/src/virtual_dom/vnode.rs
+++ b/packages/yew/src/virtual_dom/vnode.rs
@@ -6,6 +6,7 @@ use gloo::console;
 use std::cmp::PartialEq;
 use std::fmt;
 use std::iter::FromIterator;
+use wasm_bindgen::JsCast;
 
 use web_sys::{Element, Node};
 
@@ -45,8 +46,23 @@ impl VNode {
         }
     }
 
+    /// Returns the first DOM node if available
+    pub(crate) fn first_node(&self) -> Option<Node> {
+        match self {
+            VNode::VTag(vtag) => vtag.reference().cloned().map(JsCast::unchecked_into),
+            VNode::VText(vtext) => vtext
+                .reference
+                .as_ref()
+                .cloned()
+                .map(JsCast::unchecked_into),
+            VNode::VComp(vcomp) => vcomp.node_ref.get(),
+            VNode::VList(vlist) => vlist.get(0).and_then(VNode::first_node),
+            VNode::VRef(node) => Some(node.clone()),
+        }
+    }
+
     /// Returns the first DOM node that is used to designate the position of the virtual DOM node.
-    pub(crate) fn first_node(&self) -> Node {
+    pub(crate) fn unchecked_first_node(&self) -> Node {
         match self {
             VNode::VTag(vtag) => vtag
                 .reference()
@@ -67,7 +83,10 @@ impl VNode {
                     crate::virtual_dom::vcomp::get_event_log(vcomp.id),
                 );
             }),
-            VNode::VList(vlist) => vlist.get(0).expect("VList is not mounted").first_node(),
+            VNode::VList(vlist) => vlist
+                .get(0)
+                .expect("VList is not mounted")
+                .unchecked_first_node(),
             VNode::VRef(node) => node.clone(),
         }
     }
@@ -85,7 +104,7 @@ impl VNode {
                     .expect("VComp has no root vnode")
                     .move_before(parent, next_sibling);
             }
-            _ => super::insert_node(&self.first_node(), parent, next_sibling.as_ref()),
+            _ => super::insert_node(&self.unchecked_first_node(), parent, next_sibling.as_ref()),
         };
     }
 }
diff --git a/packages/yew/src/virtual_dom/vtag.rs b/packages/yew/src/virtual_dom/vtag.rs
index 34ea54bad6a..ad30b4c6d01 100644
--- a/packages/yew/src/virtual_dom/vtag.rs
+++ b/packages/yew/src/virtual_dom/vtag.rs
@@ -521,7 +521,7 @@ impl VDiff for VTag {
                     }
                 } else {
                     let el = self.create_element(parent);
-                    super::insert_node(&el, parent, Some(&ancestor.first_node()));
+                    super::insert_node(&el, parent, ancestor.first_node().as_ref());
                     ancestor.detach(parent);
                     (None, el)
                 }