Re-import iterators for Tree

matthunz · Oct 10, 2023 · 53c4376 · 53c4376
1 parent 6dfb916
commit 53c4376
Show file tree

Hide file tree

Showing 5 changed files with 199 additions and 19 deletions.
diff --git a/examples/hello.rs b/examples/hello.rs
@@ -1,6 +1,6 @@
 use viewbuilder::prelude::*;
 
-fn app(cx: Scope) -> Element  {
+fn app(cx: Scope) -> Element {
     cx.render(rsx! {
         view {
             background_color: 123,
@@ -9,6 +9,6 @@ fn app(cx: Scope) -> Element  {
     })
 }
 
-fn main()  {
+fn main() {
     viewbuilder::run(app)
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -3,7 +3,7 @@ use dioxus::prelude::{Component, VirtualDom};
 pub mod element;
 
 mod tree;
-use crate::tree::Tree;
+pub use crate::tree::Tree;
 
 pub mod prelude {
     pub use super::dioxus_elements;
@@ -18,7 +18,8 @@ pub mod dioxus_elements {
         pub const TAG_NAME: &'static str = "view";
         pub const NAME_SPACE: Option<&'static str> = None;
 
-        pub const background_color: (&'static str, Option<&'static str>, bool) = ("background_color", None, false);
+        pub const background_color: (&'static str, Option<&'static str>, bool) =
+            ("background_color", None, false);
     }
 }
 
@@ -29,5 +30,5 @@ pub fn run(app: Component) {
 
     let mut tree = Tree::new();
     tree.update(mutations);
-    dbg!(tree);
+    println!("{}", tree);
 }
diff --git a/src/tree/iter.rs b/src/tree/iter.rs
@@ -0,0 +1,66 @@
+use super::{Node, NodeKind, Tree};
+use dioxus::core::ElementId;
+
+enum Operation {
+    Key(ElementId),
+    Pop,
+}
+
+pub enum Item<'a> {
+    Node {
+        key: ElementId,
+        node: &'a Node,
+        level: usize,
+    },
+    Pop {
+        level: usize,
+    },
+}
+
+pub struct Iter<'a> {
+    tree: &'a Tree,
+    stack: Vec<Operation>,
+    count: usize,
+}
+
+impl<'a> Iter<'a> {
+    pub(crate) fn new(tree: &'a Tree, root: ElementId) -> Self {
+        Iter {
+            tree,
+            stack: vec![Operation::Key(root)],
+            count: 0,
+        }
+    }
+}
+
+impl<'a> Iterator for Iter<'a> {
+    type Item = Item<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.stack.pop().map(|item| match item {
+            Operation::Key(key) => {
+                let node = &self.tree.elements[&key];
+
+                self.stack.push(Operation::Pop);
+                if let NodeKind::Element { ref children, .. } = node.kind {
+                    for child in children.iter().copied().map(Operation::Key) {
+                        self.stack.push(child);
+                    }
+                }
+
+                let count = self.count;
+                self.count += 1;
+
+                Item::Node {
+                    key,
+                    node,
+                    level: count,
+                }
+            }
+            Operation::Pop => {
+                self.count -= 1;
+                Item::Pop { level: self.count }
+            }
+        })
+    }
+}
diff --git a/src/tree/iter_mut.rs b/src/tree/iter_mut.rs
@@ -0,0 +1,60 @@
+use super::{Node, NodeKind, Tree};
+use dioxus::core::ElementId;
+use std::marker::PhantomData;
+
+enum Operation {
+    Key(ElementId),
+    Pop,
+}
+
+pub enum ItemMut<'a> {
+    Node { node: &'a mut Node, level: usize },
+    Pop { level: usize },
+}
+
+pub struct IterMut<'a> {
+    tree: *mut Tree,
+    stack: Vec<Operation>,
+    count: usize,
+    _marker: PhantomData<&'a mut Tree>,
+}
+
+impl<'a> IterMut<'a> {
+    pub(crate) fn new(tree: &'a mut Tree, root: ElementId) -> Self {
+        IterMut {
+            tree,
+            stack: vec![Operation::Key(root)],
+            count: 0,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<'a> Iterator for IterMut<'a> {
+    type Item = ItemMut<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.stack.pop().map(|item| match item {
+            Operation::Key(key) => {
+                let tree = unsafe { &mut *self.tree };
+                let node = tree.elements.get_mut(&key).unwrap();
+
+                if let NodeKind::Element { ref children, .. } = node.kind {
+                    for child in children.iter().copied().map(Operation::Key) {
+                        self.stack.push(child);
+                    }
+                }
+                self.stack.push(Operation::Pop);
+
+                let count = self.count;
+                self.count += 1;
+
+                ItemMut::Node { node, level: count }
+            }
+            Operation::Pop => {
+                self.count -= 1;
+                ItemMut::Pop { level: self.count }
+            }
+        })
+    }
+}
diff --git a/src/tree.rs → src/tree/mod.rs b/src/tree.rs → src/tree/mod.rs
@@ -1,9 +1,16 @@
+use core::fmt;
 use dioxus::{
-    core::{ElementId, Mutation, Mutations, BorrowedAttributeValue},
+    core::{BorrowedAttributeValue, ElementId, Mutation, Mutations},
     prelude::{TemplateAttribute, TemplateNode},
 };
 use std::{any::Any, collections::HashMap, rc::Rc};
 
+mod iter;
+pub use iter::{Item, Iter};
+
+mod iter_mut;
+pub use iter_mut::{ItemMut, IterMut};
+
 #[derive(Debug)]
 struct TemplateAttributeValue {}
 
@@ -34,7 +41,7 @@ enum NodeKind {
 }
 
 #[derive(Debug)]
-struct Node {
+pub struct Node {
     kind: NodeKind,
 }
 
@@ -62,6 +69,14 @@ impl Tree {
         }
     }
 
+    pub fn iter(&self) -> Iter {
+        Iter::new(self, ElementId(0))
+    }
+
+    pub fn iter_mut(&mut self) -> IterMut {
+        IterMut::new(self, ElementId(0))
+    }
+
     pub fn update(&mut self, mutations: Mutations) {
         for template in mutations.templates {
             let roots: Vec<_> = template
@@ -70,20 +85,20 @@ impl Tree {
                 .map(|elem| match elem {
                     TemplateNode::Text { text } => TemplateElement::Text(Rc::new(text.to_string())),
                     TemplateNode::Element {
-                        tag,
-                        namespace,
+                        tag: _,
+                        namespace: _,
                         attrs,
-                        children,
+                        children: _,
                     } => {
                         let attrs = attrs
                             .into_iter()
                             .filter_map(|attr| match attr {
                                 TemplateAttribute::Static {
-                                    name,
-                                    value,
-                                    namespace,
+                                    name: _,
+                                    value: _,
+                                    namespace: _,
                                 } => todo!(),
-                                TemplateAttribute::Dynamic { id } => None,
+                                TemplateAttribute::Dynamic { id: _ } => None,
                             })
                             .collect();
                         TemplateElement::View { attrs }
@@ -114,7 +129,7 @@ impl Tree {
                                 },
                             );
                         }
-                        TemplateElement::View { attrs } => {
+                        TemplateElement::View { attrs: _ } => {
                             self.elements.insert(
                                 id,
                                 Node {
@@ -132,7 +147,7 @@ impl Tree {
                 Mutation::AppendChildren { id, m } => {
                     let elem = self.elements.get_mut(&id).unwrap();
                     if let NodeKind::Element {
-                        ref attrs,
+                        attrs: _,
                         ref mut children,
                     } = elem.kind
                     {
@@ -146,13 +161,13 @@ impl Tree {
                     name,
                     value,
                     id,
-                    ns,
+                    ns: _,
                 } => {
                     let elem = self.elements.get_mut(&id).unwrap();
                     if let NodeKind::Element { ref mut attrs, .. } = elem.kind {
-                        let val = match value {
+                        let _val = match value {
                             BorrowedAttributeValue::Int(i) => Box::new(i),
-                            _ => todo!()
+                            _ => todo!(),
                         };
                         attrs.push(Attribute {
                             name: name.to_string(),
@@ -165,3 +180,41 @@ impl Tree {
         }
     }
 }
+
+impl fmt::Display for Tree {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        for item in self.iter() {
+            match item {
+                Item::Node {
+                    key: _,
+                    node,
+                    level,
+                } => {
+                    for _ in 0..level {
+                        write!(f, "    ")?;
+                    }
+
+                    match &node.kind {
+                        NodeKind::Element { attrs, children: _ } => {
+                            writeln!(f, "Element {{")?;
+                            for attr in attrs {
+                                for _ in 0..level + 1 {
+                                    write!(f, "    ")?;
+                                }
+                                writeln!(f, "{}: {:?}", &attr.name, &attr.value)?;
+                            }
+                        }
+                        NodeKind::Text(text) => writeln!(f, "{}", &text)?,
+                    }
+                }
+                Item::Pop { level } => {
+                    for _ in 0..level {
+                        write!(f, "    ")?;
+                    }
+                    writeln!(f, "}}")?
+                }
+            }
+        }
+        Ok(())
+    }
+}