webpage/crates/service-webpage/src/ast/walk.rs

use markdown::mdast::Node;
use std::fmt::Debug;
use std::marker::PhantomPinned;

pub enum AstWalkStep<'a, T> {
	Enter(&'a T),
	Exit(&'a T),
}

impl<T: Debug> Debug for AstWalkStep<'_, T> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			Self::Enter(x) => f.debug_tuple("AstWalkStep::Enter").field(x).finish(),
			Self::Exit(x) => f.debug_tuple("AstWalkStep::Exit").field(x).finish(),
		}
	}
}

pub struct AstWalk<'a> {
	_pin: PhantomPinned,
	child_stack: Vec<usize>,
	node_stack: Vec<&'a Node>,
}

impl<'a> AstWalk<'a> {
	pub fn new(root: &'a Node) -> Self {
		let mut res = Self {
			_pin: PhantomPinned {},
			node_stack: Vec::with_capacity(32),
			child_stack: Vec::with_capacity(32),
		};

		res.node_stack.push(root);
		return res;
	}

	fn _next_inner(&mut self) -> Option<AstWalkStep<'a, Node>> {
		if self.node_stack.is_empty() {
			return None;
		}

		let current_node = *self.node_stack.last().unwrap();

		// The index of the next child we should look at.
		// If `None`, we look at the parent.
		let current_child = {
			let n_nodes = self.node_stack.len();
			let n_childs = self.child_stack.len();
			match n_nodes - n_childs {
				2.. => unreachable!(),
				1 => None,
				0 => Some(self.child_stack.pop().unwrap()),
			}
		};

		match current_child {
			None => {
				self.child_stack.push(0);
				return Some(AstWalkStep::Enter(current_node));
			}

			Some(current_child) => {
				let child = current_node
					.children()
					.map(|x| x.get(current_child))
					.flatten();

				match child {
					None => {
						self.node_stack.pop();
						return Some(AstWalkStep::Exit(current_node));
					}

					Some(x) => {
						self.child_stack.push(current_child + 1);
						self.node_stack.push(&x);
						self.child_stack.push(0);
						return Some(AstWalkStep::Enter(x));
					}
				}
			}
		}
	}
}

impl<'a> Iterator for AstWalk<'a> {
	type Item = AstWalkStep<'a, Node>;

	fn next(self: &mut Self) -> Option<Self::Item> {
		return self._next_inner();
	}
}

//
// MARK: tests
//

#[cfg(test)]
mod tests {
	use super::*;
	use markdown::mdast::{Emphasis, Paragraph, Root, Strong, Text};

	#[test]
	fn single_leaf() {
		let node = Node::Text(Text {
			value: "Hello".to_string(),
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		assert_eq!(steps.len(), 2);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Exit(_)));
	}

	#[test]
	fn single_child() {
		let node = Node::Paragraph(Paragraph {
			children: vec![Node::Text(Text {
				value: "Hello".to_string(),
				position: None,
			})],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Should be: Enter(Paragraph), Leaf(Text), Exit(Paragraph)
		assert_eq!(steps.len(), 4);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Enter(_)));
		assert!(matches!(steps[2], AstWalkStep::Exit(_)));
		assert!(matches!(steps[3], AstWalkStep::Exit(_)));
	}

	#[test]
	fn multiple_children() {
		let node = Node::Paragraph(Paragraph {
			children: vec![
				Node::Text(Text {
					value: "Hello".to_string(),
					position: None,
				}),
				Node::Text(Text {
					value: " ".to_string(),
					position: None,
				}),
				Node::Text(Text {
					value: "World".to_string(),
					position: None,
				}),
			],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Should be: Enter(Paragraph), Leaf(Text1), Leaf(Text2), Leaf(Text3), Exit(Paragraph)
		assert_eq!(steps.len(), 8);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Enter(_)));
		assert!(matches!(steps[2], AstWalkStep::Exit(_)));
		assert!(matches!(steps[3], AstWalkStep::Enter(_)));
		assert!(matches!(steps[4], AstWalkStep::Exit(_)));
		assert!(matches!(steps[5], AstWalkStep::Enter(_)));
		assert!(matches!(steps[6], AstWalkStep::Exit(_)));
		assert!(matches!(steps[7], AstWalkStep::Exit(_)));
	}

	#[test]
	fn nested_1() {
		let node = Node::Paragraph(Paragraph {
			children: vec![Node::Emphasis(Emphasis {
				children: vec![Node::Text(Text {
					value: "emphasized".to_string(),
					position: None,
				})],
				position: None,
			})],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Should be: Enter(Paragraph), Enter(Emphasis), Leaf(Text), Exit(Emphasis), Exit(Paragraph)
		assert_eq!(steps.len(), 6);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Enter(_)));
		assert!(matches!(steps[2], AstWalkStep::Enter(_)));
		assert!(matches!(steps[3], AstWalkStep::Exit(_)));
		assert!(matches!(steps[4], AstWalkStep::Exit(_)));
		assert!(matches!(steps[5], AstWalkStep::Exit(_)));
	}

