Fix SPMI to handle replays of BBINSTR jit method contexts

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