add rejit support to GetMethodDescData

Author: Max Charlamb

docs/design/datacontracts/CodeVersions.md

@@ -4,6 +4,29 @@ This contract encapsulates support for [code versioning](../features/code-versio
 
 ## APIs of contract
 
+```csharp
+internal readonly struct ILCodeVersionHandle
+{
+    internal readonly TargetPointer Module;
+    internal readonly uint MethodDefinition;
+    internal readonly TargetPointer ILCodeVersionNode;
+    internal ILCodeVersionHandle(TargetPointer module, uint methodDef, TargetPointer ilCodeVersionNodeAddress)
+    {
+        if (module != TargetPointer.Null && ilCodeVersionNodeAddress != TargetPointer.Null)
+            throw new ArgumentException("Both MethodDesc and ILCodeVersionNode cannot be non-null");
+
+        if (module != TargetPointer.Null && methodDef == 0)
+            throw new ArgumentException("MethodDefinition must be non-zero if Module is non-null");
+
+        Module = module;
+        MethodDefinition = methodDef;
+        ILCodeVersionNode = ilCodeVersionNodeAddress;
+    }
+    public static ILCodeVersionHandle Invalid => new ILCodeVersionHandle(TargetPointer.Null, 0, TargetPointer.Null);
+    public bool IsValid => Module != TargetPointer.Null || ILCodeVersionNode != TargetPointer.Null;
+}
+```
+
 ```csharp
 internal struct NativeCodeVersionHandle
 {
@@ -26,10 +49,19 @@ internal struct NativeCodeVersionHandle
 ```
 
 ```csharp
+// Return a handle to the active version of the IL code for a given method descriptor
+public virtual ILCodeVersionHandle GetActiveILCodeVersion(TargetPointer methodDesc);
+// Return a handle to the IL code version representing the given native code version
+public virtual ILCodeVersionHandle GetILCodeVersion(NativeCodeVersionHandle codeVersionHandle);
+// Return all of the IL code versions for a given method descriptor
+public virtual IEnumerable<ILCodeVersionHandle> GetILCodeVersions(TargetPointer methodDesc);
+
 // Return a handle to the version of the native code that includes the given instruction pointer
 public virtual NativeCodeVersionHandle GetNativeCodeVersionForIP(TargetCodePointer ip);
 // Return a handle to the active version of the native code for a given method descriptor
 public virtual NativeCodeVersionHandle GetActiveNativeCodeVersion(TargetPointer methodDesc);
+// Return a handle to the active version of the native code for a given method descriptor and IL code version. The IL code version and method descriptor must represent the same method
+public virtual NativeCodeVersionHandle GetActiveNativeCodeVersionForILCodeVersion(TargetPointer methodDesc, ILCodeVersionHandle ilCodeVersionHandle);
 
 // returns true if the given method descriptor supports multiple code versions
 public virtual bool CodeVersionManagerSupportsMethod(TargetPointer methodDesc);
@@ -52,11 +84,13 @@ Data descriptors used:
 | NativeCodeVersionNode | NativeCode | indicates an explicit native code version node |
 | NativeCodeVersionNode | Flags | `NativeCodeVersionNodeFlags` flags, see below |
 | NativeCodeVersionNode | VersionId | Version ID corresponding to the parent IL code version |
+| ILCodeVersioningState | FirstVersionNode | pointer to the first `ILCodeVersionNode` |
 | ILCodeVersioningState | ActiveVersionKind | an `ILCodeVersionKind` value indicating which fields of the active version are value |
 | ILCodeVersioningState | ActiveVersionNode | if the active version is explicit, the NativeCodeVersionNode for the active version |
 | ILCodeVersioningState | ActiveVersionModule | if the active version is synthetic or unknown, the pointer to the Module that defines the method |
 | ILCodeVersioningState | ActiveVersionMethodDef | if the active version is synthetic or unknown, the MethodDef token for the method |
 | ILCodeVersionNode | VersionId | Version ID of the node |
+| ILCodeVersionNode | Next | Pointer to the next `ILCodeVersionNode`|
 
 The flag indicates that the default version of the code for a method desc is active:
 ```csharp
@@ -93,6 +127,40 @@ Contracts used:
 | Loader |
 | RuntimeTypeSystem |
 
+### Finding all of the ILCodeVersions for a method
+```csharp
+IEnumerable<ILCodeVersionHandle> ICodeVersions.GetILCodeVersions(TargetPointer methodDesc)
+{
+    // CodeVersionManager::GetILCodeVersions
+    IRuntimeTypeSystem rts = _target.Contracts.RuntimeTypeSystem;
+    MethodDescHandle md = rts.GetMethodDescHandle(methodDesc);
+    TargetPointer mtAddr = rts.GetMethodTable(md);
+    TypeHandle typeHandle = rts.GetTypeHandle(mtAddr);
+    TargetPointer module = rts.GetModule(typeHandle);
+    uint methodDefToken = rts.GetMethodToken(md);
+
+    ModuleHandle moduleHandle = _target.Contracts.Loader.GetModuleHandle(module);
+    TargetPointer ilCodeVersionTable = _target.Contracts.Loader.GetLookupTables(moduleHandle).MethodDefToILCodeVersioningState;
+    TargetPointer ilVersionStateAddress = _target.Contracts.Loader.GetModuleLookupMapElement(ilCodeVersionTable, methodDefToken, out var _);
+
+    // always add the synthetic version
+    yield return new ILCodeVersionHandle(module, methodDefToken, TargetPointer.Null);
+
+    // if explicit versions exist, iterate linked list and return them
+    if (ilVersionStateAddress != TargetPointer.Null)
+    {
+        Data.ILCodeVersioningState ilState = _target.ProcessedData.GetOrAdd<Data.ILCodeVersioningState>(ilVersionStateAddress);
+        TargetPointer nodePointer = ilState.FirstVersionNode;
+        while (nodePointer != TargetPointer.Null)
+        {
+            Data.ILCodeVersionNode current = _target.ProcessedData.GetOrAdd<Data.ILCodeVersionNode>(nodePointer);
+            yield return new ILCodeVersionHandle(TargetPointer.Null, 0, nodePointer);
+            nodePointer = current.Next;
+        }
+    }
+}
+```
+
 ### Finding the start of a specific native code version
 
 ```csharp
@@ -235,7 +303,7 @@ bool IsActiveNativeCodeVersion(NativeCodeVersionHandle nativeCodeVersion)
 
 NativeCodeVersionHandle FindActiveNativeCodeVersion(ILCodeVersionHandle methodDefActiveVersion, TargetPointer methodDescAddress)
 {
-    TargetNUInt? ilVersionId = default;
+    TargetNUInt ilVersionId = GetId(ilcodeVersion);
     if (methodDefActiveVersion.Module != TargetPointer.Null)
     {
         NativeCodeVersionHandle provisionalHandle = new NativeCodeVersionHandle(methodDescAddress: methodDescAddress, codeVersionNodeAddress: TargetPointer.Null);
@@ -256,7 +324,7 @@ NativeCodeVersionHandle FindActiveNativeCodeVersion(ILCodeVersionHandle methodDe
     MethodDescHandle md = rts.GetMethodDescHandle(methodDescAddress);
     return FindFirstCodeVersion(rts, md, (codeVersion) =>
     {
-        return (!ilVersionId.HasValue || ilVersionId.Value.Value == codeVersion.ILVersionId.Value)
+        return (ilVersionId == codeVersion.ILVersionId)
             && ((NativeCodeVersionNodeFlags)codeVersion.Flags).HasFlag(NativeCodeVersionNodeFlags.IsActiveChild);
     });
 }

docs/design/datacontracts/ReJIT.md

@@ -4,8 +4,22 @@ This contract encapsulates support for [ReJIT](../features/code-versioning.md) i
 
 ## APIs of contract
 
+```csharp
+public enum RejitState
+{
+    Requested,
+    Active
+}
+```
+
 ```csharp
 bool IsEnabled();
+
+RejitState GetRejitState(ILCodeVersionHandle codeVersionHandle);
+
+TargetNUInt GetRejitId(ILCodeVersionHandle codeVersionHandle);
+
+IEnumerable<TargetNUInt> GetRejitIds(TargetPointer methodDesc)
 ```
 
 ## Version 1
@@ -14,6 +28,8 @@ Data descriptors used:
 | Data Descriptor Name | Field | Meaning |
 | --- | --- | --- |
 | ProfControlBlock | GlobalEventMask | an `ICorProfiler` `COR_PRF_MONITOR` value |
+| ILCodeVersionNode | VersionId | `ILCodeVersion` ReJIT id
+| ILCodeVersionNode | RejitState | a `RejitFlags` value |
 
 Global variables used:
 | Global Name | Type | Purpose |
@@ -23,6 +39,7 @@ Global variables used:
 Contracts used:
 | Contract Name |
 | --- |
+| CodeVersions |
 
 ```csharp
 // see src/coreclr/inc/corprof.idl
@@ -32,6 +49,21 @@ private enum COR_PRF_MONITOR
     COR_PRF_ENABLE_REJIT = 0x00040000,
 }
 
