From 047a498a4a986d9fe626a4497d60214583ac148c Mon Sep 17 00:00:00 2001
From: Andrew Arnott <andrew.arnott@microsoft.com>
Date: Mon, 23 Oct 2023 16:24:10 -0600
Subject: [PATCH] Fix VSTHRD004 mis-fire on `NoThrowAwaitable()`

Fixes #1232
---
 ...004AwaitSwitchToMainThreadAsyncAnalyzer.cs | 19 ++++++++++++++----
 ...aitSwitchToMainThreadAsyncAnalyzerTests.cs | 20 +++++++++++++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/Microsoft.VisualStudio.Threading.Analyzers/AbstractVSTHRD004AwaitSwitchToMainThreadAsyncAnalyzer.cs b/src/Microsoft.VisualStudio.Threading.Analyzers/AbstractVSTHRD004AwaitSwitchToMainThreadAsyncAnalyzer.cs
index 5156134d1..177770600 100644
--- a/src/Microsoft.VisualStudio.Threading.Analyzers/AbstractVSTHRD004AwaitSwitchToMainThreadAsyncAnalyzer.cs
+++ b/src/Microsoft.VisualStudio.Threading.Analyzers/AbstractVSTHRD004AwaitSwitchToMainThreadAsyncAnalyzer.cs
@@ -45,12 +45,23 @@ private void AnalyzeInvocation(OperationAnalysisContext context)
             methodSymbol.ContainingType.Name == Types.JoinableTaskFactory.TypeName &&
             methodSymbol.ContainingType.BelongsToNamespace(Types.JoinableTaskFactory.Namespace))
         {
-            // This is a call to JTF.SwitchToMainThreadAsync(). Is it being (directly) awaited?
-            if (!(invocation.Parent is IAwaitOperation))
+            // This is a call to JTF.SwitchToMainThreadAsync(). Is it being awaited in some ancestor?
+            for (IOperation? parentOp = invocation.Parent; parentOp is not null; parentOp = parentOp.Parent)
             {
-                Location? location = (this.LanguageUtils.IsolateMethodName(invocation) ?? invocation.Syntax).GetLocation();
-                context.ReportDiagnostic(Diagnostic.Create(Descriptor, location));
+                if (parentOp is IAwaitOperation)
+                {
+                    return;
+                }
+
+                if (parentOp is IExpressionStatementOperation or IReturnOperation)
+                {
+                    // We've reached the top of the statement without finding an await.
+                    break;
+                }
             }
+
+            Location? location = (this.LanguageUtils.IsolateMethodName(invocation) ?? invocation.Syntax).GetLocation();
+            context.ReportDiagnostic(Diagnostic.Create(Descriptor, location));
         }
     }
 }
diff --git a/test/Microsoft.VisualStudio.Threading.Analyzers.Tests/VSTHRD004AwaitSwitchToMainThreadAsyncAnalyzerTests.cs b/test/Microsoft.VisualStudio.Threading.Analyzers.Tests/VSTHRD004AwaitSwitchToMainThreadAsyncAnalyzerTests.cs
index 1271292de..54483b7d9 100644
--- a/test/Microsoft.VisualStudio.Threading.Analyzers.Tests/VSTHRD004AwaitSwitchToMainThreadAsyncAnalyzerTests.cs
+++ b/test/Microsoft.VisualStudio.Threading.Analyzers.Tests/VSTHRD004AwaitSwitchToMainThreadAsyncAnalyzerTests.cs
@@ -104,6 +104,26 @@ async Task FooAsync()
         await CSVerify.VerifyAnalyzerAsync(test);
     }
 
+    [Fact]
+    public async Task AsyncMethodWithAwaitNoThrowAwaitable_ProducesNoDiagnostic()
+    {
+        var test = @"
+using System.Threading.Tasks;
+
+class Test
+{
+    Microsoft.VisualStudio.Threading.JoinableTaskFactory jtf;
+
+    async Task FooAsync()
+    {
+        await jtf.SwitchToMainThreadAsync().NoThrowAwaitable();
+    }
+}
+";
+
+        await CSVerify.VerifyAnalyzerAsync(test);
+    }
+
     [Fact]
     public async Task TaskReturningSyncMethod_ProducesDiagnostic()
     {