Fix crossgen2 OverflowException emitting inlining info for runtime-async inliners

docs/design/coreclr/botr/readytorun-format.md

@@ -336,7 +336,7 @@ ECMA 335 does not have a natural encoding for describing an overridden method. T
 
 #### IL Body signatures
 
-ECMA 335 does not define a format that can represent the exact implementation of a method by itself. This signature holds all of the IL of the method, the EH table, the locals table, and each token (other than type references) in those tables is replaced with an index into a local stream of signatures. Those signatures are simply verbatim copies of the needed metadata to describe MemberRefs, TypeSpecs, MethodSpecs, StandaloneSignatures and strings. All of that is bundled into a large byte array. In addition, a series of TypeSignatures follows which allow the type references to be resolved, as well as a methodreference to the uninstantiated method. Assuming all of this matches with the data that is present at runtime, the fixup is considered to be satisfied. See ReadyToRunStandaloneMetadata.cs for the exact details of the format.
+ECMA 335 does not define a format that can represent the exact implementation of a method by itself. This signature holds all of the IL of the method, the EH table, the locals table, and each token (other than type references) in those tables is replaced with an index into a local stream of signatures. Those signatures are simply verbatim copies of the needed metadata to describe MemberRefs, TypeSpecs, MethodSpecs, StandaloneSignatures and strings. All of that is bundled into a large byte array. The byte that describes the locals table also records the method implementation flags that change how the same IL executes: `0x04` is set for `miAsync` (runtime-async) methods and `0x08` for `miSynchronized` methods. In addition, a series of TypeSignatures follows which allow the type references to be resolved, as well as a methodreference to the uninstantiated method. Assuming all of this matches with the data that is present at runtime, the fixup is considered to be satisfied. See ReadyToRunStandaloneMethodMetadata.cs for the exact details of the format.
 
 #### InjectStringThunks signatures
 

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/TestCases/CrossModuleInlining/AsyncGenericInlinerConsumer.cs

@@ -0,0 +1,37 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+
+/// <summary>
+/// Consumer that awaits a runtime-async method on a generic type from AsyncCrossModuleGenericLib,
+/// instantiated over a value type defined in this assembly so that the async variant of
+/// AsyncGenericWrapper&lt;LocalAsyncStruct&gt;.InvokeGetValueAsync() is compiled into this image.
+/// </summary>
+
+public struct LocalAsyncStruct { public int Value; }
+
+public static class AsyncGenericInlinerConsumer
+{
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static async Task<int> UseAsync()
+    {
+        var wrapper = new AsyncGenericWrapper<LocalAsyncStruct>(new LocalAsyncStruct { Value = 1 });
+        return await wrapper.InvokeGetValueAsync();
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static async Task<int> AwaitValueTask()
+    {
+        var wrapper = new AsyncGenericWrapper<LocalAsyncStruct>(new LocalAsyncStruct { Value = 2 });
+        return await wrapper.GetValueTask();
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static Task<int> ReturnValueTask()
+    {
+        var wrapper = new AsyncGenericWrapper<LocalAsyncStruct>(new LocalAsyncStruct { Value = 3 });
+        return wrapper.GetValueTask();
+    }
+}

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/TestCases/CrossModuleInlining/Dependencies/AsyncCrossModuleGenericLib.cs

@@ -0,0 +1,41 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+
+/// <summary>
+/// Library with a generic type whose runtime-async method inlines a utility method.
+/// When the consumer awaits the method on an instantiation over one of its own value types,
+/// crossgen2 compiles the async variant of the method into the consumer as a CrossModuleCompileable
+/// generic, making the async variant a cross-module inliner of AsyncGenericUtility.GetAsyncGenericValue.
+/// </summary>
+public static class AsyncGenericUtility
+{
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static int GetAsyncGenericValue() => 42;
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static int GetSharedInlineeValue() => 43;
+}
+
+public class AsyncGenericWrapper<T>
+{
+    private T _value;
+
+    public AsyncGenericWrapper(T value) => _value = value;
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public async Task<int> InvokeGetValueAsync()
+    {
+        await Task.Yield();
+        return AsyncGenericUtility.GetAsyncGenericValue();
+    }
+
+    // Not async, so both the task-returning method and its async variant are compiled from this IL.
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public Task<int> GetValueTask()
+    {
+        return Task.FromResult(AsyncGenericUtility.GetSharedInlineeValue());
+    }
+}

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/TestCases/R2RTestSuites.cs

@@ -1995,6 +1995,56 @@ static void Validate(ReadyToRunReader reader)
         }
     }
 
