Transform indirect calls to direct ones in the importer

Imports indirect calls as direct ones when the target method is known.

Only handles addresses from ldftn as the VM has no way to verify pointers
from static readonly fields and crashes on invalid ones.

The helpers currently contain a small dead path that I'll soon use when
extending the code to also handle delegates.

Opening as a draft so that the code can be reviewed while I finish the tests.

Fixes dotnet#44610

Author: MichalPetryka

src/coreclr/jit/compiler.h

@@ -3830,6 +3830,15 @@ class Compiler
     bool impCanPInvokeInlineCallSite(BasicBlock* block);
     void impCheckForPInvokeCall(
         GenTreeCall* call, CORINFO_METHOD_HANDLE methHnd, CORINFO_SIG_INFO* sig, unsigned mflags, BasicBlock* block);
+
+    enum SigTransform
+    {
+        LeaveIntact    = 0,
+        DeleteThis     = 1 << 0,
+        ReplaceRefThis = 1 << 1,
+    };
+
+    bool impCanSubstituteSig(CORINFO_SIG_INFO* sourceSig, CORINFO_SIG_INFO* targetSig, SigTransform transformation);
     GenTreeCall* impImportIndirectCall(CORINFO_SIG_INFO* sig, const DebugInfo& di = DebugInfo());
     void impPopArgsForUnmanagedCall(GenTreeCall* call, CORINFO_SIG_INFO* sig);
 
@@ -3858,6 +3867,7 @@ class Compiler
 
     GenTree* impInitClass(CORINFO_RESOLVED_TOKEN* pResolvedToken);
 
+    GenTree* impGetNodeFromLocal(GenTree* node);
     GenTree* impImportStaticReadOnlyField(CORINFO_FIELD_HANDLE field, CORINFO_CLASS_HANDLE ownerCls);
     GenTree* impImportCnsTreeFromBuffer(uint8_t* buffer, var_types valueType);
 

src/coreclr/jit/importer.cpp

@@ -3764,6 +3764,65 @@ GenTree* Compiler::impInitClass(CORINFO_RESOLVED_TOKEN* pResolvedToken)
     return node;
 }
 
+//------------------------------------------------------------------------
+// impGetNodeFromLocal: Tries to return the node that's assigned
+// to the provided local.
+//
+// Arguments:
+//    node - GT_LCL_VAR whose value is searched for
+//
+// Return Value:
+//    The tree representing the node assigned to the variable when possible,
+//    nullptr otherwise.
+//
+GenTree* Compiler::impGetNodeFromLocal(GenTree* node)
+{
+    assert(node != nullptr);
+    assert(node->OperIs(GT_LCL_VAR));
+
+    unsigned lclNum = node->AsLclVarCommon()->GetLclNum();
+
+    if (lvaTable[lclNum].lvSingleDef == 0)
+    {
+        return nullptr;
+    }
+
+    auto findValue = [&](Statement* stmtList) -> GenTree* {
+        for (Statement* stmt : StatementList(stmtList))
+        {
+            GenTree* root = stmt->GetRootNode();
+            if (root->OperIs(GT_ASG) && root->AsOp()->gtOp1->OperIs(GT_LCL_VAR) &&
+                root->AsOp()->gtOp1->AsLclVarCommon()->GetLclNum() == lclNum)
+            {
+                return root->AsOp()->gtOp2;
+            }
+        }
+        return nullptr;
+    };
+
+    GenTree* valueNode = findValue(impStmtList);
+    BasicBlock* bb = fgFirstBB;
+    while (valueNode == nullptr && bb != nullptr)
+    {
+        valueNode = findValue(bb->bbStmtList);
+        if (valueNode == nullptr && bb->NumSucc(this) == 1)
+        {
+            bb = bb->GetSucc(0, this);
+        }
+        else
+        {
+            bb = nullptr;
+        }
+    }
+
+    if (valueNode != nullptr && valueNode->OperIs(GT_LCL_VAR))
+    {
+        return impGetNodeFromLocal(valueNode);
+    }
+
+    return valueNode;
+}
+
 //------------------------------------------------------------------------
 // impImportStaticReadOnlyField: Tries to import 'static readonly' field
 //    as a constant if the host type is statically initialized.

src/coreclr/jit/importercalls.cpp

@@ -104,13 +104,16 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
     bool checkForSmallType  = false;
     bool bIntrinsicImported = false;
 
-    CORINFO_SIG_INFO calliSig;
+    CORINFO_SIG_INFO originalSig;
     NewCallArg       extraArg;
 
-    /*-------------------------------------------------------------------------
-     * First create the call node
-     */
+    // run transformations when instrumenting to not pollute PGO data
+    bool optimizedOrInstrumented = opts.OptimizationEnabled() || opts.IsInstrumented();
+    CORINFO_METHOD_HANDLE replacementMethod = nullptr;
+    GenTree* newThis = nullptr;
+    SigTransform sigTransformation = SigTransform::LeaveIntact;
 
