function to walk over all elements (recursive)

Author: alexkazik

src/dom/element.rs

@@ -7,6 +7,7 @@ use ownable::{IntoOwned, ToBorrowed, ToOwned};
 use serde::Serialize;
 use std::borrow::Cow;
 use std::default::Default;
+use std::ops::ControlFlow;
 
 /// Normal: `<div></div>` or Void: `<meta/>`and `<meta>`
 #[derive(Debug, Clone, Serialize, PartialEq)]
@@ -68,3 +69,66 @@ impl Default for Element<'static> {
         }
     }
 }
+
+pub enum WalkResult {
+    Continue,
+    Skip,
+    Break,
+}
+
+impl<'a> Element<'a> {
+    #[inline]
+    pub fn for_each_element<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&mut Element<'a>) -> WalkResult,
+    {
+        let _ = for_each_element(self.children.as_mut_slice(), &mut f);
+    }
+
+    #[inline]
+    pub async fn for_each_element_async<F>(&mut self, mut f: F)
+    where
+        F: AsyncFnMut(&mut Element<'a>) -> WalkResult,
+    {
+        let _ = for_each_element_async(self.children.as_mut_slice(), &mut f).await;
+    }
+}
+
+pub(crate) fn for_each_element<'a, F>(tree: &mut [Node<'a>], f: &mut F) -> ControlFlow<(), ()>
+where
+    F: FnMut(&mut Element<'a>) -> WalkResult,
+{
+    for n in tree {
+        if let Node::Element(e) = n {
+            match f(e) {
+                WalkResult::Continue => {
+                    for_each_element(e.children.as_mut_slice(), f)?;
+                }
+                WalkResult::Skip => (),
+                WalkResult::Break => return ControlFlow::Break(()),
+            }
+        }
+    }
+    ControlFlow::Continue(())
+}
+
+pub(crate) async fn for_each_element_async<'a, F>(
+    tree: &mut [Node<'a>],
+    f: &mut F,
+) -> ControlFlow<(), ()>
+where
+    F: AsyncFnMut(&mut Element<'a>) -> WalkResult,
+{
+    for n in tree {
+        if let Node::Element(e) = n {
+            match f(e).await {
+                WalkResult::Continue => {
+                    Box::pin(for_each_element_async(e.children.as_mut_slice(), f)).await?;
+                }
+                WalkResult::Skip => (),
+                WalkResult::Break => return ControlFlow::Break(()),
+            }
+        }
+    }
+    ControlFlow::Continue(())
+}

src/dom/mod.rs

@@ -18,6 +18,7 @@ pub mod span;
 pub mod vecmap;
 pub mod vecset;
 
+use crate::dom::element::{WalkResult, for_each_element, for_each_element_async};
 #[cfg(feature = "source-span")]
 use crate::dom::span::SourceSpan;
 use element::{Element, ElementVariant};
@@ -368,4 +369,20 @@ impl<'a> Dom<'a> {
         }
         Ok(attribute)
     }
+
+    #[inline]
+    pub fn for_each_element<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&mut Element) -> WalkResult,
+    {
+        let _ = for_each_element(self.children.as_mut_slice(), &mut f);
+    }
+
+    #[inline]
+    pub async fn for_each_element_async<F>(&mut self, mut f: F)
+    where
+        F: AsyncFnMut(&mut Element) -> WalkResult,
+    {
+        let _ = for_each_element_async(self.children.as_mut_slice(), &mut f).await;
+    }
 }

src/dom/node.rs

@@ -1,5 +1,6 @@
-use super::element::Element;
-use crate::Text;
+use super::element::{Element, for_each_element_async};
+use crate::dom::element::for_each_element;
+use crate::{Text, WalkResult};
 use ownable::{IntoOwned, ToBorrowed, ToOwned};
 use serde::Serialize;
 
@@ -46,6 +47,26 @@ impl<'a> Node<'a> {
             _ => None,
         }
     }
+
+    #[inline]
+    pub fn for_each_element<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&mut Element<'a>) -> WalkResult,
+    {
+        if let Node::Element(e) = self {
+            let _ = for_each_element(e.children.as_mut_slice(), &mut f);
+        }
+    }
+
+    #[inline]
+    pub async fn for_each_element_async<F>(&mut self, mut f: F)
+    where
+        F: AsyncFnMut(&mut Element<'a>) -> WalkResult,
+    {
+        if let Node::Element(e) = self {
+            let _ = for_each_element_async(e.children.as_mut_slice(), &mut f).await;
+        }
+    }
 }
 
 impl<'a> IntoIterator for &'a Node<'a> {

src/lib.rs

@@ -104,6 +104,7 @@ use grammar::Rule;
 
 pub use crate::dom::Dom;
 pub use crate::dom::DomVariant;
+pub use crate::dom::element::WalkResult;
 pub use crate::dom::element::{Element, ElementVariant};
 pub use crate::dom::html::{Attribute, Text};
 pub use crate::dom::node::Node;