Fix SPMI to handle replays of BBINSTR jit method contexts (#41386)

Move the handling of allocMethodBlockCounts to the method context.
While this data is both an input and an output, we're more interested
in the input side.

We could also treat it as an output, if we say wanted to verify that
the replay jit wrote the same sequence of IL offsets into the BlockCount
records, but that doesn't seem crucial as changes here would also lead
to code diffs.

Fix the code that checks the AddressMap for reloc hints so it doesn't
fail if the exact lookup fails, but instead falls back to a search.
This is needed because the base of the replay count buffer is recorded
as the map key, but the instrumentation probes in jitted code refer
to offsets from this value.

Fixes #37270.

src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.cpp

@@ -895,9 +895,17 @@ void* CompileResult::repAddressMap(void* replayAddress)
 {
     if (AddressMap == nullptr)
         return nullptr;
-    Agnostic_AddressMap value;
-    value = AddressMap->Get((DWORDLONG)replayAddress);
-    return (void*)value.Address;
+
+    int index = AddressMap->GetIndex((DWORDLONG)replayAddress);
+
+    if (index != -1)
+    {
+        Agnostic_AddressMap value;
+        value = AddressMap->Get((DWORDLONG)replayAddress);
+        return (void*)value.Address;
+    }
+
+    return nullptr;
 }
 void* CompileResult::searchAddressMap(void* newAddress)
 {
@@ -962,39 +970,6 @@ void CompileResult::dmpAllocUnwindInfo(DWORD key, const Agnostic_AllocUnwindInfo
            value.pUnwindBlock_index, value.funcKind);
 }
 
-void CompileResult::recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result)
-{
-    if (AllocMethodBlockCounts == nullptr)
-        AllocMethodBlockCounts = new LightWeightMap<DWORD, Agnostic_AllocMethodBlockCounts>();
-
-    Agnostic_AllocMethodBlockCounts value;
-
-    value.count  = (DWORD)count;
-    value.result = (DWORD)result;
-    value.pBlockCounts_index =
-        AllocMethodBlockCounts->AddBuffer((unsigned char*)*pBlockCounts, count * sizeof(ICorJitInfo::BlockCounts));
-
-    AllocMethodBlockCounts->Add((DWORD)0, value);
-}
-void CompileResult::dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value)
-{
-    printf("AllocMethodBlockCounts key %u, value cnt-%u ind-%u res-%08X", key, value.count, value.pBlockCounts_index,
-           value.result);
-}
-HRESULT CompileResult::repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts)
-{
-    Agnostic_AllocMethodBlockCounts value;
-    value = AllocMethodBlockCounts->Get((DWORD)0);
-
-    if (count != value.count)
-        __debugbreak();
-
-    HRESULT result = (HRESULT)value.result;
-    *pBlockCounts = (ICorJitInfo::BlockCounts*)AllocMethodBlockCounts->GetBuffer(value.pBlockCounts_index);
-    recAddressMap((void*)0x4242, (void*)*pBlockCounts, count * (sizeof(ICorJitInfo::BlockCounts)));
-    return result;
-}
-
 void CompileResult::recRecordCallSite(ULONG instrOffset, CORINFO_SIG_INFO* callSig, CORINFO_METHOD_HANDLE methodHandle)
 {
     repRecordCallSite(instrOffset, callSig, methodHandle);

src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.h

@@ -166,12 +166,6 @@ class CompileResult
         DWORD HandlerLength;
         DWORD ClassToken; // one view of symetric union
     };
-    struct Agnostic_AllocMethodBlockCounts
-    {
-        DWORD count;
-        DWORD pBlockCounts_index;
-        DWORD result;
-    };
     struct Agnostic_CORINFO_SIG_INFO2
     {
         DWORD     callConv;
@@ -328,10 +322,6 @@ class CompileResult
                             CorJitFuncKind funcKind);
     void dmpAllocUnwindInfo(DWORD key, const Agnostic_AllocUnwindInfo& value);
 
