Remove MorphAddrContext, use peel-offsets-style helper instead

Replace the threaded MorphAddrContext with fgMarkAddrModeForFieldAddr,
a pre-walk over an indirection's address tree that peels constant
offsets and chained instance FIELD_ADDRs to decide null-check elision.

The new GTF_FLD_TGT_NONFAULTING flag on FIELD_ADDR communicates that
the consuming indirection will fault on null, so the explicit null
check inserted by morph can be elided.

SPMI asmdiffs: no diffs across all 11 collections.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: EgorBo

src/coreclr/jit/compiler.h

@@ -6800,18 +6800,7 @@ class Compiler
     GenTree* fgMorphIntoHelperCall(
         GenTree* tree, int helper, bool morphArgs, GenTree* arg1 = nullptr, GenTree* arg2 = nullptr);
 
-    // A "MorphAddrContext" carries information from the surrounding context.  If we are evaluating a byref address,
-    // it is useful to know whether the address will be immediately dereferenced, or whether the address value will
-    // be used, perhaps by passing it as an argument to a called method.  This affects how null checking is done:
-    // for sufficiently small offsets, we can rely on OS page protection to implicitly null-check addresses that we
-    // know will be dereferenced.  To know that reliance on implicit null checking is sound, we must further know that
-    // all offsets between the top-level indirection and the bottom are constant, and that their sum is sufficiently
-    // small; hence the other fields of MorphAddrContext.
-    struct MorphAddrContext
-    {
-        GenTreeIndir* m_user = nullptr;  // Indirection using this address.
-        size_t        m_totalOffset = 0; // Sum of offsets between the top-level indirection and here (current context).
-    };
+    void fgMarkAddrModeForFieldAddr(GenTreeIndir* indir);
 
 #ifdef FEATURE_SIMD
     GenTree* getSIMDStructFromField(GenTree*  tree,
@@ -6846,8 +6835,8 @@ class Compiler
     void fgAssignSetVarDef(GenTree* tree);
 
 private:
-    GenTree* fgMorphFieldAddr(GenTree* tree, MorphAddrContext* mac);
-    GenTree* fgMorphExpandInstanceField(GenTree* tree, MorphAddrContext* mac);
+    GenTree* fgMorphFieldAddr(GenTree* tree);
+    GenTree* fgMorphExpandInstanceField(GenTree* tree);
     GenTree* fgMorphExpandTlsFieldAddr(GenTree* tree);
     bool fgCanFastTailCall(GenTreeCall* call, const char** failReason);
 #if FEATURE_FASTTAILCALL
@@ -6902,7 +6891,7 @@ class Compiler
     GenTree* fgMorphInitBlock(GenTree* tree);
     GenTree* fgMorphCopyBlock(GenTree* tree);
 private:
-    GenTree* fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optAssertionPropDone = nullptr);
+    GenTree* fgMorphSmpOp(GenTree* tree, bool* optAssertionPropDone = nullptr);
     bool fgTryReplaceStructLocalWithFields(GenTree** use);
     GenTree* fgMorphFinalizeIndir(GenTreeIndir* indir);
     GenTree* fgOptimizeCast(GenTreeCast* cast);
@@ -6937,7 +6926,7 @@ class Compiler
     GenTree* fgMorphReduceAddOps(GenTree* tree);
 
 public:
-    GenTree* fgMorphTree(GenTree* tree, MorphAddrContext* mac = nullptr);
+    GenTree* fgMorphTree(GenTree* tree);
 
 private:
     void fgAssertionGen(GenTree* tree);

src/coreclr/jit/gentree.h

@@ -469,6 +469,8 @@ enum GenTreeFlags : unsigned
 
     GTF_FLD_TLS                 = 0x80000000, // GT_FIELD_ADDR -- field address is a Windows x86 TLS reference
     GTF_FLD_DEREFERENCED        = 0x40000000, // GT_FIELD_ADDR -- used to preserve previous behavior
+    GTF_FLD_TGT_NONFAULTING     = 0x20000000, // GT_FIELD_ADDR -- the consuming indirection will fault on null,
+                                              // so the explicit null check inserted by morph can be elided.
 
     GTF_INX_RNGCHK              = 0x80000000, // GT_INDEX_ADDR -- this array address should be range-checked
     GTF_INX_ADDR_NONNULL        = 0x40000000, // GT_INDEX_ADDR -- this array address is not null

