From 1143fcf0f6489d2766f959faadc66a2123ffc055 Mon Sep 17 00:00:00 2001
From: Steven Kirk <steven.kirk@tricycleinc.com>
Date: Wed, 1 Oct 2014 16:10:39 +0200
Subject: [PATCH] Fix tree item container generator.

---
 .../Generators/ItemContainerGenerator.cs      | 34 +++++++++++++---
 .../Generators/TreeItemContainerGenerator.cs  | 39 ++++++++++++++++++-
 2 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/Perspex.Controls/Generators/ItemContainerGenerator.cs b/Perspex.Controls/Generators/ItemContainerGenerator.cs
index 129b52d3ba..ff0e31894b 100644
--- a/Perspex.Controls/Generators/ItemContainerGenerator.cs
+++ b/Perspex.Controls/Generators/ItemContainerGenerator.cs
@@ -84,9 +84,10 @@ namespace Perspex.Controls.Generators
                 {
                     Control container = this.CreateContainerOverride(item);
                     container.TemplatedParent = null;
-                    this.containersByItem.Add(item, container);
-                    this.itemsByContainer.Add(container, item);
+                    this.AddInternal(item, container);
                     result.Add(container);
+
+                    System.Diagnostics.Debug.WriteLine("{0} : Generated container for {1} {2}: {3}", this.GetHashCode(), item, item.GetHashCode(), container);
                 }
             }
             finally
@@ -103,10 +104,7 @@ namespace Perspex.Controls.Generators
 
             foreach (var item in items)
             {
-                Control container = this.containersByItem[item];
-                this.containersByItem.Remove(item);
-                this.itemsByContainer.Remove(container);
-                result.Add(container);
+                result.Add(this.RemoveByItemInternal(item));
             }
 
             return result;
@@ -122,5 +120,29 @@ namespace Perspex.Controls.Generators
         {
             return this.Owner.ApplyDataTemplate(item);
         }
+
+        protected void AddInternal(object item, Control container)
+        {
+            this.containersByItem.Add(item, container);
+            this.itemsByContainer.Add(container, item);
+        }
+
+        protected object RemoveByContainerInternal(Control container)
+        {
+            object item = this.itemsByContainer[container];
+            this.containersByItem.Remove(item);
+            this.itemsByContainer.Remove(container);
+            System.Diagnostics.Debug.WriteLine("{0} : Removed container for {1} {2}", this.GetHashCode(), item, item.GetHashCode());
+            return item;
+        }
+
+        protected Control RemoveByItemInternal(object item)
+        {
+            Control container = this.containersByItem[item];
+            this.containersByItem.Remove(item);
+            this.itemsByContainer.Remove(container);
+            System.Diagnostics.Debug.WriteLine("{0} : Removed container for {1} {2}", this.GetHashCode(), item, item.GetHashCode());
+            return container;
+        }
     }
 }
diff --git a/Perspex.Controls/Generators/TreeItemContainerGenerator.cs b/Perspex.Controls/Generators/TreeItemContainerGenerator.cs
index fe068e95d5..721a994d50 100644
--- a/Perspex.Controls/Generators/TreeItemContainerGenerator.cs
+++ b/Perspex.Controls/Generators/TreeItemContainerGenerator.cs
@@ -6,13 +6,31 @@
 
 namespace Perspex.Controls.Generators
 {
-    public class TreeItemContainerGenerator<T> : ItemContainerGenerator where T : TreeViewItem, new()
+    using System;
+    using System.Collections;
+    using System.Collections.Generic;
+    using System.Linq;
+
+    public class TreeItemContainerGenerator<T> : ItemContainerGenerator, IItemContainerGenerator where T : TreeViewItem, new()
     {
         public TreeItemContainerGenerator(ItemsControl owner)
             : base(owner)
         {
         }
 
+        IEnumerable<Control> IItemContainerGenerator.Remove(IEnumerable items)
+        {
+            var result = new List<Control>();
+
+            foreach (var item in items)
+            {
+                var container = (T)this.GetContainerForItem(item);
+                this.Remove(container, result);
+            }
+
+            return result;
+        }
+
         protected override Control CreateContainerOverride(object item)
         {
             T result = item as T;
@@ -52,5 +70,24 @@ namespace Perspex.Controls.Generators
 
             return treeTemplate;
         }
+
+        private void Remove(T container, List<Control> removed)
+        {
+            if (container.Items != null)
+            {
+                foreach (var childItem in container.Items)
+                {
+                    var childContainer = (T)this.GetContainerForItem(childItem);
+
+                    if (childContainer != null)
+                    {
+                        this.Remove(childContainer, removed);
+                    }
+                }
+            }
+
+            this.RemoveByContainerInternal(container);
+            removed.Add(container);
+        }
     }
 }