+// see src/coreclr/vm/codeversion.h
+[Flags]
+public enum RejitFlags : uint
+{
+    kStateRequested = 0x00000000,
+
+    kStateGettingReJITParameters = 0x00000001,
+
+    kStateActive = 0x00000002,
+
+    kStateMask = 0x0000000F,
+
+    kSuppressParams = 0x80000000
+}
+
 bool IsEnabled()
 {
     TargetPointer address = target.ReadGlobalPointer("ProfilerControlBlock");
@@ -40,4 +72,56 @@ bool IsEnabled()
     bool clrConfigEnabledReJit = /* host process does not have environment variable DOTNET_ProfAPI_ReJitOnAttach set to 0 */;
     return profEnabledReJIT || clrConfigEnabledReJIT;
 }
+
+RejitState GetRejitState(ILCodeVersionHandle codeVersion)
+{
+    // ILCodeVersion::GetRejitState
+    if (codeVersion is not explicit)
+    {
+        // for non explicit ILCodeVersions, ReJITState is always kStateActive
+        return RejitState.Active;
+    }
+    else
+    {
+        // ILCodeVersionNode::GetRejitState
+        ILCodeVersionNode codeVersionNode = AsNode(codeVersion);
+        return ((RejitFlags)ilCodeVersionNode.RejitState & RejitFlags.kStateMask) switch
+        {
+            RejitFlags.kStateRequested => RejitState.Requested,
+            RejitFlags.kStateActive => RejitState.Active,
+            _ => throw new NotImplementedException($"Unknown ReJIT state: {ilCodeVersionNode.RejitState}"),
+        };
+    }
+}
+
+TargetNUInt GetRejitId(ILCodeVersionHandle codeVersion)
+{
+    // ILCodeVersion::GetVersionId
+    if (codeVersion is not explicit)
+    {
+        // for non explicit ILCodeVersions, ReJITId is always 0
+        return new TargetNUInt(0);
+    }
+    else
+    {
+        // ILCodeVersionNode::GetVersionId
+        ILCodeVersionNode codeVersionNode = AsNode(codeVersion);
+        return codeVersionNode.VersionId;
+    }
+}
+
+IEnumerable<TargetNUInt> GetRejitIds(TargetPointer methodDesc)
+{
+    // ReJitManager::GetReJITIDs
+    ICodeVersions cv = _target.Contracts.CodeVersions;
+    IEnumerable<ILCodeVersionHandle> ilCodeVersions = cv.GetILCodeVersions(methodDesc);
+
+    foreach (ILCodeVersionHandle ilCodeVersionHandle in ilCodeVersions)
+    {
+        if (GetRejitState(ilCodeVersionHandle) == RejitState.Active)
+        {
+            yield return GetRejitId(ilCodeVersionHandle);
+        }
+    }
+}
 ```

src/coreclr/debug/runtimeinfo/datadescriptor.h

@@ -483,6 +483,7 @@ CDAC_TYPE_END(CodeHeapListNode)
 
 CDAC_TYPE_BEGIN(ILCodeVersioningState)
 CDAC_TYPE_INDETERMINATE(ILCodeVersioningState)
+CDAC_TYPE_FIELD(ILCodeVersioningState, /*pointer*/, FirstVersionNode, cdac_data<ILCodeVersioningState>::FirstVersionNode)
 CDAC_TYPE_FIELD(ILCodeVersioningState, /*uint32*/, ActiveVersionKind, cdac_data<ILCodeVersioningState>::ActiveVersionKind)
 CDAC_TYPE_FIELD(ILCodeVersioningState, /*pointer*/, ActiveVersionNode, cdac_data<ILCodeVersioningState>::ActiveVersionNode)
 CDAC_TYPE_FIELD(ILCodeVersioningState, /*pointer*/, ActiveVersionModule, cdac_data<ILCodeVersioningState>::ActiveVersionModule)
@@ -501,6 +502,8 @@ CDAC_TYPE_END(NativeCodeVersionNode)
 CDAC_TYPE_BEGIN(ILCodeVersionNode)
 CDAC_TYPE_INDETERMINATE(ILCodeVersionNode)
 CDAC_TYPE_FIELD(ILCodeVersionNode, /*nuint*/, VersionId, cdac_data<ILCodeVersionNode>::VersionId)
+CDAC_TYPE_FIELD(ILCodeVersionNode, /*pointer*/, Next, cdac_data<ILCodeVersionNode>::Next)
+CDAC_TYPE_FIELD(ILCodeVersionNode, /*uint32*/, RejitState, cdac_data<ILCodeVersionNode>::RejitState)
 CDAC_TYPE_END(ILCodeVersionNode)
 
 CDAC_TYPE_BEGIN(ProfControlBlock)

src/coreclr/vm/codeversion.h

@@ -419,6 +419,8 @@ template<>
 struct cdac_data<ILCodeVersionNode>
 {
     static constexpr size_t VersionId = offsetof(ILCodeVersionNode, m_rejitId);
+    static constexpr size_t Next = offsetof(ILCodeVersionNode, m_pNextILVersionNode);
+    static constexpr size_t RejitState = offsetof(ILCodeVersionNode, m_rejitState);
 };
 
 class ILCodeVersionCollection
@@ -543,6 +545,7 @@ class ILCodeVersioningState
 template<>
 struct cdac_data<ILCodeVersioningState>
 {
+    static constexpr size_t FirstVersionNode = offsetof(ILCodeVersioningState, m_pFirstVersionNode);
     static constexpr size_t ActiveVersionKind = offsetof(ILCodeVersioningState, m_activeVersion.m_storageKind);
     static constexpr size_t ActiveVersionNode = offsetof(ILCodeVersioningState, m_activeVersion.m_pVersionNode);
     static constexpr size_t ActiveVersionModule = offsetof(ILCodeVersioningState, m_activeVersion.m_synthetic.m_pModule);

src/native/managed/cdacreader/Microsoft.Diagnostics.DataContractReader.Abstractions/Contracts/Extensions/ICodeVersionsExtensions.cs

@@ -0,0 +1,13 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.Diagnostics.DataContractReader.Contracts.Extensions;
+
+internal static class ICodeVersionsExtensions
+{
+    internal static NativeCodeVersionHandle GetActiveNativeCodeVersion(this ICodeVersions cv, TargetPointer methodDesc)
+    {
+        ILCodeVersionHandle iLCodeVersionHandle = cv.GetActiveILCodeVersion(methodDesc);
+        return cv.GetActiveNativeCodeVersionForILCodeVersion(methodDesc, iLCodeVersionHandle);
+    }
+}

src/native/managed/cdacreader/Microsoft.Diagnostics.DataContractReader.Abstractions/Contracts/Extensions/IReJITExtensions.cs

@@ -0,0 +1,24 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+
+namespace Microsoft.Diagnostics.DataContractReader.Contracts.Extensions;
+
+internal static class IReJITExtensions
+{
+    public static IEnumerable<TargetNUInt> GetRejitIds(this IReJIT rejit, Target target, TargetPointer methodDesc)
+    {
+        ICodeVersions cv = target.Contracts.CodeVersions;
+
+        IEnumerable<ILCodeVersionHandle> ilCodeVersions = cv.GetILCodeVersions(methodDesc);
+
+        foreach (ILCodeVersionHandle ilCodeVersionHandle in ilCodeVersions)
+        {
+            if (rejit.GetRejitState(ilCodeVersionHandle) == RejitState.Active)
+            {
+                yield return rejit.GetRejitId(ilCodeVersionHandle);
+            }
+        }
+    }
+}

src/native/managed/cdacreader/Microsoft.Diagnostics.DataContractReader.Abstractions/Contracts/ICodeVersions.cs

@@ -2,23 +2,51 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
 
 namespace Microsoft.Diagnostics.DataContractReader.Contracts;
 
 internal interface ICodeVersions : IContract
 {
     static string IContract.Name { get; } = nameof(CodeVersions);
 
+    public virtual ILCodeVersionHandle GetActiveILCodeVersion(TargetPointer methodDesc) => throw new NotImplementedException();
+
+    public virtual ILCodeVersionHandle GetILCodeVersion(NativeCodeVersionHandle codeVersionHandle) => throw new NotImplementedException();
+
+    public virtual IEnumerable<ILCodeVersionHandle> GetILCodeVersions(TargetPointer methodDesc) => throw new NotImplementedException();
+
     public virtual NativeCodeVersionHandle GetNativeCodeVersionForIP(TargetCodePointer ip) => throw new NotImplementedException();
-    public virtual NativeCodeVersionHandle GetActiveNativeCodeVersion(TargetPointer methodDesc) => throw new NotImplementedException();
 
-    public virtual bool CodeVersionManagerSupportsMethod(TargetPointer methodDesc) => throw new NotImplementedException();
+    public virtual NativeCodeVersionHandle GetActiveNativeCodeVersionForILCodeVersion(TargetPointer methodDesc, ILCodeVersionHandle ilCodeVersionHandle) => throw new NotImplementedException();
 
     public virtual TargetCodePointer GetNativeCode(NativeCodeVersionHandle codeVersionHandle) => throw new NotImplementedException();
 
+    public virtual bool CodeVersionManagerSupportsMethod(TargetPointer methodDesc) => throw new NotImplementedException();
+}
+
+internal readonly struct ILCodeVersionHandle
+{
+    internal readonly TargetPointer Module;
+    internal readonly uint MethodDefinition;
+    internal readonly TargetPointer ILCodeVersionNode;
+    internal ILCodeVersionHandle(TargetPointer module, uint methodDef, TargetPointer ilCodeVersionNodeAddress)
+    {
+        if (module != TargetPointer.Null && ilCodeVersionNodeAddress != TargetPointer.Null)
+            throw new ArgumentException("Both MethodDesc and ILCodeVersionNode cannot be non-null");
+
+        if (module != TargetPointer.Null && methodDef == 0)
+            throw new ArgumentException("MethodDefinition must be non-zero if Module is non-null");
+
+        Module = module;
+        MethodDefinition = methodDef;
+        ILCodeVersionNode = ilCodeVersionNodeAddress;
+    }
+    public static ILCodeVersionHandle Invalid => new ILCodeVersionHandle(TargetPointer.Null, 0, TargetPointer.Null);
+    public bool IsValid => Module != TargetPointer.Null || ILCodeVersionNode != TargetPointer.Null;
 }
 
-internal struct NativeCodeVersionHandle
+internal readonly struct NativeCodeVersionHandle
 {
     // no public constructors
     internal readonly TargetPointer MethodDescAddress;

src/native/managed/cdacreader/Microsoft.Diagnostics.DataContractReader.Abstractions/Contracts/IReJIT.cs

@@ -2,13 +2,25 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
 
 namespace Microsoft.Diagnostics.DataContractReader.Contracts;
 
+public enum RejitState
+{
+    Requested,
+    Active
+}
+
 internal interface IReJIT : IContract
 {
     static string IContract.Name { get; } = nameof(ReJIT);
+
     bool IsEnabled() => throw new NotImplementedException();
+
+    RejitState GetRejitState(ILCodeVersionHandle codeVersionHandle) => throw new NotImplementedException();
+
+    TargetNUInt GetRejitId(ILCodeVersionHandle codeVersionHandle) => throw new NotImplementedException();
 }
 
 internal readonly struct ReJIT : IReJIT