src/coreclr/jit/morph.cpp

@@ -3527,19 +3527,113 @@ unsigned Compiler::fgGetFieldMorphingTemp(GenTreeFieldAddr* fieldNode)
     return lclNum;
 }
 
+//------------------------------------------------------------------------
+// fgMarkAddrModeForFieldAddr: Walk an indirection's unmorphed address tree to
+// identify FIELD_ADDR nodes whose explicit null check can be elided because
+// this indirection will fault on null instead.
+//
+// Replaces the legacy "MorphAddrContext" thread-through mechanism: peels
+// constant offsets off ADD nodes (matching the previous "ADD(_, const)"
+// propagation) and recurses through chained instance FIELD_ADDR nodes,
+// accumulating the cumulative constant offset that would be added on top of
+// the FIELD_ADDR's base object.
+//
+// For each instance FIELD_ADDR whose base object could be null:
+//   * If the cumulative offset (including this field's offset) is small enough
+//     that the OS page guard catches the access, mark the FIELD_ADDR with
+//     GTF_FLD_TGT_NONFAULTING so morphing can elide the explicit null check,
+//     and clear GTF_IND_NONFAULTING on the consuming indirection so it faults
+//     on null instead.
+//   * Otherwise, the FIELD_ADDR will introduce an explicit NULLCHECK and stop
+//     propagation: mark the indirection as having an ordering side effect (so
+//     downstream optimizations treat the inserted NULLCHECK conservatively),
+//     and stop walking past this FIELD_ADDR.
+//
+// Arguments:
+//    indir - The indirection whose address operand will be morphed.
+//
+void Compiler::fgMarkAddrModeForFieldAddr(GenTreeIndir* indir)
+{
+    GenTree* addr   = indir->Addr();
+    size_t   offset = 0;
+
+    while (true)
+    {
+        // Peel ADD(addr, CNS_INT) -- only follow the op1 spine, like the legacy mac propagation did.
+        if (addr->OperIs(GT_ADD) && !addr->gtOverflow())
+        {
+            GenTree* op2 = addr->gtGetOp2();
+            if (op2->IsCnsIntOrI() && op2->TypeIs(TYP_I_IMPL) && !op2->AsIntCon()->IsIconHandle())
+            {
+                ClrSafeInt<size_t> total(offset);
+                total += op2->AsIntCon()->IconValue();
+                if (total.IsOverflow())
+                {
+                    return;
+                }
+                offset = total.Value();
+                addr   = addr->gtGetOp1();
+                continue;
+            }
+            return;
+        }
+
+        // Peel through a chained instance FIELD_ADDR.
+        if (addr->OperIs(GT_FIELD_ADDR) && addr->AsFieldAddr()->IsInstance())
+        {
+            GenTreeFieldAddr* fld       = addr->AsFieldAddr();
+            GenTree*          objRef    = fld->GetFldObj();
+            unsigned          fldOffset = fld->gtFldOffset;
+
+            if (fgAddrCouldBeNull(objRef))
+            {
+                ClrSafeInt<size_t> total(offset);
+                total += fldOffset;
+
+                bool elide = !total.IsOverflow() && !fgIsBigOffset(total.Value());
+                if (elide)
+                {
+                    fld->gtFlags |= GTF_FLD_TGT_NONFAULTING;
+                    // We can elide the null check only by letting it happen as part of the
+                    // consuming indirection, so it is no longer non-faulting.
+                    indir->gtFlags &= ~GTF_IND_NONFAULTING;
+                }
+                else
+                {
+                    // The FIELD_ADDR will introduce an explicit NULLCHECK, which acts as a
+                    // barrier in the consuming indirection's address chain.
+                    indir->SetHasOrderingSideEffect();
+                    return;
+                }
+            }
+
+            ClrSafeInt<size_t> totalOff(offset);
+            totalOff += fldOffset;
+            if (totalOff.IsOverflow())
+            {
+                return;
+            }
+            offset = totalOff.Value();
+            addr   = objRef;
+            continue;
+        }
+
+        return;
+    }
+}
+
 //------------------------------------------------------------------------
 // fgMorphFieldAddr: Fully morph a FIELD_ADDR tree.
 //
 // Expands the field node into explicit additions.
 //
 // Arguments:
 //    tree - The FIELD_ADDR tree
