JIT: Handle QMARK properly for LCL_ADDR propagation

src/coreclr/jit/lclmorph.cpp

@@ -225,10 +225,11 @@ class LocalEqualsLocalAddrAssertions
     AssertionToIndexMap                       m_map;
     uint64_t*                                 m_lclAssertions;
     uint64_t*                                 m_outgoingAssertions;
-    uint64_t                                  m_currentAssertions = 0;
     BitVec                                    m_localsToExpose;
 
 public:
+    uint64_t CurrentAssertions = 0;
+
     LocalEqualsLocalAddrAssertions(Compiler* comp)
         : m_comp(comp)
         , m_assertions(comp->getAllocator(CMK_LocalAddressVisitor))
@@ -285,28 +286,28 @@ class LocalEqualsLocalAddrAssertions
     {
         if ((m_assertions.Height() == 0) || (block->bbPreds == nullptr) || m_comp->bbIsHandlerBeg(block))
         {
-            m_currentAssertions = 0;
+            CurrentAssertions = 0;
             return;
         }
 
-        m_currentAssertions = UINT64_MAX;
+        CurrentAssertions = UINT64_MAX;
         for (BasicBlock* pred : block->PredBlocks())
         {
             assert(m_comp->m_dfsTree->Contains(pred));
             if (pred->bbPostorderNum <= block->bbPostorderNum)
             {
-                m_currentAssertions = 0;
+                CurrentAssertions = 0;
                 break;
             }
 
-            m_currentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
+            CurrentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
         }
 
 #ifdef DEBUG
-        if (m_currentAssertions != 0)
+        if (CurrentAssertions != 0)
         {
             JITDUMP(FMT_BB " incoming assertions:\n", block->bbNum);
-            uint64_t assertions = m_currentAssertions;
+            uint64_t assertions = CurrentAssertions;
             do
             {
                 uint32_t index = BitOperations::BitScanForward(assertions);
@@ -328,7 +329,7 @@ class LocalEqualsLocalAddrAssertions
     //
     void EndBlock(BasicBlock* block)
     {
-        m_outgoingAssertions[block->bbPostorderNum] = m_currentAssertions;
+        m_outgoingAssertions[block->bbPostorderNum] = CurrentAssertions;
     }
 
     //-------------------------------------------------------------------
@@ -389,7 +390,7 @@ class LocalEqualsLocalAddrAssertions
             }
         }
 
-        m_currentAssertions |= uint64_t(1) << index;
+        CurrentAssertions |= uint64_t(1) << index;
     }
 
     //-------------------------------------------------------------------
@@ -400,7 +401,7 @@ class LocalEqualsLocalAddrAssertions
     //
     void Clear(unsigned dstLclNum)
     {
-        m_currentAssertions &= ~m_lclAssertions[dstLclNum];
+        CurrentAssertions &= ~m_lclAssertions[dstLclNum];
     }
 
     //-----------------------------------------------------------------------------------
@@ -415,7 +416,7 @@ class LocalEqualsLocalAddrAssertions
     //
     const LocalEqualsLocalAddrAssertion* GetCurrentAssertion(unsigned lclNum)
     {
-        uint64_t curAssertion = m_currentAssertions & m_lclAssertions[lclNum];
+        uint64_t curAssertion = CurrentAssertions & m_lclAssertions[lclNum];
         assert(genMaxOneBit(curAssertion));
         if (curAssertion == 0)
         {
@@ -797,6 +798,9 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
             }
             break;
 
+            case GT_QMARK:
+                return HandleQMarkSubTree(use);
+
             default:
                 break;
         }
@@ -1038,6 +1042,81 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
     }
 
 private:
+    //------------------------------------------------------------------------
+    // HandleQMarkSubTree: Process a sub-tree rooted at a GT_QMARK.
+    //
+    // Arguments:
+    //   use - the use of the qmark
+    //
+    // Returns:
+    //   The walk result.
+    //
+    // Remarks:
+    //   GT_QMARK needs special handling due to the conditional nature of it.
+    //   Particularly when we are optimizing and propagating LCL_ADDRs we need
+    //   to take care that assertions created inside the conditionally executed
+    //   parts are handled appropriately. This function inlines the pre and
+    //   post-order visit logic here to make that handling work.
+    //
+    fgWalkResult HandleQMarkSubTree(GenTree** use)
+    {
+        assert((*use)->OperIs(GT_QMARK));
+        GenTreeQmark* qmark = (*use)->AsQmark();
+
+        // We have to inline the pre/postorder visit here to handle
+        // assertions properly.
+        assert(!qmark->IsReverseOp());
+        if (WalkTree(&qmark->gtOp1, qmark) == Compiler::WALK_ABORT)
+        {
+            return Compiler::WALK_ABORT;
+        }
+
+        if (m_lclAddrAssertions != nullptr)
+        {
+            uint64_t origAssertions = m_lclAddrAssertions->CurrentAssertions;
+
+            if (WalkTree(&qmark->gtOp2->AsOp()->gtOp1, qmark->gtOp2) == Compiler::WALK_ABORT)
+            {
+                return Compiler::WALK_ABORT;
+            }
+
+            uint64_t op1Assertions                 = m_lclAddrAssertions->CurrentAssertions;
+            m_lclAddrAssertions->CurrentAssertions = origAssertions;
+
+            if (WalkTree(&qmark->gtOp2->AsOp()->gtOp2, qmark->gtOp2) == Compiler::WALK_ABORT)
+            {
+                return Compiler::WALK_ABORT;
+            }
+
+            uint64_t op2Assertions                 = m_lclAddrAssertions->CurrentAssertions;
+            m_lclAddrAssertions->CurrentAssertions = op1Assertions & op2Assertions;
+        }
+        else
+        {
+            if ((WalkTree(&qmark->gtOp2->AsOp()->gtOp1, qmark->gtOp2) == Compiler::WALK_ABORT) ||
+                (WalkTree(&qmark->gtOp2->AsOp()->gtOp2, qmark->gtOp2) == Compiler::WALK_ABORT))
+            {
+                return Compiler::WALK_ABORT;
+            }
+        }
+
+        assert(TopValue(0).Node() == qmark->gtGetOp2()->gtGetOp2());
+        assert(TopValue(1).Node() == qmark->gtGetOp2()->gtGetOp1());
+        assert(TopValue(2).Node() == qmark->gtGetOp1());
+
+        EscapeValue(TopValue(0), qmark->gtGetOp2());
+        PopValue();
+
+        EscapeValue(TopValue(0), qmark->gtGetOp2());
+        PopValue();
+
+        EscapeValue(TopValue(0), qmark);
+        PopValue();
+
+        PushValue(use);
+        return Compiler::WALK_SKIP_SUBTREES;
+    }
+
     void PushValue(GenTree** use)
     {
         m_valueStack.Push(use);