Use stack

Author: CyrusNajmabadi

src/EditorFeatures/CSharp/Highlighting/KeywordHighlighters/AsyncAwaitHighlighter.cs

@@ -4,19 +4,24 @@
 using System.Collections.Generic;
 using System.ComponentModel.Composition;
 using System.Linq;
+using System.Net.Http.Headers;
 using System.Threading;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Extensions;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Editor.Implementation.Highlighting;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Editor.CSharp.KeywordHighlighting.KeywordHighlighters
 {
     [ExportHighlighter(LanguageNames.CSharp)]
     internal class AsyncAwaitHighlighter : AbstractKeywordHighlighter
     {
+        private static readonly ObjectPool<Stack<SyntaxNode>> s_stackPool
+            = new ObjectPool<Stack<SyntaxNode>>(() => new Stack<SyntaxNode>());
+
         [ImportingConstructor]
         public AsyncAwaitHighlighter()
         {
@@ -28,13 +33,40 @@ protected override bool IsHighlightableNode(SyntaxNode node)
         protected override IEnumerable<TextSpan> GetHighlightsForNode(SyntaxNode node, CancellationToken cancellationToken)
         {
             var spans = new List<TextSpan>();
-            foreach (var child in node.DescendantNodesAndSelf(n => !n.IsReturnableConstruct()))
+
+            foreach (var current in WalkChildren(node))
             {
-                HighlightRelatedKeywords(node, spans);
+                cancellationToken.ThrowIfCancellationRequested();
+                HighlightRelatedKeywords(current, spans);
             }
+
             return spans;
         }
 
+        private IEnumerable<SyntaxNode> WalkChildren(SyntaxNode node)
+        {
+            using (var pooledObject = s_stackPool.GetPooledObject())
+            {
+                var stack = pooledObject.Object;
+                stack.Push(node);
+
+                while (stack.Count > 0)
+                {
+                    var current = stack.Pop();
+                    yield return current;
+
+                    foreach (var child in current.ChildNodes())
+                    {
+                        // Only recurse if we have anything to do
+                        if (!child.IsReturnableConstruct())
+                        {
+                            stack.Push(child);
+                        }
+                    }
+                }
+            }
+        }
+
         private static void HighlightRelatedKeywords(SyntaxNode node, List<TextSpan> spans)
         {
             // Highlight async keyword