-//    mac  - The morphing context, used to elide adding null checks
 //
 // Return Value:
 //    The fully morphed "tree".
 //
-GenTree* Compiler::fgMorphFieldAddr(GenTree* tree, MorphAddrContext* mac)
+GenTree* Compiler::fgMorphFieldAddr(GenTree* tree)
 {
     assert(tree->OperIs(GT_FIELD_ADDR));
 
@@ -3549,7 +3643,7 @@ GenTree* Compiler::fgMorphFieldAddr(GenTree* tree, MorphAddrContext* mac)
 
     if (fieldNode->IsInstance())
     {
-        tree = fgMorphExpandInstanceField(tree, mac);
+        tree = fgMorphExpandInstanceField(tree);
     }
     else if (fieldNode->IsTlsStatic())
     {
@@ -3560,11 +3654,10 @@ GenTree* Compiler::fgMorphFieldAddr(GenTree* tree, MorphAddrContext* mac)
         assert(!"Normal statics are expected to be handled in the importer");
     }
 
-    // Pass down the current mac; if non null we are computing an address
     GenTree* result;
     if (tree->OperIsSimple())
     {
-        result = fgMorphSmpOp(tree, mac);
+        result = fgMorphSmpOp(tree);
         result->SetMorphed(this);
 
         // Quirk: preserve previous behavior with this NO_CSE.
@@ -3575,7 +3668,7 @@ GenTree* Compiler::fgMorphFieldAddr(GenTree* tree, MorphAddrContext* mac)
     }
     else
     {
-        result = fgMorphTree(tree, mac);
+        result = fgMorphTree(tree);
     }
 
     JITDUMP("\nFinal value of Compiler::fgMorphFieldAddr after morphing:\n");
@@ -3591,12 +3684,11 @@ GenTree* Compiler::fgMorphFieldAddr(GenTree* tree, MorphAddrContext* mac)
 //
 // Arguments:
 //    tree - The FIELD_ADDR tree
-//    mac  - The morphing context, used to elide adding null checks
 //
 // Return Value:
 //    The expanded "tree" of an arbitrary shape.
 //
-GenTree* Compiler::fgMorphExpandInstanceField(GenTree* tree, MorphAddrContext* mac)
+GenTree* Compiler::fgMorphExpandInstanceField(GenTree* tree)
 {
     assert(tree->OperIs(GT_FIELD_ADDR) && tree->AsFieldAddr()->IsInstance());
 
@@ -3673,29 +3765,18 @@ GenTree* Compiler::fgMorphExpandInstanceField(GenTree* tree, MorphAddrContext* m
 
     if (fgAddrCouldBeNull(objRef))
     {
-        // A non-null context here implies our [+ some offset] parent is an indirection, one that
-        // will implicitly null-check the produced address.
-        addExplicitNullCheck = (mac == nullptr) || fgIsBigOffset(mac->m_totalOffset + fieldOffset);
+        // The pre-walk performed by fgMarkAddrModeForFieldAddr sets GTF_FLD_TGT_NONFAULTING when
+        // the consuming indirection will implicitly null-check the produced address (the cumulative
+        // constant offset down to this FIELD_ADDR is small enough that the OS page guard catches it).
+        addExplicitNullCheck = (tree->gtFlags & GTF_FLD_TGT_NONFAULTING) == 0;
 
         // The transformation here turns a value dependency (FIELD_ADDR being a
         // known non-null operand) into a control-flow dependency (introducing
         // explicit COMMA(NULLCHECK, ...)). This effectively "disconnects" the
-        // null check from the parent of the FIELD_ADDR node. For the cases
-        // where we made use of non-nullness we need to make the dependency
-        // explicit now.
-        if (addExplicitNullCheck)
-        {
-            if (mac != nullptr)
-            {
-                mac->m_user->SetHasOrderingSideEffect();
-            }
-        }
-        else
-        {
-            // We can elide the null check only by letting it happen as part of
-            // the consuming indirection, so it is no longer non-faulting.
-            mac->m_user->gtFlags &= ~GTF_IND_NONFAULTING;
-        }
+        // null check from the parent of the FIELD_ADDR node.
+        //
+        // For the elided case, the consuming indirection has already been marked
+        // (~GTF_IND_NONFAULTING) by fgMarkAddrModeForFieldAddr; nothing else to do here.
     }
 
     if (addExplicitNullCheck)
@@ -6920,14 +7001,13 @@ GenTreeOp* Compiler::fgMorphCommutative(GenTreeOp* tree)
 //
 // Arguments:
 //    tree - tree to morph
-//    mac  - address context for morphing
 //    optAssertionPropDone - [out, optional] set true if local assertions
 //       were killed/genned while morphing this tree
 //
 // Returns:
 //    Tree, possibly updated
 //
-GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optAssertionPropDone)
+GenTree* Compiler::fgMorphSmpOp(GenTree* tree, bool* optAssertionPropDone)
 {
     assert(tree->OperKind() & GTK_SMPOP);
 
@@ -7011,7 +7091,7 @@ GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optA
             break;
 
         case GT_FIELD_ADDR:
-            return fgMorphFieldAddr(tree, mac);
+            return fgMorphFieldAddr(tree);
 
         case GT_INDEX_ADDR:
             return fgMorphIndexAddr(tree->AsIndexAddr());
@@ -7140,7 +7220,7 @@ GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optA
             {
                 assert(tree->OperIs(GT_DIV));
                 tree->ChangeOper(GT_UDIV, GenTree::PRESERVE_VN);
-                return fgMorphSmpOp(tree, mac);
+                return fgMorphSmpOp(tree);
             }
 
 #if !defined(TARGET_64BIT) && !defined(TARGET_WASM)
@@ -7207,7 +7287,7 @@ GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optA
             {
                 assert(tree->OperIs(GT_MOD));
                 tree->ChangeOper(GT_UMOD, GenTree::PRESERVE_VN);
-                return fgMorphSmpOp(tree, mac);
+                return fgMorphSmpOp(tree);
             }
 
             // Do not use optimizations (unlike UMOD's idiv optimizing during codegen) for signed mod.
@@ -7533,34 +7613,14 @@ GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optA
         // the other operands in case this is a store. However, doing so unconditionally preserves previous
         // behavior and "fixes up" field store importation that places the null check in the wrong location
         // (before the 'value' operand is evaluated).
-        MorphAddrContext indMac;
         if (tree->OperIsIndir() && !tree->OperIsAtomicOp())
         {
-            // Communicate to FIELD_ADDR morphing that the parent is an indirection.
-            indMac.m_user = tree->AsIndir();
-            mac           = &indMac;
-        }
-        // For additions, if we already have a context, keep track of whether all offsets added
-        // to the address are constant, and their sum does not overflow.
-        else if ((mac != nullptr) && tree->OperIs(GT_ADD) && op2->IsCnsIntOrI())
-        {
-            ClrSafeInt<size_t> offset(mac->m_totalOffset);
-            offset += op2->AsIntCon()->IconValue();
-            if (!offset.IsOverflow())
-            {
-                mac->m_totalOffset = offset.Value();
-            }
-            else
-            {
-                mac = nullptr;
-            }
-        }
-        else // Reset the context.
-        {
-            mac = nullptr;
+            // Examine the unmorphed address tree to identify any FIELD_ADDR whose null check
+            // can be elided because this indirection will fault on null instead.
+            fgMarkAddrModeForFieldAddr(tree->AsIndir());
         }
 
-        tree->AsOp()->gtOp1 = op1 = fgMorphTree(op1, mac);
+        tree->AsOp()->gtOp1 = op1 = fgMorphTree(op1);
 
         // If we are exiting the "then" part of a Qmark-Colon we must
         // save the state of the current assertions table so that we
@@ -12448,7 +12508,7 @@ GenTreeOp* Compiler::fgMorphLongMul(GenTreeOp* mul)
  *  Transform the given tree for code generation and return an equivalent tree.
  */
 
-GenTree* Compiler::fgMorphTree(GenTree* tree, MorphAddrContext* mac)
+GenTree* Compiler::fgMorphTree(GenTree* tree)
 {
     assert(tree);
     tree->ClearMorphed();
@@ -12558,7 +12618,7 @@ GenTree* Compiler::fgMorphTree(GenTree* tree, MorphAddrContext* mac)
 
     if (kind & GTK_SMPOP)
     {
-        tree = fgMorphSmpOp(tree, mac, &optAssertionPropDone);
+        tree = fgMorphSmpOp(tree, &optAssertionPropDone);
         goto DONE;
     }