allow async interruptions on safepoints

src/coreclr/gcinfo/gcinfoencoder.cpp

@@ -943,6 +943,8 @@ bool GcInfoEncoder::IsAlwaysScratch(GcSlotDesc &slotDesc)
     else
         return FALSE;
 
+    // TODO: VS add ARM64
+
 #elif defined(TARGET_AMD64)
 
     _ASSERTE( m_SizeOfStackOutgoingAndScratchArea != (UINT32)-1 );
@@ -964,6 +966,8 @@ bool GcInfoEncoder::IsAlwaysScratch(GcSlotDesc &slotDesc)
             | (1 << 14)  // r14
             | (1 << 15); // r15
 
+        PreservedRegMask |= 1; // rax - may contain return value
+
         return !(PreservedRegMask & (1 << regNum));
     }
     else if (!slotDesc.IsUntracked() && (slotDesc.Slot.Stack.Base == GC_SP_REL) &&

src/coreclr/inc/gcinfodecoder.h

@@ -504,6 +504,8 @@ class GcInfoDecoder
     bool IsInterruptible();
 
 #ifdef PARTIALLY_INTERRUPTIBLE_GC_SUPPORTED
+    bool IsSafePoint();
+
     // This is used for gccoverage
     bool IsSafePoint(UINT32 codeOffset);
 

src/coreclr/jit/codegenlinear.cpp

@@ -714,8 +714,7 @@ void CodeGen::genCodeForBBlist()
                 }
                 // Do likewise for blocks that end in DOES_NOT_RETURN calls
                 // that were not caught by the above rules. This ensures that
-                // gc register liveness doesn't change across call instructions
-                // in fully-interruptible mode.
+                // gc register liveness doesn't change to some random state after call instructions
                 else
                 {
                     GenTree* call = block->lastNode();

src/coreclr/nativeaot/Runtime/unix/UnixNativeCodeManager.cpp

@@ -193,7 +193,13 @@ bool UnixNativeCodeManager::IsSafePoint(PTR_VOID pvAddress)
         codeOffset
     );
 
-    return decoder.IsInterruptible();
+    if (decoder.IsInterruptible())
+        return true;
+
+    if (decoder.IsSafePoint())
+        return true;
+
+    return false;
 }
 
 void UnixNativeCodeManager::EnumGcRefs(MethodInfo *    pMethodInfo,
@@ -219,6 +225,21 @@ void UnixNativeCodeManager::EnumGcRefs(MethodInfo *    pMethodInfo,
         // the reasons for this adjustment are explained in EECodeManager::EnumGcRefs
         codeOffset--;
     }
+    else
+    {
+        // CONSIDER: We can optimize this by remembering the need to adjust in IsSafePoint and propagating into here.
+        //           Or, better yet, maybe we should change the decoder to not require this adjustment.
+        //           The scenario that adjustment tries to handle (fallthrough into BB with random liveness)
+        //           does not seem possible.
+        GcInfoDecoder decoder1(
+            GCInfoToken(gcInfo),
+            GcInfoDecoderFlags(DECODE_INTERRUPTIBILITY),
+            codeOffset
+        );
+
+        if (decoder1.IsSafePoint())
+            codeOffset--;
+    }
 
     GcInfoDecoder decoder(
         GCInfoToken(gcInfo),

src/coreclr/nativeaot/Runtime/windows/CoffNativeCodeManager.cpp

@@ -415,7 +415,13 @@ bool CoffNativeCodeManager::IsSafePoint(PTR_VOID pvAddress)
         codeOffset
     );
 
-    return decoder.IsInterruptible();
+    if (decoder.IsInterruptible())
+        return true;
+
+    if (decoder.IsSafePoint())
+        return true;
+
+    return false;
 #else
     // Extract the necessary information from the info block header
     hdrInfo info;
@@ -452,12 +458,27 @@ void CoffNativeCodeManager::EnumGcRefs(MethodInfo *    pMethodInfo,
         // the reasons for this adjustment are explained in EECodeManager::EnumGcRefs
         codeOffset--;
     }
+    else
+    {
+        // CONSIDER: We can optimize this by remembering the need to adjust in IsSafePoint and propagating into here.
+        //           Or, better yet, maybe we should change the decoder to not require this adjustment.
+        //           The scenario that adjustment tries to handle (fallthrough into BB with random liveness)
+        //           does not seem possible.
+        GcInfoDecoder decoder1(
+            GCInfoToken(gcInfo),
+            GcInfoDecoderFlags(DECODE_INTERRUPTIBILITY),
+            codeOffset
+        );
+
+        if (decoder1.IsSafePoint())
+            codeOffset--;
+    }
 
     GcInfoDecoder decoder(
         GCInfoToken(gcInfo),
         GcInfoDecoderFlags(DECODE_GC_LIFETIMES | DECODE_SECURITY_OBJECT | DECODE_VARARG),
         codeOffset
-        );
+    );
 
     if (!decoder.EnumerateLiveSlots(
         pRegisterSet,

src/coreclr/vm/gcinfodecoder.cpp

@@ -367,7 +367,11 @@ GcInfoDecoder::GcInfoDecoder(
     {
         if(m_NumSafePoints)
         {
-            m_SafePointIndex = FindSafePoint(m_InstructionOffset);
+            // Safepoints are encoded with a -1 adjustment
+            // DECODE_GC_LIFETIMES adjusts the offset accordingly, but DECODE_INTERRUPTIBILITY does not
+            // adjust here
+            UINT32 offset = flags & DECODE_INTERRUPTIBILITY ? m_InstructionOffset - 1 : m_InstructionOffset;
+            m_SafePointIndex = FindSafePoint(offset);
         }
     }
     else if(flags & (DECODE_FOR_RANGES_CALLBACK | DECODE_INTERRUPTIBILITY))
@@ -393,6 +397,12 @@ bool GcInfoDecoder::IsInterruptible()
     return m_IsInterruptible;
 }
 
+bool GcInfoDecoder::IsSafePoint()
+{
+    _ASSERTE(m_Flags & (DECODE_INTERRUPTIBILITY | DECODE_GC_LIFETIMES));
+    return m_SafePointIndex != m_NumSafePoints;
+}
+
 bool GcInfoDecoder::HasMethodDescGenericsInstContext()
 {
     _ASSERTE( m_Flags & DECODE_GENERICS_INST_CONTEXT );