+    /// <summary>
+    /// Tests cross-module generic compilation where the runtime-async variant of a method from an
+    /// --opt-cross-module library is compiled into the consumer and inlines another library method.
+    /// The inlining info must reference the IL body fixup that was recorded for the async variant.
+    /// </summary>
+    [ConditionalFact(typeof(TestPaths), nameof(TestPaths.IsNotWasmTarget))]
+    public void AsyncCrossModuleGenericInliner()
+    {
+        var asyncCrossModuleGenericLib = new CompiledAssembly
+        {
+            AssemblyName = "AsyncCrossModuleGenericLib",
+            SourceResourceNames = ["CrossModuleInlining/Dependencies/AsyncCrossModuleGenericLib.cs"],
+            Features = { RuntimeAsyncFeature },
+        };
+        var consumer = new CompiledAssembly
+        {
+            AssemblyName = "AsyncGenericInlinerConsumer",
+            SourceResourceNames = ["CrossModuleInlining/AsyncGenericInlinerConsumer.cs"],
+            References = [asyncCrossModuleGenericLib],
+            Features = { RuntimeAsyncFeature },
+        };
+
+        new R2RTestRunner(_output).Run(new R2RTestCase(
+            nameof(AsyncCrossModuleGenericInliner),
+            [
+                new(consumer.AssemblyName,
+                [
+                    new CrossgenAssembly(asyncCrossModuleGenericLib)
+                    {
+                        Kind = Crossgen2InputKind.Reference,
+                        Options = [Crossgen2AssemblyOption.CrossModuleOptimization],
+                    },
+                    new CrossgenAssembly(consumer),
+                ])
+                {
+                    Validate = Validate,
+                },
+            ]));
+
+        static void Validate(ReadyToRunReader reader)
+        {
+            string diag;
+            Assert.True(R2RAssert.HasManifestRef(reader, "AsyncCrossModuleGenericLib", out diag), diag);
+            Assert.True(R2RAssert.HasCrossModuleInliningInfo(reader, out diag), diag);
+            Assert.True(R2RAssert.HasCrossModuleInliners(reader, "GetAsyncGenericValue", ["InvokeGetValueAsync"], out diag), diag);
+            Assert.True(R2RAssert.HasAsyncVariant(reader, "GetValueTask", out diag), diag);
+            Assert.True(R2RAssert.HasCrossModuleInlinerCount(reader, "GetSharedInlineeValue", "GetValueTask", 1, out diag), diag);
+        }
+    }
+
     [ConditionalFact(typeof(TestPaths), nameof(TestPaths.IsNotWasmTarget))]
     public void VirtualMethodGenericsNonGVM()
     {

src/coreclr/tools/aot/ILCompiler.ReadyToRun.Tests/TestCasesRunner/R2RResultChecker.cs

@@ -637,6 +637,48 @@ public static bool HasCrossModuleInliners(
         return false;
     }
 
+    /// <summary>
+    /// Returns true if the CrossModuleInlineInfo entry for an inlinee matching <paramref name="inlineeMethodName"/>
+    /// has exactly <paramref name="expectedCount"/> cross-module inliners whose resolved names contain
+    /// <paramref name="inlinerMethodName"/>.
+    /// </summary>
+    public static bool HasCrossModuleInlinerCount(
+        ReadyToRunReader reader,
+        string inlineeMethodName,
+        string inlinerMethodName,
+        int expectedCount,
+        out string diagnostic)
+    {
+        if (!TryGetCrossModuleInliningInfoSection(reader, out var inliningInfo, out diagnostic))
+            return false;
+
+        foreach (var entry in inliningInfo.GetEntries())
+        {
+            string inlineeName = inliningInfo.ResolveMethodName(entry.Inlinee);
+            if (!inlineeName.Contains(inlineeMethodName, StringComparison.OrdinalIgnoreCase))
+                continue;
+
+            var matchingInliners = new List<string>();
+            foreach (var inliner in entry.Inliners)
+            {
+                if (!inliner.IsCrossModule)
+                    continue;
+
+                string inlinerName = inliningInfo.ResolveMethodName(inliner);
+                if (inlinerName.Contains(inlinerMethodName, StringComparison.OrdinalIgnoreCase))
+                    matchingInliners.Add(inlinerName);
+            }
+
+            diagnostic =
+                $"Inlinee '{inlineeName}': expected {expectedCount} cross-module inliner(s) matching '{inlinerMethodName}', " +
+                $"found {matchingInliners.Count}:\n  {string.Join("\n  ", matchingInliners)}";
+            return matchingInliners.Count == expectedCount;
+        }
+
+        diagnostic = $"No CrossModuleInlineInfo entry found for inlinee matching '{inlineeMethodName}'.";
+        return false;
+    }
+
     /// <summary>
     /// Returns true if any inlining info section (CrossModuleInlineInfo or InliningInfo2) records
     /// that <paramref name="inlinerMethodName"/> inlined <paramref name="inlineeMethodName"/>.

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRun/ILBodyFixupSignature.cs

@@ -51,6 +51,15 @@ public ILBodyFixupSignature(ReadyToRunFixupKind fixupKind, MethodDesc signatureM
             _signatureMethod = signatureMethod;
         }
 
+        /// <summary>
+        /// Returns the method a compiled method's own IL body fixup is created for, or null if the compiled method cannot have one.
+        /// </summary>
+        public static MethodDesc GetSignatureMethodForCompiledMethod(MethodDesc compiledMethod)
+        {
+            MethodDesc typicalDef = compiledMethod.GetTypicalMethodDefinition();
+            return typicalDef is EcmaMethod or AsyncMethodVariant ? typicalDef : null;
+        }
+
         public override int ClassCode => 308579267;
 
         protected override void OnMarked(NodeFactory context)

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRun/InliningInfoNode.cs

@@ -60,7 +60,8 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
             if (relocsOnly)
                 return new ObjectData(Array.Empty<byte>(), Array.Empty<Relocation>(), 1, new ISymbolDefinitionNode[] { this });
 
-            Dictionary<MethodDesc, HashSet<MethodDesc>> inlineeToInliners = new Dictionary<MethodDesc, HashSet<MethodDesc>>();
+            // Inliners are keyed by metadata definition; the value is the identity whose Check_IL_Body fixup was recorded when compiling it.
+            Dictionary<MethodDesc, Dictionary<EcmaMethod, MethodDesc>> inlineeToInliners = new Dictionary<MethodDesc, Dictionary<EcmaMethod, MethodDesc>>();
 
             // Build a map from inlinee to the list of inliners
             // We are only interested in the generic definitions of these.
@@ -83,6 +84,21 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
                 // Only encode inlining info for inliners within the active module, or if cross module inline format is in use
                 Debug.Assert(AllowCrossModuleInlines || (inlinerDefinition.Module == _module));
 
+                MethodDesc inlinerIdentity;
+                if (AllowCrossModuleInlines && !factory.CompilationModuleGroup.VersionsWithMethodBody(inlinerDefinition))
+                {
+                    inlinerIdentity = ILBodyFixupSignature.GetSignatureMethodForCompiledMethod(inliner);
+
+                    // Cross-module inliners are encoded by their own Check_IL_Body import. Thunks and stubs don't have one.
+                    if (inlinerIdentity is null || inlinerIdentity.IsCompilerGeneratedILBodyForAsync())
+                        continue;
+                }
+                else
+                {
+                    // Inliners in the version bubble are encoded by metadata RID.
+                    inlinerIdentity = inlinerDefinition;
+                }
+
                 bool inlinerReportAllVersionsWithInlinee = !AllowCrossModuleInlines || factory.CompilationModuleGroup.CrossModuleCompileable(inlinerDefinition);
 
                 foreach (MethodDesc inlinee in inlinees)
@@ -123,12 +139,17 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
                         }
                     }
 
