AsyncAwaitHighlighter

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

-// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
-
-using System.Collections.Generic;
-using System.Threading;
-using Microsoft.CodeAnalysis.CSharp;
-using Microsoft.CodeAnalysis.CSharp.Syntax;
-using Microsoft.CodeAnalysis.Text;
-using Roslyn.Utilities;
-
-namespace Microsoft.CodeAnalysis.Editor.CSharp.KeywordHighlighting.KeywordHighlighters
-{
-    [ExportHighlighter(LanguageNames.CSharp)]
-    internal class AsyncAnonymousFunctionHighlighter : AbstractAsyncHighlighter<AnonymousFunctionExpressionSyntax>
-    {
-        protected override IEnumerable<TextSpan> GetHighlights(AnonymousFunctionExpressionSyntax node, CancellationToken cancellationToken)
-        {
-            if (node.AsyncKeyword.Kind() != SyntaxKind.AsyncKeyword)
-            {
-                return SpecializedCollections.EmptyEnumerable<TextSpan>();
-            }
-
-            var spans = new List<TextSpan>();
-
-            HighlightRelatedKeywords(node, spans);
-
-            return spans;
-        }
-    }
-}

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

@@ -2,6 +2,7 @@
 
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Extensions;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
@@ -10,9 +11,20 @@
 
 namespace Microsoft.CodeAnalysis.Editor.CSharp.KeywordHighlighting.KeywordHighlighters
 {
-    internal abstract class AbstractAsyncHighlighter<TNode> : AbstractKeywordHighlighter<TNode> where TNode : SyntaxNode
+    [ExportHighlighter(LanguageNames.CSharp)]
+    internal class AsyncAwaitHighlighter : AbstractKeywordHighlighter
     {
-        protected void HighlightRelatedKeywords(SyntaxNode node, List<TextSpan> spans)
+        protected override bool HighlightNode(SyntaxNode node)
+            => node.IsReturnableConstruct();
+
+        protected override IEnumerable<TextSpan> GetHighlightsForNode(SyntaxNode node, CancellationToken cancellationToken)
+        {
+            var spans = new List<TextSpan>();
+            HighlightRelatedKeywords(node, spans);
+            return spans;
+        }
+
+        private static void HighlightRelatedKeywords(SyntaxNode node, List<TextSpan> spans)
         {
             // Highlight async keyword
             switch (node)

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

-// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
-
-using System.Collections.Generic;
-using System.Threading;
-using Microsoft.CodeAnalysis.CSharp;
-using Microsoft.CodeAnalysis.CSharp.Syntax;
-using Microsoft.CodeAnalysis.Text;
-using Roslyn.Utilities;
-
-namespace Microsoft.CodeAnalysis.Editor.CSharp.KeywordHighlighting.KeywordHighlighters
-{
-    [ExportHighlighter(LanguageNames.CSharp)]
-    internal class AsyncLocalFunctionHighlighter : AbstractAsyncHighlighter<LocalFunctionStatementSyntax>
-    {
-        protected override IEnumerable<TextSpan> GetHighlights(LocalFunctionStatementSyntax node, CancellationToken cancellationToken)
-        {
-            if (!node.Modifiers.Any(SyntaxKind.AsyncKeyword))
-            {
-                return SpecializedCollections.EmptyEnumerable<TextSpan>();
-            }
-
-            var spans = new List<TextSpan>();
-
-            HighlightRelatedKeywords(node, spans);
-
-            return spans;
-        }
-    }
-}

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

-// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
-
-using System.Collections.Generic;
-using System.Threading;
-using Microsoft.CodeAnalysis.CSharp;
-using Microsoft.CodeAnalysis.CSharp.Syntax;
-using Microsoft.CodeAnalysis.Text;
-using Roslyn.Utilities;
-
-namespace Microsoft.CodeAnalysis.Editor.CSharp.KeywordHighlighting.KeywordHighlighters
-{
-    [ExportHighlighter(LanguageNames.CSharp)]
-    internal class AsyncMethodHighlighter : AbstractAsyncHighlighter<MethodDeclarationSyntax>
-    {
-        protected override IEnumerable<TextSpan> GetHighlights(MethodDeclarationSyntax node, CancellationToken cancellationToken)
-        {
-            if (!node.Modifiers.Any(SyntaxKind.AsyncKeyword))
-            {
-                return SpecializedCollections.EmptyEnumerable<TextSpan>();
-            }
-
-            var spans = new List<TextSpan>();
-
-            HighlightRelatedKeywords(node, spans);
-
-            return spans;
-        }
-    }
-}

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

-// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
-
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading;
-using Microsoft.CodeAnalysis.CSharp.Extensions;
-using Microsoft.CodeAnalysis.CSharp.Syntax;
-using Microsoft.CodeAnalysis.Text;
-using Roslyn.Utilities;
-
-namespace Microsoft.CodeAnalysis.Editor.CSharp.KeywordHighlighting.KeywordHighlighters
-{
-    [ExportHighlighter(LanguageNames.CSharp)]
-    internal class AwaitHighlighter : AbstractAsyncHighlighter<AwaitExpressionSyntax>
-    {
-        protected override IEnumerable<TextSpan> GetHighlights(AwaitExpressionSyntax awaitExpression, CancellationToken cancellationToken)
-        {
-            var parent = awaitExpression
-                             .AncestorsAndSelf()
-                             .FirstOrDefault(n => n.IsReturnableConstruct());
-
-            if (parent == null)
-            {
-                return SpecializedCollections.EmptyEnumerable<TextSpan>();
-            }
-
-            var spans = new List<TextSpan>();
-
-            HighlightRelatedKeywords(parent, spans);
-
-            return spans;
-        }
-    }
-}

src/EditorFeatures/CSharpTest/KeywordHighlighting/AsyncAnonymousFunctionHighlighterTests.cs

@@ -11,7 +11,7 @@ public class AsyncAnonymousFunctionHighlighterTests : AbstractCSharpKeywordHighl
     {
         internal override IHighlighter CreateHighlighter()
         {
-            return new AsyncAnonymousFunctionHighlighter();
+            return new AsyncAwaitHighlighter();
         }
 
         [Fact, Trait(Traits.Feature, Traits.Features.KeywordHighlighting)]

src/EditorFeatures/CSharpTest/KeywordHighlighting/AsyncLocalFunctionHighlighterTests.cs

@@ -11,7 +11,7 @@ public class AsyncLocalFunctionHighlighterTests : AbstractCSharpKeywordHighlight
     {
         internal override IHighlighter CreateHighlighter()
         {
-            return new AsyncLocalFunctionHighlighter();
+            return new AsyncAwaitHighlighter();
         }
 
         [Fact, Trait(Traits.Feature, Traits.Features.KeywordHighlighting)]

src/EditorFeatures/CSharpTest/KeywordHighlighting/AsyncMethodHighlighterTests.cs

@@ -11,7 +11,7 @@ public class AsyncMethodHighlighterTests : AbstractCSharpKeywordHighlighterTests
     {
         internal override IHighlighter CreateHighlighter()
         {
-            return new AsyncMethodHighlighter();
+            return new AsyncAwaitHighlighter();
         }
 
         [Fact, Trait(Traits.Feature, Traits.Features.KeywordHighlighting)]

src/EditorFeatures/CSharpTest/KeywordHighlighting/AwaitHighlighterTests.cs

@@ -12,7 +12,7 @@ public class AwaitHighlighterTests : AbstractCSharpKeywordHighlighterTests
     {
         internal override IHighlighter CreateHighlighter()
         {
-            return new AwaitHighlighter();
+            return new AsyncAwaitHighlighter();
         }
 
         [Fact, Trait(Traits.Feature, Traits.Features.KeywordHighlighting)]

src/EditorFeatures/Core/Implementation/KeywordHighlighting/AbstractKeywordHighlighter.cs

@@ -8,9 +8,20 @@
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.Highlighting
 {
-    internal abstract class AbstractKeywordHighlighter<TNode> : IHighlighter
-        where TNode : SyntaxNode
+    internal abstract class AbstractKeywordHighlighter<TNode> : AbstractKeywordHighlighter where TNode : SyntaxNode
     {
+        protected sealed override bool HighlightNode(SyntaxNode node) => node is TNode;
+
+        protected sealed override IEnumerable<TextSpan> GetHighlightsForNode(SyntaxNode node, CancellationToken cancellationToken)
+            => GetHighlights((TNode)node, cancellationToken);
+
+        protected abstract IEnumerable<TextSpan> GetHighlights(TNode node, CancellationToken cancellationToken);
+    }
+
+    internal abstract class AbstractKeywordHighlighter : IHighlighter
+    {
+        protected abstract bool HighlightNode(SyntaxNode node);
+
         public IEnumerable<TextSpan> GetHighlights(
             SyntaxNode root, int position, CancellationToken cancellationToken)
         {
@@ -18,9 +29,9 @@ internal abstract class AbstractKeywordHighlighter<TNode> : IHighlighter
             {
                 for (var parent = token.Parent; parent != null; parent = parent.Parent)
                 {
-                    if (parent is TNode parentTNode)
+                    if (HighlightNode(parent))
                     {
-                        var highlights = GetHighlights(parentTNode, cancellationToken);
+                        var highlights = GetHighlightsForNode(parent, cancellationToken);
 
                         // Only return them if any of them matched
                         if (highlights.Any(span => span.IntersectsWith(position)))
@@ -35,7 +46,7 @@ internal abstract class AbstractKeywordHighlighter<TNode> : IHighlighter
             return SpecializedCollections.EmptyEnumerable<TextSpan>();
         }
 
-        protected abstract IEnumerable<TextSpan> GetHighlights(TNode node, CancellationToken cancellationToken);
+        protected abstract IEnumerable<TextSpan> GetHighlightsForNode(SyntaxNode node, CancellationToken cancellationToken);
 
         protected TextSpan EmptySpan(int position)
         {