Traverser: $enter and $leave callbacks

nette · Mar 10, 2022 · f7e5415 · f7e5415
1 parent 997b071
commit f7e5415
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 10 deletions.
diff --git a/src/Neon/Traverser.php b/src/Neon/Traverser.php
@@ -16,17 +16,24 @@ final class Traverser
 	public const DontTraverseChildren = 1;
 	public const StopTraversal = 2;
 
-	/** @var callable(Node): (Node|int|null) */
-	private $callback;
+	/** @var callable(Node): (Node|int|null)|null */
+	private $enter;
+
+	/** @var callable(Node): (Node|int|null)|null */
+	private $leave;
 
 	/** @var bool */
 	private $stop;
 
 
-	/** @param  callable(Node): (Node|int|null)  $callback */
-	public function traverse(Node $node, callable $callback): Node
+	/**
+	 * @param  callable(Node): (Node|int|null)|null  $enter
+	 * @param  callable(Node): (Node|int|null)|null  $leave
+	 */
+	public function traverse(Node $node, ?callable $enter = null, ?callable $leave = null): Node
 	{
-		$this->callback = $callback;
+		$this->enter = $enter;
+		$this->leave = $leave;
 		$this->stop = false;
 		return $this->traverseNode($node);
 	}
@@ -35,8 +42,8 @@ public function traverse(Node $node, callable $callback): Node
 	private function traverseNode(Node $node): Node
 	{
 		$children = true;
-		if ($this->callback) {
-			$res = ($this->callback)($node);
+		if ($this->enter) {
+			$res = ($this->enter)($node);
 			if ($res instanceof Node) {
 				$node = $res;
 
@@ -58,6 +65,15 @@ private function traverseNode(Node $node): Node
 			}
 		}
 
+		if (!$this->stop && $this->leave) {
+			$res = ($this->leave)($node);
+			if ($res instanceof Node) {
+				$node = $res;
+			} elseif ($res === self::StopTraversal) {
+				$this->stop = true;
+			}
+		}
+
 		return $node;
 	}
 }
diff --git a/tests/Neon/Traverser.order.phpt b/tests/Neon/Traverser.order.phpt
@@ -19,14 +19,19 @@ a: 1
 $log = [];
 $traverser->traverse(
 	$node,
-	function ($node) use (&$log) { $log[] = ['enter', get_class($node)]; }
+	function ($node) use (&$log) { $log[] = ['enter', get_class($node)]; },
+	function ($node) use (&$log) { $log[] = ['leave', get_class($node)]; }
 );
 
 Assert::equal([
 	['enter', Node\BlockArrayNode::class],
 	['enter', Node\ArrayItemNode::class],
 	['enter', Node\LiteralNode::class],
+	['leave', Node\LiteralNode::class],
 	['enter', Node\LiteralNode::class],
+	['leave', Node\LiteralNode::class],
+	['leave', Node\ArrayItemNode::class],
+	['leave', Node\BlockArrayNode::class],
 ], $log);
 
 
@@ -39,12 +44,15 @@ $traverser->traverse(
 		return $node instanceof Node\ArrayItemNode
 			? Neon\Traverser::DontTraverseChildren
 			: null;
-	}
+	},
+	function ($node) use (&$log) { $log[] = ['leave', get_class($node)]; }
 );
 
 Assert::equal([
 	['enter', Node\BlockArrayNode::class],
 	['enter', Node\ArrayItemNode::class],
+	['leave', Node\ArrayItemNode::class],
+	['leave', Node\BlockArrayNode::class],
 ], $log);
 
 
@@ -55,10 +63,29 @@ $traverser->traverse(
 	function ($node) use (&$log) {
 		$log[] = ['enter', get_class($node)];
 		return $node instanceof Node\ArrayItemNode ? Neon\Traverser::StopTraversal : null;
-	}
+	},
+	function ($node) use (&$log) { $log[] = ['enter', get_class($node)]; }
 );
 
 Assert::equal([
 	['enter', Node\BlockArrayNode::class],
 	['enter', Node\ArrayItemNode::class],
 ], $log);
+
+
+
+$log = [];
+$traverser->traverse(
+	$node,
+	null,
+	function ($node) use (&$log) {
+		$log[] = ['leave', get_class($node)];
+		return $node instanceof Node\ArrayItemNode ? Neon\Traverser::StopTraversal : null;
+	}
+);
+
+Assert::equal([
+	['leave', Node\LiteralNode::class],
+	['leave', Node\LiteralNode::class],
+	['leave', Node\ArrayItemNode::class],
+], $log);