-                    if (!inlineeToInliners.TryGetValue(inlineeDefinition, out HashSet<MethodDesc> inliners))
+                    if (!inlineeToInliners.TryGetValue(inlineeDefinition, out Dictionary<EcmaMethod, MethodDesc> inliners))
                     {
-                        inliners = new HashSet<MethodDesc>();
+                        inliners = new Dictionary<EcmaMethod, MethodDesc>();
                         inlineeToInliners.Add(inlineeDefinition, inliners);
                     }
-                    inliners.Add(inlinerDefinition);
+
+                    // Both variants of a method may inline the same inlinee; report it once, independent of enumeration order.
+                    if (!inliners.TryGetValue(inlinerDefinition, out MethodDesc existingInliner) || (inlinerIdentity is EcmaMethod && existingInliner is not EcmaMethod))
+                    {
+                        inliners[inlinerDefinition] = inlinerIdentity;
+                    }
                 }
             }
 
@@ -176,12 +197,8 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
                         sig.Append(new UnsignedConstant((uint)factory.ManifestMetadataTable.ModuleToIndex(ecmaInlinee.Module)));
                     }
 
-                    // We're only concerned with metadata here, so we can convert all to EcmaMethod and lose info about AsyncVariant vs Task-Returning
-                    List<EcmaMethod> sortedInliners = new List<EcmaMethod>(inlineeWithInliners.Value.Count);
-                    foreach (var inliner in inlineeWithInliners.Value)
-                    {
-                        sortedInliners.Add((EcmaMethod)inliner.GetPrimaryMethodDesc());
-                    }
+                    // We're only concerned with metadata here, so we can use the EcmaMethod keys and lose info about AsyncVariant vs Task-Returning
+                    List<EcmaMethod> sortedInliners = new List<EcmaMethod>(inlineeWithInliners.Value.Keys);
                     sortedInliners.MergeSort((a, b) =>
                     {
                         if (a == b)
@@ -242,8 +259,8 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
                     bool isCrossModuleInlinee = !factory.CompilationModuleGroup.VersionsWithMethodBody(inlinee);
                     Debug.Assert(!isCrossModuleInlinee || factory.CompilationModuleGroup.CrossModuleInlineable(inlinee));
 
-                    EcmaMethod[] sortedInliners = new EcmaMethod[inlineeWithInliners.Value.Count];
-                    inlineeWithInliners.Value.CopyTo(sortedInliners);
+                    MethodDesc[] sortedInliners = new MethodDesc[inlineeWithInliners.Value.Count];
+                    inlineeWithInliners.Value.Values.CopyTo(sortedInliners, 0);
 
                     sortedInliners.MergeSort((a, b) =>
                     {
@@ -270,14 +287,16 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
                         }
                         else
                         {
-                            int aRid = MetadataTokens.GetRowNumber(a.Handle);
-                            int bRid = MetadataTokens.GetRowNumber(b.Handle);
+                            EcmaMethod ecmaA = (EcmaMethod)a.GetPrimaryMethodDesc();
+                            EcmaMethod ecmaB = (EcmaMethod)b.GetPrimaryMethodDesc();
+                            int aRid = MetadataTokens.GetRowNumber(ecmaA.Handle);
+                            int bRid = MetadataTokens.GetRowNumber(ecmaB.Handle);
                             if (aRid < bRid)
                                 return -1;
                             else if (aRid > bRid)
                                 return 1;
 
-                            result = a.Module.CompareTo(b.Module);
+                            result = ecmaA.Module.CompareTo(ecmaB.Module);
                         }
                         Debug.Assert(result != 0);
                         return result;
@@ -336,7 +355,7 @@ public override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
                         uint baseRid = 0;
                         for (; inlinerIndex < sortedInliners.Length; inlinerIndex++)
                         {
-                            var inliner = sortedInliners[inlinerIndex];
+                            var inliner = (EcmaMethod)sortedInliners[inlinerIndex].GetPrimaryMethodDesc();
                             uint inlinerRid = (uint)MetadataTokens.GetRowNumber(inliner.Handle);
                             uint ridDelta = inlinerRid - baseRid;
                             baseRid = inlinerRid;

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunStandaloneMethodMetadata.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Reflection;
 using System.Reflection.Metadata;
 using System.Reflection.Metadata.Ecma335;
 
@@ -19,6 +20,13 @@ namespace ILCompiler
     // That code can be found in src\coreclr\vm\readytorunstandalonemethodmetadata.cpp
     public class ReadyToRunStandaloneMethodMetadata
     {
+        // Values of the byte that follows the EH clauses. Must match the VM.
+        private const byte LocalsNotInitialized = 0;
+        private const byte LocalsInitialized = 1;
+        private const byte NoLocals = 2;
+        private const byte AsyncImplFlag = 4;
+        private const byte SynchronizedImplFlag = 8;
+
         public byte[] ConstantData;
         public TypeDesc[] TypeRefs;
 
@@ -92,14 +100,17 @@ public static ReadyToRunStandaloneMethodMetadata Compute(EcmaMethod wrappedMetho
                     }
                 }
 
-                if (localsBlob.Length == 0)
-                {
-                    // No locals. Encode a 2 to indicate this
-                    _nonCodeAlternateBlob.WriteByte(2);
-                }
-                else
+                // Impl flags that change how the same IL executes are part of the IL body identity.
+                byte localsAndImplFlags = localsBlob.Length == 0 ? NoLocals : (_methodBody.LocalVariablesInitialized ? LocalsInitialized : LocalsNotInitialized);
+                MethodImplAttributes implAttributes = metadataReader.GetMethodDefinition(wrappedMethod.Handle).ImplAttributes;
+                if ((implAttributes & MethodImplAttributes.Async) != 0)
+                    localsAndImplFlags |= AsyncImplFlag;
+                if ((implAttributes & MethodImplAttributes.Synchronized) != 0)
+                    localsAndImplFlags |= SynchronizedImplFlag;
+                _nonCodeAlternateBlob.WriteByte(localsAndImplFlags);
+
+                if (localsBlob.Length != 0)
                 {
-                    _nonCodeAlternateBlob.WriteByte(_methodBody.LocalVariablesInitialized ? (byte)1 : (byte)0);
                     EcmaSignatureTranslator sigTranslator = new EcmaSignatureTranslator(localsBlob, _nonCodeAlternateBlob, GetAlternateStreamToken);
                     sigTranslator.ParseLocalsSignature();
                 }