+    // handle special import cases
     if (opcode == CEE_CALLI)
     {
         if (IsTargetAbi(CORINFO_NATIVEAOT_ABI))
@@ -125,25 +128,102 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
         }
 
         /* Get the call site sig */
-        eeGetSig(pResolvedToken->token, pResolvedToken->tokenScope, pResolvedToken->tokenContext, &calliSig);
+        eeGetSig(pResolvedToken->token, pResolvedToken->tokenScope, pResolvedToken->tokenContext, &originalSig);
 
-        callRetTyp = JITtype2varType(calliSig.retType);
+        if (!optimizedOrInstrumented)
+        {
+            // ignore
+        }
+        else if (originalSig.getCallConv() == CORINFO_CALLCONV_DEFAULT)
+        {
+            JITDUMP("\n\nimpImportCall trying to import calli as call\n");
+            GenTree* fptr = impStackTop().val;
+            if (fptr->OperIs(GT_LCL_VAR))
+            {
+                JITDUMP("impImportCall trying to import calli as call - trying to substitute LCL_VAR\n");
+                GenTree* lclValue = impGetNodeFromLocal(fptr);
+                if (lclValue != nullptr)
+                {
+                    fptr = lclValue;
+                }
+            }
+            if (fptr->OperIs(GT_FTN_ADDR))
+            {
+                replacementMethod = fptr->AsFptrVal()->gtFptrMethod;
+            }
+            else
+            {
+                JITDUMP("impImportCall failed to import calli as call - address node not found\n");
+            }
+        }
+        else
+        {
+            JITDUMP("impImportCall failed to import calli as call - call conv %u is not managed\n",
+                originalSig.getCallConv());
+        }
+    }
+
+    if (replacementMethod != nullptr)
+    {
+        JITDUMP("impImportCall trying to transform call - found target method %s\n",
+            eeGetMethodName(replacementMethod));
+        CORINFO_SIG_INFO methodSig;
+        CORINFO_CLASS_HANDLE targetClass = info.compCompHnd->getMethodClass(replacementMethod);
+        info.compCompHnd->getMethodSig(replacementMethod, &methodSig, targetClass);
 
-        call = impImportIndirectCall(&calliSig, di);
+        unsigned replacementFlags = info.compCompHnd->getMethodAttribs(replacementMethod);
+
+        if ((replacementFlags & CORINFO_FLG_PINVOKE) != 0)
+        {
+            JITDUMP("impImportCall aborting transformation - found PInvoke\n");
+        }
+        else if (impCanSubstituteSig(&originalSig, &methodSig, sigTransformation))
+        {
+            impPopStack();
+            if (newThis != nullptr)
+            {
+                assert(sigTransformation == SigTransform::ReplaceRefThis);
+                CORINFO_CLASS_HANDLE thisCls = NO_CLASS_HANDLE;
+                info.compCompHnd->getArgType(&methodSig, methodSig.args, &thisCls);
+                impPushOnStack(newThis, typeInfo(TI_REF, thisCls));
+            }
+            JITDUMP("impImportCall transforming call\n");
+            pResolvedToken->hMethod = replacementMethod;
+            pResolvedToken->hClass = targetClass;
+
+            callInfo->sig = methodSig;
+            callInfo->hMethod = replacementMethod;
+            callInfo->methodFlags = replacementFlags;
+            callInfo->classFlags = info.compCompHnd->getClassAttribs(targetClass);
+
+            return impImportCall(CEE_CALL, pResolvedToken, nullptr, nullptr,
+                prefixFlags, callInfo, rawILOffset);
+        }
+    }
+
+    /*-------------------------------------------------------------------------
+     * First create the call node
+     */
+
+    if (opcode == CEE_CALLI)
+    {
+        callRetTyp = JITtype2varType(originalSig.retType);
+
+        call = impImportIndirectCall(&originalSig, di);
 
         // We don't know the target method, so we have to infer the flags, or
         // assume the worst-case.
-        mflags = (calliSig.callConv & CORINFO_CALLCONV_HASTHIS) ? 0 : CORINFO_FLG_STATIC;
+        mflags = (originalSig.callConv & CORINFO_CALLCONV_HASTHIS) ? 0 : CORINFO_FLG_STATIC;
 
 #ifdef DEBUG
         if (verbose)
         {
-            unsigned structSize = (callRetTyp == TYP_STRUCT) ? eeTryGetClassSize(calliSig.retTypeSigClass) : 0;
+            unsigned structSize = (callRetTyp == TYP_STRUCT) ? eeTryGetClassSize(originalSig.retTypeSigClass) : 0;
             printf("\nIn Compiler::impImportCall: opcode is %s, kind=%d, callRetType is %s, structSize is %u\n",
                    opcodeNames[opcode], callInfo->kind, varTypeName(callRetTyp), structSize);
         }
 #endif
-        sig = &calliSig;
+        sig = &originalSig;
     }
     else // (opcode != CEE_CALLI)
     {
@@ -1622,6 +1702,114 @@ GenTree* Compiler::impFixupCallStructReturn(GenTreeCall* call, CORINFO_CLASS_HAN
 #endif // FEATURE_MULTIREG_RET
 }
 
+//-----------------------------------------------------------------------------------
+//  impCanSubstituteSig: Checks whether it's safe to replace a call with another one.
+//  This DOES NOT check if the calls are actually compatible, it only checks if their trees are.
+//  Use ONLY when code will call the method with target sig anyway.
+//
+//  Arguments:
+//    sourceSig      - original call signature
+//    targetSig      - new call signature
+//    transformation - transformations performed on the original signature
+//
+//  Return Value:
+//    Whether it's safe to change the IR to call the target method
+//
+bool Compiler::impCanSubstituteSig(CORINFO_SIG_INFO* sourceSig, CORINFO_SIG_INFO* targetSig, SigTransform transformation)
+{
+    const SigTransform thisChangeMask = (SigTransform)(SigTransform::DeleteThis | SigTransform::ReplaceRefThis);
+    assert((transformation & thisChangeMask) != thisChangeMask);
+
+    if (sourceSig->getCallConv() != targetSig->getCallConv())
+    {
+        JITDUMP("impCanSubstituteSig returning false - call conv %u != %u\n", sourceSig->callConv, targetSig->callConv);
+        return false;
+    }
+
+    unsigned sourceArgCount = sourceSig->numArgs;
+    if ((transformation & SigTransform::DeleteThis) != 0)
+    {
+        sourceArgCount--;
+    }
+
+    if (sourceArgCount != targetSig->numArgs)
+    {
+        JITDUMP("impCanSubstituteSig returning false - args count %u != %u\n", sourceArgCount, targetSig->numArgs);
+        return false;
+    }
+
+    if (sourceSig->retType != targetSig->retType)
+    {
+        JITDUMP("impCanSubstituteSig returning false - return type %u != %u\n",
+            (unsigned)sourceSig->retType, (unsigned)targetSig->retType);
+        return false;
+    }
+
+    if (sourceSig->retType == CORINFO_TYPE_VALUECLASS || sourceSig->retType == CORINFO_TYPE_REFANY)
+    {
+        ClassLayout* layoutRetA = typGetObjLayout(sourceSig->retTypeClass);
+        ClassLayout* layoutRetB = typGetObjLayout(targetSig->retTypeClass);
+
+        if (!ClassLayout::AreCompatible(layoutRetA, layoutRetB))
+        {
+            JITDUMP("impCanSubstituteSig returning false - return type %u is incompatible with %u\n",
+                (unsigned)sourceSig->retType, (unsigned)targetSig->retType);
+            return false;
+        }
+    }
+
+    CORINFO_ARG_LIST_HANDLE sourceArg = sourceSig->args;
+    CORINFO_ARG_LIST_HANDLE targetArg = targetSig->args;
+
+    assert((transformation & (SigTransform::DeleteThis | SigTransform::ReplaceRefThis)) == 0 ||
+        eeGetArgType(sourceArg, sourceSig) == TYP_REF);
+
+    if ((transformation & SigTransform::DeleteThis) != 0)
+    {
+        sourceArg = info.compCompHnd->getArgNext(sourceArg);
+    }
+
+    if ((transformation & SigTransform::ReplaceRefThis) != 0 && eeGetArgType(targetArg, targetSig) != TYP_REF)
+    {
+        JITDUMP("impCanSubstituteSig returning false - this is not TYP_REF\n");
+        return false;
+    }
+
+    for (unsigned i = 0; i < targetSig->numArgs; i++,
+        sourceArg = info.compCompHnd->getArgNext(sourceArg),
+        targetArg = info.compCompHnd->getArgNext(targetArg))
+    {
+        var_types sourceType = eeGetArgType(sourceArg, sourceSig);
+        var_types targetType = eeGetArgType(targetArg, targetSig);
+        if (sourceType != targetType)
+        {
+            JITDUMP("impCanSubstituteSig returning false - parameter %u type %s != %s\n",
+                i, varTypeName(sourceType), varTypeName(targetType));
+            return false;
+        }
+
+        if (varTypeIsStruct(sourceType) && varTypeIsStruct(targetType))
+        {
+            CORINFO_CLASS_HANDLE sourceClassHnd = NO_CLASS_HANDLE;
+            CORINFO_CLASS_HANDLE targetClassHnd = NO_CLASS_HANDLE;
+            info.compCompHnd->getArgType(sourceSig, sourceArg, &sourceClassHnd);
+            info.compCompHnd->getArgType(targetSig, targetArg, &targetClassHnd);
+
+            ClassLayout* sourceLayout = typGetObjLayout(sourceClassHnd);
+            ClassLayout* targetLayout = typGetObjLayout(targetClassHnd);
+
+            if (!ClassLayout::AreCompatible(sourceLayout, targetLayout))
+            {
+                JITDUMP("impCanSubstituteSig returning false - parameter %u type %s is inconmpatible with %s\n",
+                    i, varTypeName(sourceType), varTypeName(targetType));
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
 GenTreeCall* Compiler::impImportIndirectCall(CORINFO_SIG_INFO* sig, const DebugInfo& di)
 {
     var_types callRetTyp = JITtype2varType(sig->retType);