From 6aa6e40804bb4c42f33e193d9cd20187c77b7b0b Mon Sep 17 00:00:00 2001
From: EgorBo <egorbo@gmail.com>
Date: Thu, 25 Jan 2024 22:36:24 +0100
Subject: [PATCH] Disable late casts

---
 src/coreclr/jit/compiler.cpp |  3 --
 src/coreclr/jit/importer.cpp | 77 +++++++++++++++++++++++++++++++++---
 2 files changed, 71 insertions(+), 9 deletions(-)

diff --git a/src/coreclr/jit/compiler.cpp b/src/coreclr/jit/compiler.cpp
index bd8cd590eea537..9726894719b89d 100644
--- a/src/coreclr/jit/compiler.cpp
+++ b/src/coreclr/jit/compiler.cpp
@@ -5050,9 +5050,6 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     // Expand thread local access
     DoPhase(this, PHASE_EXPAND_TLS, &Compiler::fgExpandThreadLocalAccess);
 
-    // Expand casts
-    DoPhase(this, PHASE_EXPAND_CASTS, &Compiler::fgLateCastExpansion);
-
     // Insert GC Polls
     DoPhase(this, PHASE_INSERT_GC_POLLS, &Compiler::fgInsertGCPolls);
 
diff --git a/src/coreclr/jit/importer.cpp b/src/coreclr/jit/importer.cpp
index 701e0463367b89..a0981b548ecdf4 100644
--- a/src/coreclr/jit/importer.cpp
+++ b/src/coreclr/jit/importer.cpp
@@ -5518,6 +5518,77 @@ GenTree* Compiler::impCastClassOrIsInstToTree(
             // If the class is exact, the jit can expand the IsInst check inline.
             canExpandInline = isClassExact;
         }
+
+        // Check if this cast helper have some profile data
+        if (impIsCastHelperMayHaveProfileData(helper))
+        {
+            const int               maxLikelyClasses = 32;
+            LikelyClassMethodRecord likelyClasses[maxLikelyClasses];
+            unsigned                likelyClassCount =
+                getLikelyClasses(likelyClasses, maxLikelyClasses, fgPgoSchema, fgPgoSchemaCount, fgPgoData, ilOffset);
+
+            if (likelyClassCount > 0)
+            {
+#ifdef DEBUG
+                for (UINT32 i = 0; i < likelyClassCount; i++)
+                {
+                    const char* className = eeGetClassName((CORINFO_CLASS_HANDLE)likelyClasses[i].handle);
+                    JITDUMP("  %u) %p (%s) [likelihood:%u%%]\n", i + 1, likelyClasses[i].handle, className,
+                            likelyClasses[i].likelihood);
+                }
+
+                // Optional stress mode to pick a random known class, rather than
+                // the most likely known class.
+                if (JitConfig.JitRandomGuardedDevirtualization() != 0)
+                {
+                    // Reuse the random inliner's random state.
+                    CLRRandom* const random =
+                        impInlineRoot()->m_inlineStrategy->GetRandom(JitConfig.JitRandomGuardedDevirtualization());
+
+                    unsigned index          = static_cast<unsigned>(random->Next(static_cast<int>(likelyClassCount)));
+                    likelyClasses[0].handle = likelyClasses[index].handle;
+                    likelyClasses[0].likelihood = 100;
+                    likelyClassCount            = 1;
+                }
+#endif
+
+                LikelyClassMethodRecord likelyClass = likelyClasses[0];
+                CORINFO_CLASS_HANDLE    likelyCls   = (CORINFO_CLASS_HANDLE)likelyClass.handle;
+
+                // if there is a dominating candidate with >= 40% likelihood, use it
+                const unsigned likelihoodMinThreshold = 40;
+                if ((likelyCls != NO_CLASS_HANDLE) && (likelyClass.likelihood > likelihoodMinThreshold))
+                {
+                    TypeCompareState castResult =
+                        info.compCompHnd->compareTypesForCast(likelyCls, pResolvedToken->hClass);
+
+                    // If case of MustNot we still can optimize isinst (only), e.g.:
+                    //
+                    //   bool objIsDisposable = obj is IDisposable;
+                    //
+                    // when the profile tells us that obj is mostly Int32, hence, never implements that interface.
+                    // for castclass it makes little sense as it will always throw a cast exception anyway.
+                    if ((castResult == TypeCompareState::Must) ||
+                        (castResult == TypeCompareState::MustNot && !isCastClass))
+                    {
+                        bool isAbstract = (info.compCompHnd->getClassAttribs(likelyCls) &
+                                           (CORINFO_FLG_INTERFACE | CORINFO_FLG_ABSTRACT)) != 0;
+                        // If it's abstract it means we most likely deal with a stale PGO data so bail out.
+                        if (!isAbstract)
+                        {
+                            JITDUMP("Adding \"is %s (%X)\" check as a fast path for %s using PGO data.\n",
+                                    eeGetClassName(likelyCls), likelyCls, isCastClass ? "castclass" : "isinst");
+
+                            reversedMTCheck    = castResult == TypeCompareState::MustNot;
+                            canExpandInline    = true;
+                            partialExpand      = true;
+                            exactCls           = likelyCls;
+                            fastPathLikelihood = likelyClass.likelihood;
+                        }
+                    }
+                }
+            }
+        }
     }
 
     const bool expandInline = canExpandInline && shouldExpandInline;
@@ -5548,12 +5619,6 @@ GenTree* Compiler::impCastClassOrIsInstToTree(
                 compCurBB->SetFlags(BBF_HAS_HISTOGRAM_PROFILE);
             }
         }
-        else if (impIsCastHelperMayHaveProfileData(helper))
-        {
-            // Leave a note for fgLateCastExpand to expand this helper call
-            call->gtCallMoreFlags |= GTF_CALL_M_CAST_CAN_BE_EXPANDED;
-            call->gtCastHelperILOffset = ilOffset;
-        }
         return call;
     }