From afcdab91016b8d57ed6a5dfa319a875a57b0349c Mon Sep 17 00:00:00 2001
From: Steven Kirk <grokys@gmail.com>
Date: Thu, 11 Feb 2016 00:03:04 +0100
Subject: [PATCH] Improved GetTemplatedChildren

Now doesn't need to enumerate the whole descendent tree.
---
 .../Templates/TemplateExtensions.cs           | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/Perspex.Controls/Templates/TemplateExtensions.cs b/src/Perspex.Controls/Templates/TemplateExtensions.cs
index dbd5003b2d..018ec033fd 100644
--- a/src/Perspex.Controls/Templates/TemplateExtensions.cs
+++ b/src/Perspex.Controls/Templates/TemplateExtensions.cs
@@ -14,19 +14,28 @@ namespace Perspex.Controls.Templates
     {
         public static IEnumerable<IControl> GetTemplateChildren(this ITemplatedControl control)
         {
-            var visual = control as IVisual;
-
-            if (visual != null)
+            foreach (IControl child in GetTemplateChildren((IControl)control, control))
             {
-                // TODO: This searches the whole descendent tree - it can stop when it exits the
-                // template.
-                return visual.GetVisualDescendents()
-                    .OfType<IControl>()
-                    .Where(x => x.TemplatedParent == control);
+                yield return child;
             }
-            else
+        }
+
+        private static IEnumerable<IControl> GetTemplateChildren(IControl control, ITemplatedControl templatedParent)
+        {
+            foreach (IControl child in control.GetVisualChildren())
             {
-                return Enumerable.Empty<IControl>();
+                if (child.TemplatedParent == templatedParent)
+                {
+                    yield return child;
+                }
+
+                if (child.TemplatedParent != null)
+                {
+                    foreach (var descendent in GetTemplateChildren(child, templatedParent))
+                    {
+                        yield return descendent;
+                    }
+                }
             }
         }
     }