fix: find async state machines in optimized builds

ArchUnit.sln

@@ -65,6 +65,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "ClassAssembly", "TestAssemb
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "SlicesTestAssembly", "TestAssemblies\SlicesTestAssembly\SlicesTestAssembly.csproj", "{E421FA12-BBB9-4A8B-BD29-8757836079A5}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "OptimizedAssembly", "TestAssemblies\OptimizedAssembly\OptimizedAssembly.csproj", "{74857E33-09C5-43C7-AA27-E9373DEC1380}"
+EndProject
 Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Solution Items", "Solution Items", "{A780B3D6-C525-4270-ABF1-63BA65EA1593}"
 	ProjectSection(SolutionItems) = preProject
 		README.md = README.md
@@ -441,6 +443,18 @@ Global
 		{E421FA12-BBB9-4A8B-BD29-8757836079A5}.Release|x64.Build.0 = Release|Any CPU
 		{E421FA12-BBB9-4A8B-BD29-8757836079A5}.Release|x86.ActiveCfg = Release|Any CPU
 		{E421FA12-BBB9-4A8B-BD29-8757836079A5}.Release|x86.Build.0 = Release|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Debug|x64.ActiveCfg = Debug|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Debug|x64.Build.0 = Debug|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Debug|x86.ActiveCfg = Debug|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Debug|x86.Build.0 = Debug|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Release|Any CPU.Build.0 = Release|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Release|x64.ActiveCfg = Release|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Release|x64.Build.0 = Release|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Release|x86.ActiveCfg = Release|Any CPU
+		{74857E33-09C5-43C7-AA27-E9373DEC1380}.Release|x86.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -459,6 +473,7 @@ Global
 		{1D0187EB-9D04-4CF3-AE63-0C0E97FCB49C} = {B1191F18-91CB-4387-B775-A5EB64D3AC30}
 		{2C04EB93-5AE0-474C-B328-145A734A5E2B} = {B1191F18-91CB-4387-B775-A5EB64D3AC30}
 		{E421FA12-BBB9-4A8B-BD29-8757836079A5} = {B1191F18-91CB-4387-B775-A5EB64D3AC30}
+		{74857E33-09C5-43C7-AA27-E9373DEC1380} = {B1191F18-91CB-4387-B775-A5EB64D3AC30}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {D1A8274D-32D6-44DB-9BB3-1A5B273709AF}

ArchUnitNET/Loader/MonoCecilMemberExtensions.cs

@@ -289,6 +289,32 @@ internal static bool IsAsync(this MethodDefinition methodDefinition)
             );
         }
 