	#[test]
	fn nested_2() {
		// Create: Paragraph -> [Text, Strong -> Text, Text]
		let node = Node::Paragraph(Paragraph {
			children: vec![
				Node::Text(Text {
					value: "Before ".to_string(),
					position: None,
				}),
				Node::Strong(Strong {
					children: vec![Node::Text(Text {
						value: "bold".to_string(),
						position: None,
					})],
					position: None,
				}),
				Node::Text(Text {
					value: " after".to_string(),
					position: None,
				}),
			],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Expected order:
		// 0: Enter(Paragraph)
		// 1: Leaf(Text "Before ")
		// 2: Enter(Strong)
		// 3: Leaf(Text "bold")
		// 4: Exit(Strong)
		// 5: Leaf(Text " after")
		// 6: Exit(Paragraph)
		assert_eq!(steps.len(), 10);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Enter(_)));
		assert!(matches!(steps[2], AstWalkStep::Exit(_)));
		assert!(matches!(steps[3], AstWalkStep::Enter(_)));
		assert!(matches!(steps[4], AstWalkStep::Enter(_)));
		assert!(matches!(steps[5], AstWalkStep::Exit(_)));
		assert!(matches!(steps[6], AstWalkStep::Exit(_)));
		assert!(matches!(steps[7], AstWalkStep::Enter(_)));
		assert!(matches!(steps[8], AstWalkStep::Exit(_)));
		assert!(matches!(steps[9], AstWalkStep::Exit(_)));
	}

	#[test]
	fn nested_3() {
		let node = Node::Paragraph(Paragraph {
			children: vec![Node::Emphasis(Emphasis {
				children: vec![Node::Strong(Strong {
					children: vec![Node::Text(Text {
						value: "deeply nested".to_string(),
						position: None,
					})],
					position: None,
				})],
				position: None,
			})],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Should be: Enter(Paragraph), Enter(Emphasis), Enter(Strong), Leaf(Text), Exit(Strong), Exit(Emphasis), Exit(Paragraph)
		assert_eq!(steps.len(), 8);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Enter(_)));
		assert!(matches!(steps[2], AstWalkStep::Enter(_)));
		assert!(matches!(steps[3], AstWalkStep::Enter(_)));
		assert!(matches!(steps[4], AstWalkStep::Exit(_)));
		assert!(matches!(steps[5], AstWalkStep::Exit(_)));
		assert!(matches!(steps[6], AstWalkStep::Exit(_)));
		assert!(matches!(steps[7], AstWalkStep::Exit(_)));
	}

	#[test]
	fn empty_parent() {
		let node = Node::Paragraph(Paragraph {
			children: vec![],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Should be: Enter(Paragraph), Exit(Paragraph)
		assert_eq!(steps.len(), 2);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Exit(_)));
	}

	#[test]
	fn multiple_paragraphs() {
		let node = Node::Root(Root {
			children: vec![
				Node::Paragraph(Paragraph {
					children: vec![Node::Text(Text {
						value: "First".to_string(),
						position: None,
					})],
					position: None,
				}),
				Node::Paragraph(Paragraph {
					children: vec![Node::Text(Text {
						value: "Second".to_string(),
						position: None,
					})],
					position: None,
				}),
			],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		// Expected order:
		// 0: Enter(Root)
		// 1: Enter(Paragraph)
		// 2: Enter(Text "First")
		// 2: Exit(Text "First")
		// 3: Exit(Paragraph)
		// 4: Enter(Paragraph)
		// 5: Enter(Text "Second")
		// 5: Exit(Text "Second")
		// 6: Exit(Paragraph)
		// 7: Exit(Root)
		assert_eq!(steps.len(), 10);
		assert!(matches!(steps[0], AstWalkStep::Enter(_)));
		assert!(matches!(steps[1], AstWalkStep::Enter(_)));
		assert!(matches!(steps[2], AstWalkStep::Enter(_)));
		assert!(matches!(steps[3], AstWalkStep::Exit(_)));
		assert!(matches!(steps[4], AstWalkStep::Exit(_)));
		assert!(matches!(steps[5], AstWalkStep::Enter(_)));
		assert!(matches!(steps[6], AstWalkStep::Enter(_)));
		assert!(matches!(steps[7], AstWalkStep::Exit(_)));
		assert!(matches!(steps[8], AstWalkStep::Exit(_)));
		assert!(matches!(steps[9], AstWalkStep::Exit(_)));
	}

	#[test]
	fn enter_exit() {
		let node = Node::Root(Root {
			children: vec![Node::Paragraph(Paragraph {
				children: vec![
					Node::Emphasis(Emphasis {
						children: vec![Node::Text(Text {
							value: "a".to_string(),
							position: None,
						})],
						position: None,
					}),
					Node::Strong(Strong {
						children: vec![Node::Text(Text {
							value: "b".to_string(),
							position: None,
						})],
						position: None,
					}),
				],
				position: None,
			})],
			position: None,
		});

		let walker = AstWalk::new(&node);
		let steps: Vec<_> = walker.collect();

		let enter_count = steps
			.iter()
			.filter(|s| matches!(s, AstWalkStep::Enter(_)))
			.count();
		let exit_count = steps
			.iter()
			.filter(|s| matches!(s, AstWalkStep::Exit(_)))
			.count();

		assert_eq!(enter_count, exit_count);
		assert_eq!(enter_count, 6);
	}
}