From 55e9ab59bdd73abba398687fab994a38fa3cb28b Mon Sep 17 00:00:00 2001
From: Andy Ayers <andya@microsoft.com>
Date: Thu, 30 Jan 2025 18:37:48 -0800
Subject: [PATCH] JIT: fix case where implied subrange assertions can get lost
 in morph (#112020)

If a method used short bit vectors during morph, modifications to the
assertion set were being made to a copy, not to the original.

Also, restrict implied boolean subrange assertions to integer vars.
---
 src/coreclr/jit/morph.cpp | 39 ++++++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp
index a787fece19fb6d..90f1bc9161ebea 100644
--- a/src/coreclr/jit/morph.cpp
+++ b/src/coreclr/jit/morph.cpp
@@ -12412,26 +12412,31 @@ void Compiler::fgAssertionGen(GenTree* tree)
     // If this tree creates an assignment of 0 or 1 to an int local, also create a [0..1] subrange
     // assertion for that local, in case this local is used as a bool.
     //
-    auto addImpliedBoolSubrangeAssertion = [=](AssertionIndex index, ASSERT_TP assertions) {
+    auto addImpliedAssertions = [=](AssertionIndex index, ASSERT_TP& assertions) {
         AssertionDsc* const assertion = optGetAssertion(index);
         if ((assertion->assertionKind == OAK_EQUAL) && (assertion->op1.kind == O1K_LCLVAR) &&
             (assertion->op2.kind == O2K_CONST_INT))
         {
-            ssize_t iconVal = assertion->op2.u1.iconVal;
-            if ((iconVal == 0) || (iconVal == 1))
-            {
-                AssertionDsc extraAssertion   = {OAK_SUBRANGE};
-                extraAssertion.op1.kind       = O1K_LCLVAR;
-                extraAssertion.op1.lcl.lclNum = assertion->op1.lcl.lclNum;
-                extraAssertion.op2.kind       = O2K_SUBRANGE;
-                extraAssertion.op2.u2         = IntegralRange(SymbolicIntegerValue::Zero, SymbolicIntegerValue::One);
+            LclVarDsc* const lclDsc = lvaGetDesc(assertion->op1.lcl.lclNum);
 
-                AssertionIndex extraIndex = optFinalizeCreatingAssertion(&extraAssertion);
-                if (extraIndex != NO_ASSERTION_INDEX)
+            if (varTypeIsIntegral(lclDsc->TypeGet()))
+            {
+                ssize_t iconVal = assertion->op2.u1.iconVal;
+                if ((iconVal == 0) || (iconVal == 1))
                 {
-                    unsigned const bvIndex = extraIndex - 1;
-                    BitVecOps::AddElemD(apTraits, assertions, bvIndex);
-                    announce(extraIndex, "[bool range] ");
+                    AssertionDsc extraAssertion   = {OAK_SUBRANGE};
+                    extraAssertion.op1.kind       = O1K_LCLVAR;
+                    extraAssertion.op1.lcl.lclNum = assertion->op1.lcl.lclNum;
+                    extraAssertion.op2.kind       = O2K_SUBRANGE;
+                    extraAssertion.op2.u2 = IntegralRange(SymbolicIntegerValue::Zero, SymbolicIntegerValue::One);
+
+                    AssertionIndex extraIndex = optFinalizeCreatingAssertion(&extraAssertion);
+                    if (extraIndex != NO_ASSERTION_INDEX)
+                    {
+                        unsigned const bvIndex = extraIndex - 1;
+                        BitVecOps::AddElemD(apTraits, assertions, bvIndex);
+                        announce(extraIndex, "[bool range] ");
+                    }
                 }
             }
         }
@@ -12485,7 +12490,7 @@ void Compiler::fgAssertionGen(GenTree* tree)
             announce(ifTrueAssertionIndex, "[if true] ");
             unsigned const bvIndex = ifTrueAssertionIndex - 1;
             BitVecOps::AddElemD(apTraits, apLocalIfTrue, bvIndex);
-            addImpliedBoolSubrangeAssertion(ifTrueAssertionIndex, apLocalIfTrue);
+            addImpliedAssertions(ifTrueAssertionIndex, apLocalIfTrue);
         }
 
         if (ifFalseAssertionIndex != NO_ASSERTION_INDEX)
@@ -12493,7 +12498,7 @@ void Compiler::fgAssertionGen(GenTree* tree)
             announce(ifFalseAssertionIndex, "[if false] ");
             unsigned const bvIndex = ifFalseAssertionIndex - 1;
             BitVecOps::AddElemD(apTraits, apLocal, ifFalseAssertionIndex - 1);
-            addImpliedBoolSubrangeAssertion(ifFalseAssertionIndex, apLocal);
+            addImpliedAssertions(ifFalseAssertionIndex, apLocal);
         }
     }
     else
@@ -12502,7 +12507,7 @@ void Compiler::fgAssertionGen(GenTree* tree)
         announce(apIndex, "");
         unsigned const bvIndex = apIndex - 1;
         BitVecOps::AddElemD(apTraits, apLocal, bvIndex);
-        addImpliedBoolSubrangeAssertion(apIndex, apLocal);
+        addImpliedAssertions(apIndex, apLocal);
     }
 }