- Make prefixing more generic

1 files changed, 48 insertions, 4 deletions

diff --git a/ext/spl/examples/recursivetreeiterator.inc b/ext/spl/examples/recursivetreeiterator.inc
index 169b3a8022..ab526e4ca1 100755
--- a/ext/spl/examples/recursivetreeiterator.inc
+++ b/ext/spl/examples/recursivetreeiterator.inc
@@ -26,17 +26,61 @@ class RecursiveTreeIterator extends RecursiveIteratorIterator
 		$this->callToString = (bool)($cit_flags & CachingIterator::CALL_TOSTRING);
 	}
 
+	/** @return prefix used if $level < depth and hasNext($level) == true
+	 */
+	function getPrefix1()
+	{
+		return '| ';
+	}
+
+	/** @return prefix used if $level < depth and hasNext($level) == false
+	 */
+	function getPrefix2()
+	{
+		return '  ';
+	}
+
+	/** @return prefix used if $level == depth and hasNext($level) == true
+	 */
+	function getPrefix3()
+	{
+		return '|-';
+	}
+
+	/** @return prefix used if $level == depth and hasNext($level) == false
+	 */
+	function getPrefix4()
+	{
+		return '\-';
+	}
+
+	function getPrefix($level)
+	{
+		if ($level < 0 || $level > $this->getDepth())
+		{
+			throw new OutOfBoundsException('Parameter $level must be >= 0 and <= depth');
+		}
+		if ($level < $this->getDepth())
+		{
+			return $this->getSubIterator($level)->hasNext() ? $this->getPrefix1() : $this->getPrefix2();
+		}
+		else
+		{
+			return $this->getSubIterator($level)->hasNext() ? $this->getPrefix3() : $this->getPrefix4();
+		}
+	}
+
 	/** @return the current element prefixed with ASCII graphics
 	 */
 	function current()
 	{
 		$tree = '';
-		for ($l=0; $l < $this->getDepth(); $l++) {
-			$tree .= $this->getSubIterator($l)->hasNext() ? '| ' : '  ';
+		for ($l=0; $l <= $this->getDepth(); $l++)
+		{
+			$tree .= $this->getPrefix($l);
 		}
 		
-		return $tree . ($this->getSubIterator($l)->hasNext() ? '|-' : '\-')
-		       . ($this->callToString ? $this->getSubIterator($l)->__toString() : $this->getSubIterator($l)->current());
+		return $tree . ($this->callToString ? $this->__toString() : parent::current());
 	}
 
 	/** Aggregates the inner iterator