From 994540d97f1e2ec636bf59bf5887fcd37a0dc771 Mon Sep 17 00:00:00 2001
From: Matt Constable <mattconstable@fastmail.com>
Date: Thu, 14 Jan 2021 17:37:22 +0000
Subject: [PATCH] Add cycle detection to tree rendering

---
 README.jp.md                                  |  1 +
 src/Spectre.Console.Tests/Unit/TreeTests.cs   | 27 ++++++++++++++++++-
 .../Widgets/CircularTreeException.cs          | 15 +++++++++++
 src/Spectre.Console/Widgets/Tree.cs           |  8 +++++-
 4 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 src/Spectre.Console/Widgets/CircularTreeException.cs

diff --git a/README.jp.md b/README.jp.md
index 5047304..598119e 100644
--- a/README.jp.md
+++ b/README.jp.md
@@ -111,6 +111,7 @@ Spectre.Consoleでできることを見るために、
 │ Panels     │ examples/Panels/Panels.csproj         │ Demonstrates how to render items in panels.          │
 │ Rules      │ examples/Rules/Rules.csproj           │ Demonstrates how to render horizontal rules (lines). │
 │ Tables     │ examples/Tables/Tables.csproj         │ Demonstrates how to render tables in a console.      │
+│ Trees      │ examples/Trees/Trees.csproj           │ Demonstrates how to render trees in a console.       │
 ╰────────────┴───────────────────────────────────────┴──────────────────────────────────────────────────────╯
 ```
 
diff --git a/src/Spectre.Console.Tests/Unit/TreeTests.cs b/src/Spectre.Console.Tests/Unit/TreeTests.cs
index eb5001d..87bf400 100644
--- a/src/Spectre.Console.Tests/Unit/TreeTests.cs
+++ b/src/Spectre.Console.Tests/Unit/TreeTests.cs
@@ -1,5 +1,6 @@
 using System.Linq;
 using System.Threading.Tasks;
+using Shouldly;
 using Spectre.Console.Testing;
 using Spectre.Verify.Extensions;
 using VerifyXunit;
@@ -56,5 +57,29 @@ namespace Spectre.Console.Tests.Unit
             // Then
             return Verifier.Verify(console.Output);
         }
+
+        [Fact]
+        public void Should_Throw_If_Tree_Contains_Cycles()
+        {
+            // Given
+            var console = new FakeConsole(width: 80);
+
+            var child2 = new TreeNode(new Text("child 2"));
+            var child3 = new TreeNode(new Text("child 3"));
+            var child1 = new TreeNode(new Text("child 1"));
+            child1.AddNodes(child2, child3);
+            var root = new TreeNode(new Text("Branch Node"));
+            root.AddNodes(child1);
+            child2.AddNode(root);
+
+            var tree = new Tree("root node");
+            tree.AddNodes(root);
+
+            // When
+            var result = Record.Exception(() => console.Render(tree));
+
+            // Then
+            result.ShouldBeOfType<CircularTreeException>();
+        }
     }
-}
\ No newline at end of file
+}
diff --git a/src/Spectre.Console/Widgets/CircularTreeException.cs b/src/Spectre.Console/Widgets/CircularTreeException.cs
new file mode 100644
index 0000000..872f2b2
--- /dev/null
+++ b/src/Spectre.Console/Widgets/CircularTreeException.cs
@@ -0,0 +1,15 @@
+using System;
+
+namespace Spectre.Console
+{
+    /// <summary>
+    /// Indicates that the tree being rendered includes a cycle, and cannot be rendered.
+    /// </summary>
+    public sealed class CircularTreeException : Exception
+    {
+        internal CircularTreeException(string message)
+            : base(message)
+        {
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/Spectre.Console/Widgets/Tree.cs b/src/Spectre.Console/Widgets/Tree.cs
index 8ccacd3..93dfdff 100644
--- a/src/Spectre.Console/Widgets/Tree.cs
+++ b/src/Spectre.Console/Widgets/Tree.cs
@@ -6,7 +6,8 @@ using Spectre.Console.Rendering;
 namespace Spectre.Console
 {
     /// <summary>
-    /// A renderable tree.
+    /// Representation of non-circular tree data.
+    /// Each node added to the tree may only be present in it a single time, in order to facilitate cycle detection.
     /// </summary>
     public sealed class Tree : Renderable, IHasTreeNodes
     {
@@ -54,6 +55,7 @@ namespace Spectre.Console
         protected override IEnumerable<Segment> Render(RenderContext context, int maxWidth)
         {
             var result = new List<Segment>();
+            var visitedNodes = new HashSet<TreeNode>();
 
             var stack = new Stack<Queue<TreeNode>>();
             stack.Push(new Queue<TreeNode>(new[] { _root }));
@@ -77,6 +79,10 @@ namespace Spectre.Console
 
                 var isLastChild = stackNode.Count == 1;
                 var current = stackNode.Dequeue();
+                if (!visitedNodes.Add(current))
+                {
+                    throw new CircularTreeException("Cycle detected in tree - unable to render.");
+                }
 
                 stack.Push(stackNode);