-    void recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result);
-    void dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value);
-    HRESULT repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts);
-
     void recRecordCallSite(ULONG instrOffset, CORINFO_SIG_INFO* callSig, CORINFO_METHOD_HANDLE methodHandle);
     void dmpRecordCallSite(DWORD key, const Agnostic_RecordCallSite& value);
     void repRecordCallSite(ULONG instrOffset, CORINFO_SIG_INFO* callSig, CORINFO_METHOD_HANDLE methodHandle);

src/coreclr/src/ToolBox/superpmi/superpmi-shared/crlwmlist.h

@@ -19,7 +19,6 @@
 #endif
 
 LWM(AddressMap, DWORDLONG, CompileResult::Agnostic_AddressMap)
-LWM(AllocMethodBlockCounts, DWORD, CompileResult::Agnostic_AllocMethodBlockCounts)
 LWM(AllocGCInfo, DWORD, CompileResult::Agnostic_AllocGCInfo)
 LWM(AllocMem, DWORD, CompileResult::Agnostic_AllocMemDetails)
 DENSELWM(AllocUnwindInfo, CompileResult::Agnostic_AllocUnwindInfo)

src/coreclr/src/ToolBox/superpmi/superpmi-shared/lightweightmap.h

@@ -78,6 +78,33 @@ class LightWeightMapBuffer
         return newOffset + sizeof(unsigned int);
     }
 
+    const unsigned char* CreateBuffer(unsigned int len)
+    {
+        if (len == 0)
+        {
+            return nullptr;
+        }
+
+        if (locked)
+        {
+            LogError("Added item that extended the buffer after it was locked by a call to GetBuffer()");
+            __debugbreak();
+        }
+
+        unsigned int   newbuffsize = bufferLength + sizeof(unsigned int) + len;
+        unsigned char* newbuffer   = new unsigned char[newbuffsize];
+        unsigned int   newOffset   = bufferLength;
+        if (bufferLength > 0)
+            memcpy(newbuffer, buffer, bufferLength);
+        memset(newbuffer + bufferLength + sizeof(unsigned int), 0, len);
+        *((unsigned int*)(newbuffer + bufferLength)) = len;
+        bufferLength += sizeof(unsigned int) + len;
+        if (buffer != nullptr)
+            delete[] buffer;
+        buffer = newbuffer;
+        return buffer + newOffset + sizeof(unsigned int);
+    }
+
     unsigned char* GetBuffer(unsigned int offset)
     {
         if (offset == (unsigned int)-1)

src/coreclr/src/ToolBox/superpmi/superpmi-shared/lwmlist.h

@@ -18,6 +18,7 @@
 #define DENSELWM(map, value) LWM(map, this_is_an_error, value)
 #endif
 
+LWM(AllocMethodBlockCounts, DWORD, Agnostic_AllocMethodBlockCounts)
 LWM(AppendClassName, Agnostic_AppendClassName, DWORD)
 LWM(AreTypesEquivalent, DLDL, DWORD)
 LWM(AsCorInfoType, DWORDLONG, DWORD)

src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.cpp

@@ -5133,6 +5133,50 @@ DWORD MethodContext::repGetFieldThreadLocalStoreID(CORINFO_FIELD_HANDLE field, v
     return (DWORD)value.B;
 }
 
+
+void MethodContext::recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result)
+{
+    if (AllocMethodBlockCounts == nullptr)
+        AllocMethodBlockCounts = new LightWeightMap<DWORD, Agnostic_AllocMethodBlockCounts>();
+
+    Agnostic_AllocMethodBlockCounts value;
+
+    value.address = (DWORDLONG)*pBlockCounts;
+    value.count  = (DWORD)count;
+    value.result = (DWORD)result;
+
+    AllocMethodBlockCounts->Add((DWORD)0, value);
+}
+void MethodContext::dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value)
+{
+    printf("AllocMethodBlockCounts key %u, value addr-%016llX cnt-%u res-%08X", key, value.address, value.count, value.result);
+}
+HRESULT MethodContext::repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts)
+{
+    Agnostic_AllocMethodBlockCounts value;
+    value = AllocMethodBlockCounts->Get((DWORD)0);
+
+    if (count != value.count)
+    {
+        LogWarning("AllocMethodBlockCount mismatch: record %d, replay %d", value.count, count);
+    }
+
+    HRESULT result = (HRESULT)value.result;
+
+    // Allocate a scratch buffer, linked to method context via AllocMethodBlockCounts, so it gets
+    // cleaned up when the method context does.
+    //
+    // We won't bother recording this via AddBuffer because currently SPMI will never look at it.
+    // But we need a writeable buffer because the jit will store IL offsets inside.
+    //
+    // Todo, perhaps: record the buffer as a compile result instead, and defer copying until
+    // jit completion so we can snapshot the offsets the jit writes.
+    //
+    *pBlockCounts = (ICorJitInfo::BlockCounts*)AllocMethodBlockCounts->CreateBuffer(count * sizeof(ICorJitInfo::BlockCounts));
+    cr->recAddressMap((void*)value.address, (void*)*pBlockCounts, count * (sizeof(ICorJitInfo::BlockCounts)));
+    return result;
+}
+
 void MethodContext::recGetMethodBlockCounts(CORINFO_METHOD_HANDLE        ftnHnd,
                                             UINT32 *                     pCount,
                                             ICorJitInfo::BlockCounts**   pBlockCounts,

src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.h

@@ -422,6 +422,12 @@ class MethodContext
         DWORDLONG method;
         DWORDLONG delegateCls;
     };
