add some search functions

2 files changed, 129 insertions, 1 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 3077220..8b5c47d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,5 @@
 use core::fmt;
-use tag::Tag;
+pub use tag::{Tag, TagIterator, TagIteratorMut};
 
 //mod query;
 mod tag;
@@ -25,6 +25,68 @@ impl Html {
 		}
 	}
 
+	pub fn child_tags(&self) -> TagIterator {
+		TagIterator {
+			inner: self.nodes.iter(),
+		}
+	}
+
+	pub fn child_tags_mut(&mut self) -> TagIteratorMut {
+		TagIteratorMut {
+			inner: self.nodes.iter_mut(),
+		}
+	}
+
+	pub fn get_by_tag_name_mut(&mut self, looking: &str) -> Option<&mut Tag> {
+		// depth first
+		fn find_node<'a>(tag: &'a mut Tag, looking: &str) -> Option<&'a mut Tag> {
+			if tag.name.as_str() == looking {
+				return Some(tag);
+			}
+
+			for child in tag.child_tags_mut() {
+				if let Some(tag) = find_node(child, looking) {
+					return Some(tag);
+				}
+			}
+
+			None
+		}
+
+		for child in self.child_tags_mut() {
+			if let Some(tag) = find_node(child, looking) {
+				return Some(tag);
+			}
+		}
+
+		None
+	}
+
+	pub fn get_parent_that_contains_tag_name_mut(&mut self, looking: &str) -> Option<&mut Tag> {
+		// depth first
+		fn find_node<'a>(tag: &'a mut Tag, looking: &str) -> Option<&'a mut Tag> {
+			if tag.has_tag(looking) {
+				return Some(tag);
+			}
+
+			for child_tag in tag.child_tags_mut() {
+				if let Some(tag) = find_node(child_tag, looking) {
+					return Some(tag);
+				}
+			}
+
+			None
+		}
+
+		for child_tag in self.child_tags_mut() {
+			if let Some(tag) = find_node(child_tag, looking) {
+				return Some(tag);
+			}
+		}
+
+		None
+	}
+
 	fn parse_node(raw: &str) -> Consumed {
 		match Self::is_tag(raw) {
 			Some(_) => {
diff --git a/src/tag.rs b/src/tag.rs
index e325b1f..0a97061 100644
--- a/src/tag.rs
+++ b/src/tag.rs
@@ -108,6 +108,36 @@ impl Tag {
 	pub fn set_inner_text<S: Into<String>>(&mut self, txt: S) {
 		self.children = vec![Node::Text(txt.into())];
 	}
+
+	pub fn child_tags(&self) -> TagIterator {
+		TagIterator {
+			inner: self.children.iter(),
+		}
+	}
+
+	pub fn child_tags_mut(&mut self) -> TagIteratorMut {
+		TagIteratorMut {
+			inner: self.children.iter_mut(),
+		}
+	}
+
+	pub fn has_tag(&self, tag: &str) -> bool {
+		self.child_tags().any(|t| t.name == tag)
+	}
+
+	pub fn by_tag_mut<'a>(&'a mut self, looking: &str) -> Option<&'a mut Tag> {
+		for tag in self.child_tags_mut() {
+			if tag.name == looking {
+				return Some(tag);
+			}
+
+			if let Some(found) = tag.by_tag_mut(looking) {
+				return Some(found);
+			}
+		}
+
+		None
+	}
 }
 
 impl fmt::Display for Tag {
@@ -138,6 +168,42 @@ impl fmt::Display for Tag {
 	}
 }
 
+pub struct TagIterator<'a> {
+	pub(crate) inner: std::slice::Iter<'a, Node>,
+}
+
+impl<'a> Iterator for TagIterator<'a> {
+	type Item = &'a Tag;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		loop {
+			match self.inner.next() {
+				None => break None,
+				Some(Node::Tag(ref tag)) => break Some(tag),
+				Some(_) => continue,
+			}
+		}
+	}
+}
+
+pub struct TagIteratorMut<'a> {
+	pub(crate) inner: std::slice::IterMut<'a, Node>,
+}
+
+impl<'a> Iterator for TagIteratorMut<'a> {
+	type Item = &'a mut Tag;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		loop {
+			match self.inner.next() {
+				None => break None,
+				Some(Node::Tag(ref mut tag)) => break Some(tag),
+				Some(_) => continue,
+			}
+		}
+	}
+}
+
 #[cfg(test)]
 mod test {
 	use crate::Tag;