From ded4efd6461bdb715638652ba812d45bdd924c81 Mon Sep 17 00:00:00 2001
From: Bruno Garcia <bruno@brunogarcia.com>
Date: Sun, 3 Jan 2021 12:13:39 -0500
Subject: [PATCH] ns2.1 and TryResolveStateMachineMethod

---
 directory.build.props                         |  1 +
 src/Ben.Demystifier/Ben.Demystifier.csproj    |  3 +-
 .../EnhancedStackTrace.Frames.cs              | 54 ++++++++++++++-----
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/directory.build.props b/directory.build.props
index 1588fce..6025539 100644
--- a/directory.build.props
+++ b/directory.build.props
@@ -1,4 +1,5 @@
 <Project>
   <PropertyGroup>
+    <LangVersion>latest</LangVersion>
   </PropertyGroup>
 </Project>
diff --git a/src/Ben.Demystifier/Ben.Demystifier.csproj b/src/Ben.Demystifier/Ben.Demystifier.csproj
index 22997fd..71db297 100644
--- a/src/Ben.Demystifier/Ben.Demystifier.csproj
+++ b/src/Ben.Demystifier/Ben.Demystifier.csproj
@@ -15,8 +15,7 @@
   </PropertyGroup>
 
   <PropertyGroup>
-    <TargetFrameworks>netstandard2.0;net45</TargetFrameworks>
-    <!--<TargetFrameworks>netstandard2.0;netstandard2.1;net45</TargetFrameworks>-->
+    <TargetFrameworks>netstandard2.1;netstandard2.0;net45</TargetFrameworks>
     <SignAssembly>true</SignAssembly>
     <AssemblyOriginatorKeyFile>key.snk</AssemblyOriginatorKeyFile>
   </PropertyGroup>
diff --git a/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs b/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs
index 03fd069..933b88c 100644
--- a/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs
+++ b/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs
@@ -804,38 +804,64 @@ namespace System.Diagnostics
             return false;
         }
 
+        // https://github.com/dotnet/runtime/blob/c985bdcec2a9190e733bcada413a193d5ff60c0d/src/libraries/System.Private.CoreLib/src/System/Diagnostics/StackTrace.cs#L375-L430
         private static bool TryResolveStateMachineMethod(ref MethodBase method, out Type declaringType)
         {
-            Debug.Assert(method != null);
-            Debug.Assert(method.DeclaringType != null);
-
+            if (method.DeclaringType is null)
+            {
+                declaringType = null!;
+                return false;
+            }
             declaringType = method.DeclaringType;
 
             var parentType = declaringType.DeclaringType;
-            if (parentType == null)
+            if (parentType is null)
             {
                 return false;
             }
 
-            var methods = parentType.GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static | BindingFlags.Instance | BindingFlags.DeclaredOnly);
+            static MethodInfo[] GetDeclaredMethods(Type type) =>
+                type.GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static | BindingFlags.Instance | BindingFlags.DeclaredOnly);
 
-            foreach (var candidateMethod in methods)
+            var methods = GetDeclaredMethods(parentType);
+            if (methods == null)
             {
-                var attributes = candidateMethod.GetCustomAttributes<StateMachineAttribute>();
+                return false;
+            }
 
-                foreach (var asma in attributes)
+            foreach (MethodInfo candidateMethod in methods)
+            {
+                var attributes = candidateMethod.GetCustomAttributes<StateMachineAttribute>(inherit: false);
+                // ReSharper disable once ConditionIsAlwaysTrueOrFalse - Taken from CoreFX
+                if (attributes is null)
+                {
+                    continue;
+                }
+
+                bool foundAttribute = false, foundIteratorAttribute = false;
+                foreach (StateMachineAttribute asma in attributes)
                 {
                     if (asma.StateMachineType == declaringType)
                     {
-                        method = candidateMethod;
-                        declaringType = candidateMethod.DeclaringType;
-                        // Mark the iterator as changed; so it gets the + annotation of the original method
-                        // async statemachines resolve directly to their builder methods so aren't marked as changed
-                        return asma is IteratorStateMachineAttribute;
+                        foundAttribute = true;
+                        foundIteratorAttribute |= asma is IteratorStateMachineAttribute
+#if NETSTANDARD2_1
+                            || asma is AsyncIteratorStateMachineAttribute
+#endif
+                            ;
                     }
                 }
-            }
 
+                if (foundAttribute)
+                {
+                    // If this is an iterator (sync or async), mark the iterator as changed, so it gets the + annotation
+                    // of the original method. Non-iterator async state machines resolve directly to their builder methods
+                    // so aren't marked as changed.
+                    method = candidateMethod;
+                    declaringType = candidateMethod.DeclaringType!;
+                    return foundIteratorAttribute;
+                }
+            }
             return false;
         }