+    struct Agnostic_AllocMethodBlockCounts
+    {
+        DWORDLONG address;
+        DWORD count;
+        DWORD result;
+    };
     struct Agnostic_GetMethodBlockCounts
     {
         DWORD count;
@@ -1168,6 +1174,10 @@ class MethodContext
     void dmpGetFieldThreadLocalStoreID(DWORDLONG key, DLD value);
     DWORD repGetFieldThreadLocalStoreID(CORINFO_FIELD_HANDLE field, void** ppIndirection);
 
+    void recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result);
+    void dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value);
+    HRESULT repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts);
+
     void recGetMethodBlockCounts(CORINFO_METHOD_HANDLE        ftnHnd,
                                  UINT32 *                     pCount,
                                  ICorJitInfo::BlockCounts**   pBlockCounts,
@@ -1338,6 +1348,7 @@ class MethodContext
 // *************************************************************************************
 enum mcPackets
 {
+    Packet_AllocMethodBlockCounts     = 131,
     Packet_AppendClassName            = 149, // Added 8/6/2014 - needed for SIMD
     Packet_AreTypesEquivalent         = 1,
     Packet_AsCorInfoType              = 2,
@@ -1493,7 +1504,6 @@ enum mcPackets
     Packet_ShouldEnforceCallvirtRestriction              = 112, // Retired 2/18/2020
 
     PacketCR_AddressMap                        = 113,
-    PacketCR_AllocMethodBlockCounts            = 131,
     PacketCR_AllocGCInfo                       = 114,
     PacketCR_AllocMem                          = 115,
     PacketCR_AllocUnwindInfo                   = 132,

src/coreclr/src/ToolBox/superpmi/superpmi-shim-collector/icorjitinfo.cpp

@@ -2030,7 +2030,7 @@ HRESULT interceptor_ICJI::allocMethodBlockCounts(UINT32          count, // The n
 {
     mc->cr->AddCall("allocMethodBlockCounts");
     HRESULT result = original_ICorJitInfo->allocMethodBlockCounts(count, pBlockCounts);
-    mc->cr->recAllocMethodBlockCounts(count, pBlockCounts, result);
+    mc->recAllocMethodBlockCounts(count, pBlockCounts, result);
     return result;
 }
 

src/coreclr/src/ToolBox/superpmi/superpmi/icorjitinfo.cpp

@@ -1784,7 +1784,7 @@ HRESULT MyICJI::allocMethodBlockCounts(UINT32          count, // The number of b
                                        BlockCounts**   pBlockCounts)
 {
     jitInstance->mc->cr->AddCall("allocMethodBlockCounts");
-    return jitInstance->mc->cr->repAllocMethodBlockCounts(count, pBlockCounts);
+    return jitInstance->mc->repAllocMethodBlockCounts(count, pBlockCounts);
 }
 
 // get profile information to be used for optimizing the current method.  The format