+        /// <summary>
+        /// Resolves the compiler-generated state machine type of an async or iterator method
+        /// from its <see cref="System.Runtime.CompilerServices.AsyncStateMachineAttribute"/> or
+        /// <see cref="System.Runtime.CompilerServices.IteratorStateMachineAttribute"/>. Unlike
+        /// looking for the <c>newobj</c> that creates it, this also works for the struct state
+        /// machines the compiler emits for async methods in optimized builds.
+        /// </summary>
+        [CanBeNull]
+        internal static TypeDefinition GetStateMachineType(this MethodDefinition methodDefinition)
+        {
+            var stateMachineAttribute = methodDefinition.CustomAttributes.FirstOrDefault(att =>
+                att.AttributeType.FullName
+                    == typeof(System.Runtime.CompilerServices.AsyncStateMachineAttribute).FullName
+                || att.AttributeType.FullName
+                    == typeof(System.Runtime.CompilerServices.IteratorStateMachineAttribute).FullName
+            );
+            if (stateMachineAttribute == null || !stateMachineAttribute.HasConstructorArguments)
+            {
+                return null;
+            }
+
+            var stateMachineType =
+                stateMachineAttribute.ConstructorArguments[0].Value as TypeReference;
+            return stateMachineType?.Resolve();
+        }
+
         internal static bool IsCompilerGenerated(this MemberReference memberReference)
         {
             if (memberReference.Name.HasCompilerGeneratedName())

ArchUnitNET/Loader/TypeProcessor.cs

@@ -864,23 +864,18 @@ private static void HandleIterator(
             DomainResolver domainResolver
         )
         {
-            var compilerGeneratedGeneratorObject = methodBody
-                .Instructions.Where(inst => inst.IsNewObjectOp())
-                .Select(inst => ((MethodReference)inst.Operand).DeclaringType.Resolve())
-                .FirstOrDefault(type =>
-                    type != null
-                    && type.Methods.Any(method => method.Name == nameof(IEnumerator.MoveNext))
-                );
+            var compilerGeneratedGeneratorObject = methodBody.Method.GetStateMachineType();
+            var moveNextMethod = compilerGeneratedGeneratorObject?.Methods.FirstOrDefault(method =>
+                method.Name == nameof(IEnumerator.MoveNext)
+            );
 
-            if (compilerGeneratedGeneratorObject == null)
+            if (moveNextMethod == null)
             {
                 methodDefinition = methodBody.Method;
                 return;
             }
 
-            methodDefinition = compilerGeneratedGeneratorObject.Methods.First(method =>
-                method.Name == nameof(IEnumerator.MoveNext)
-            );
+            methodDefinition = moveNextMethod;
             visitedMethodReferences.Add(methodDefinition);
             methodBody = methodDefinition.Body;
 
@@ -909,25 +904,18 @@ private static void HandleAsync(
             DomainResolver domainResolver
         )
         {
-            var compilerGeneratedGeneratorObject = methodBody
-                .Instructions.Where(inst => inst.IsNewObjectOp())
-                .Select(inst => ((MethodReference)inst.Operand).DeclaringType.Resolve())
-                .FirstOrDefault(type =>
-                    type != null
-                    && type.Methods.Any(method =>
-                        method.Name == nameof(IAsyncStateMachine.MoveNext)
-                    )
-                );
+            var compilerGeneratedGeneratorObject = methodBody.Method.GetStateMachineType();
+            var moveNextMethod = compilerGeneratedGeneratorObject?.Methods.FirstOrDefault(method =>
+                method.Name == nameof(IAsyncStateMachine.MoveNext)
+            );
 
-            if (compilerGeneratedGeneratorObject == null)
+            if (moveNextMethod == null)
             {
                 methodDefinition = methodBody.Method;
                 return;
             }
 
-            methodDefinition = compilerGeneratedGeneratorObject.Methods.First(method =>
-                method.Name == nameof(IAsyncStateMachine.MoveNext)
-            );
+            methodDefinition = moveNextMethod;
 
             visitedMethodReferences.Add(methodDefinition);
             methodBody = methodDefinition.Body;

ArchUnitNETTests/ArchUnitNETTests.csproj

@@ -36,6 +36,7 @@
     <ProjectReference Include="..\TestAssemblies\PropertyMemberAssembly\PropertyMemberAssembly.csproj" />
     <ProjectReference Include="..\TestAssemblies\ClassAssembly\ClassAssembly.csproj" />
     <ProjectReference Include="..\TestAssemblies\SlicesTestAssembly\SlicesTestAssembly.csproj" />
+    <ProjectReference Include="..\TestAssemblies\OptimizedAssembly\OptimizedAssembly.csproj" />
   </ItemGroup>
   <ItemGroup>
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.14.1" />

ArchUnitNETTests/Dependencies/OptimizedStateMachineDependenciesTests.cs

@@ -0,0 +1,79 @@
+using System.Linq;
+using System.Reflection;
+using System.Runtime.CompilerServices;
+using ArchUnitNET.Domain;
+using ArchUnitNET.Domain.Extensions;
+using OptimizedAssembly;
+using Xunit;
+using static ArchUnitNET.Fluent.ArchRuleDefinition;
+
+namespace ArchUnitNETTests.Dependencies
+{
+    /// <summary>
+    /// Async and iterator methods in an optimized build. There the compiler emits the async
+    /// state machine as a struct, which is initialized in place rather than created with
+    /// <c>newobj</c>, so it has to be found through the method's state machine attribute.
+    /// </summary>
+    public class OptimizedStateMachineDependenciesTests
+    {
+        private static readonly Architecture Architecture =
+            StaticTestArchitectures.OptimizedArchitecture;
+
+        private readonly MethodMember _calledMethod;
+
+        public OptimizedStateMachineDependenciesTests()
+        {
+            _calledMethod = Architecture
+                .GetClassOfType(typeof(CalledClass))
+                .GetMethodMembersWithName(nameof(CalledClass.CalledMethod) + "()")
+                .First();
+        }
+
+        [Fact]
+        public void AsyncStateMachineIsAStruct()
+        {
+            var stateMachineType = typeof(ClassWithAsyncMethod)
+                .GetMethod(nameof(ClassWithAsyncMethod.AsyncMethod))
+                .GetCustomAttribute<AsyncStateMachineAttribute>()
+                .StateMachineType;
+
+            Assert.True(stateMachineType.IsValueType);
+        }
+
+        [Fact]
+        public void AsyncMethodBodyDependenciesAreFound()
+        {
+            var asyncMethod = Architecture
+                .GetClassOfType(typeof(ClassWithAsyncMethod))
+                .GetMethodMembersWithName(nameof(ClassWithAsyncMethod.AsyncMethod) + "()")
+                .First();
+
+            Assert.Contains(_calledMethod, asyncMethod.GetCalledMethods());
+
+            var rule = Classes()
+                .That()
+                .HaveFullName(typeof(ClassWithAsyncMethod).FullName)
+                .Should()
+                .NotDependOnAny(typeof(CalledClass));
+            Assert.False(rule.HasNoViolations(Architecture));
+        }
+
+        [Fact]
+        public void IteratorMethodBodyDependenciesAreFound()
+        {
+            var iteratorMethod = Architecture
+                .GetClassOfType(typeof(ClassWithIteratorMethod))
+                .GetMethodMembersWithName(nameof(ClassWithIteratorMethod.IteratorMethod) + "()")
+                .First();
+
+            Assert.Contains(_calledMethod, iteratorMethod.GetCalledMethods());
+
+            var rule = Classes()
+                .That()
+                .HaveFullName(typeof(ClassWithIteratorMethod).FullName)
+                .Should()
+                .NotDependOnAny(typeof(CalledClass));
+            Assert.False(rule.HasNoViolations(Architecture));
+        }
+    }
+}

ArchUnitNETTests/Domain/Dependencies/Members/MethodCallDependencyTests.cs

@@ -96,7 +96,7 @@ MethodCallDependency expectedDependency
             Assert.Contains(expectedDependency, originMember.GetMethodCallDependencies());
         }
 
-        [SkipInReleaseBuildTheory]
+        [Theory]
         [ClassData(typeof(MethodDependencyTestBuild.MethodCallDependencyInAsyncMethodTestData))]
         public void MethodCallDependenciesAreFoundInAsyncMethod(
             IMember originMember,

ArchUnitNETTests/SkipInReleaseBuild.cs

@@ -8,16 +8,6 @@ public SkipInReleaseBuild()
         {
 #if !DEBUG
             Skip = "This test only works in debug build";
-#endif
-        }
-    }
-
-    public sealed class SkipInReleaseBuildTheory : TheoryAttribute
-    {
-        public SkipInReleaseBuildTheory()
-        {
-#if !DEBUG
-            Skip = "This test only works in debug build";
 #endif
         }
     }

ArchUnitNETTests/StaticTestArchitectures.cs

@@ -102,6 +102,10 @@ public static class StaticTestArchitectures
             )
             .Build();
 
+        public static readonly Architecture OptimizedArchitecture = new ArchLoader()
+            .LoadAssemblies(typeof(OptimizedAssembly.ClassWithAsyncMethod).Assembly)
+            .Build();
+
         public static readonly Architecture ArchUnitNETTestAssemblyArchitecture = new ArchLoader()
             .LoadAssemblies(typeof(Class1).Assembly)
             .Build();

TestAssemblies/OptimizedAssembly/OptimizedAssembly.csproj

@@ -0,0 +1,15 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFramework>net10.0</TargetFramework>
+    <ImplicitUsings>enable</ImplicitUsings>
+    <Nullable>enable</Nullable>
+    <TreatWarningsAsErrors>true</TreatWarningsAsErrors>
+    <IsPackable>false</IsPackable>
+    <SignAssembly>True</SignAssembly>
+    <AssemblyOriginatorKeyFile>..\..\strongNameKey.snk</AssemblyOriginatorKeyFile>
+    <AssemblyVersion>1.0.0.0</AssemblyVersion>
+    <!-- Always optimized, even in Debug, so the compiler emits the code shapes of a Release
+         build (e.g. struct instead of class state machines for async methods). -->
+    <Optimize>true</Optimize>
+  </PropertyGroup>
+</Project>

TestAssemblies/OptimizedAssembly/StateMachineClasses.cs

@@ -0,0 +1,24 @@
+namespace OptimizedAssembly;
+
+public class CalledClass
+{
+    public static void CalledMethod() { }
+}
+
+public class ClassWithAsyncMethod
+{
+    public async Task AsyncMethod()
+    {
+        await Task.Yield();
+        CalledClass.CalledMethod();
+    }
+}
+
+public class ClassWithIteratorMethod
+{
+    public IEnumerable<int> IteratorMethod()
+    {
+        CalledClass.CalledMethod();
+        yield return 0;
+    }
+}