Refactor node refs

examples/counter.rs

@@ -32,11 +32,11 @@ fn main() {
                 .flex_direction(FlexDirection::Row)
                 .child(button(&mut tree, "More!", move |tree| {
                     inc_count.fetch_add(1, Ordering::SeqCst);
-                    tree.set_text(text, inc_count.load(Ordering::SeqCst).to_string())
+                    tree.node(text).set_text( inc_count.load(Ordering::SeqCst).to_string())
                 }))
                 .child(button(&mut tree, "Less!", move |tree| {
                     dec_count.fetch_sub(1, Ordering::SeqCst);
-                    tree.set_text(text, dec_count.load(Ordering::SeqCst).to_string())
+                    tree.node(text).set_text( dec_count.load(Ordering::SeqCst).to_string())
                 }))
                 .build(&mut tree),
         )

src/tree/mod.rs

@@ -8,7 +8,7 @@ use accesskit::{NodeClassSet, NodeId, TreeUpdate};
 use kurbo::Point;
 use skia_safe::Canvas;
 use slotmap::SlotMap;
-use std::{borrow::Cow, num::NonZeroU128};
+use std::{borrow::Cow, num::NonZeroU128, ops::{Index, IndexMut}};
 use taffy::{prelude::Size, style_helpers::TaffyMaxContent, Taffy};
 
 mod iter;
@@ -160,14 +160,7 @@ impl Tree {
         NodeRef { key, tree: self }
     }
 
-    pub fn set_text(&mut self, key: NodeKey, content: impl Into<Cow<'static, str>>) {
-        if let NodeData::Text(ref mut dst) = self.nodes.nodes[key].data {
-            *dst = content.into();
-        } else {
-            todo!()
-        }
-    }
-
+
     pub fn layout(&mut self, root: NodeKey) {
         if self.inner.changes.is_empty() {
             return;
@@ -288,6 +281,7 @@ impl Tree {
     }
 }
 
+
 #[derive(Default)]
 pub struct Nodes {
     pub nodes: SlotMap<NodeKey, Node>,
@@ -302,3 +296,18 @@ impl Nodes {
         IterMut::new(self, root)
     }
 }
+
+
+impl Index<NodeKey> for Nodes {
+    type Output = Node;
+
+    fn index(&self, index: NodeKey) -> &Self::Output {
+        &self.nodes[index]
+    }
+}
+
+impl IndexMut<NodeKey> for Nodes {
+    fn index_mut(&mut self, index: NodeKey) -> &mut Self::Output {
+        &mut self.nodes[index]
+    }
+}

src/tree/node_ref.rs

@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use crate::{element::ElementData, node::NodeData, Node, NodeKey, Tree};
 use skia_safe::Color4f;
 use taffy::{prelude::Size, style::Dimension};
@@ -34,6 +36,19 @@ impl<'a> NodeRef<'a> {
         }
     }
 
+    /// Update the text of a text node.
+    /// 
+    /// ## Panics
+    /// This function will panic if the current reference is to an element,
+    /// not to a text node. 
+    pub fn set_text(&mut self, content: impl Into<Cow<'static, str>>) {
+        if let NodeData::Text(ref mut dst) = self.node().data {
+            *dst = content.into();
+        } else {
+            todo!()
+        }
+    }
+
     /// Update the background color.
     pub fn set_background_color(&mut self, color: Color4f) {
         self.as_mut().background_color = Some(color);