Highlight await and return in top-level statements

src/EditorFeatures/CSharpTest/KeywordHighlighting/AwaitHighlighterTests.cs

@@ -323,5 +323,14 @@ class C
     }
 }");
         }
+
+        [WorkItem(60400, "https://github.com/dotnet/roslyn/issues/60400")]
+        [Fact, Trait(Traits.Feature, Traits.Features.KeywordHighlighting)]
+        public async Task TestTopLevelStatements()
+        {
+            await TestAsync(
+@"[|await|] Task.Delay(1000);
+{|Cursor:[|await|]|} Task.Run(() => { })");
+        }
     }
 }

src/EditorFeatures/CSharpTest/KeywordHighlighting/ReturnStatementHighlighterTests.cs

@@ -462,5 +462,13 @@ void M()
     }
 }");
         }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.KeywordHighlighting)]
+        public async Task TestInTopLevelStatements()
+        {
+            await TestAsync(
+@"if (args.Length > 0) [|return|] 0;
+{|Cursor:[|return|]|} 1;");
+        }
     }
 }

src/Features/CSharp/Portable/Highlighting/KeywordHighlighters/AsyncAwaitHighlighter.cs

@@ -31,7 +31,7 @@ public AsyncAwaitHighlighter()
         }
 
         protected override bool IsHighlightableNode(SyntaxNode node)
-            => node.IsReturnableConstruct();
+            => node.IsReturnableConstructOrTopLevelCompilationUnit();
 
         protected override void AddHighlightsForNode(SyntaxNode node, List<TextSpan> highlights, CancellationToken cancellationToken)
         {

src/Features/CSharp/Portable/Highlighting/KeywordHighlighters/ReturnStatementHighlighter.cs

@@ -32,7 +32,7 @@ protected override void AddHighlights(
         {
             var parent = returnStatement
                              .GetAncestorsOrThis<SyntaxNode>()
-                             .FirstOrDefault(n => n.IsReturnableConstruct());
+                             .FirstOrDefault(n => n.IsReturnableConstructOrTopLevelCompilationUnit());
 
             if (parent == null)
             {

src/Features/Core/Portable/Highlighting/Keywords/AbstractKeywordHighlighter.cs

@@ -5,11 +5,9 @@
 #nullable disable
 
 using System.Collections.Generic;
-using System.Linq;
 using System.Threading;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Text;
-using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Highlighting
 {

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/CSharp/Extensions/SyntaxNodeExtensions.cs

@@ -626,6 +626,9 @@ public static bool IsReturnableConstruct(this SyntaxNode node)
             return false;
         }
 
+        public static bool IsReturnableConstructOrTopLevelCompilationUnit(this SyntaxNode node)
+            => node.IsReturnableConstruct() || (node is CompilationUnitSyntax compilationUnit && compilationUnit.Members.Any(SyntaxKind.GlobalStatement));
+
         public static bool SpansPreprocessorDirective<TSyntaxNode>(this IEnumerable<TSyntaxNode> list) where TSyntaxNode : SyntaxNode
             => CSharpSyntaxFacts.Instance.SpansPreprocessorDirective(list);