Fix for Issue 21231

When transferring a Zero offset from one GenTree node to another, we need to check if there already is a FieldSeq and append to it.
Added third parameter 'kind' to JitHashTable::Set, and Added enum SetKind
Only allow Set to overwrite an existing entry when kind is set to Overwrite.
Added validation for all calls to JitHashTable::Set
asserting that we don't expect the key to already exist or that we passed Overwrite indicating that we expect to handle it properly.

Added two test cases for Issue 21231

7 files changed, 46 insertions, 18 deletions

diff --git a/src/jit/compiler.hpp b/src/jit/compiler.hpp
index a9dc3bbc5f..c1be6acbbe 100644
--- a/src/jit/compiler.hpp
+++ b/src/jit/compiler.hpp
@@ -3291,7 +3291,7 @@ inline void Compiler::LoopDsc::AddModifiedField(Compiler* comp, CORINFO_FIELD_HA
         lpFieldsModified =
             new (comp->getAllocatorLoopHoist()) Compiler::LoopDsc::FieldHandleSet(comp->getAllocatorLoopHoist());
     }
-    lpFieldsModified->Set(fldHnd, true);
+    lpFieldsModified->Set(fldHnd, true, FieldHandleSet::Overwrite);
 }
 
 inline void Compiler::LoopDsc::AddModifiedElemType(Compiler* comp, CORINFO_CLASS_HANDLE structHnd)
@@ -3301,7 +3301,7 @@ inline void Compiler::LoopDsc::AddModifiedElemType(Compiler* comp, CORINFO_CLASS
         lpArrayElemTypesModified =
             new (comp->getAllocatorLoopHoist()) Compiler::LoopDsc::ClassHandleSet(comp->getAllocatorLoopHoist());
     }
-    lpArrayElemTypesModified->Set(structHnd, true);
+    lpArrayElemTypesModified->Set(structHnd, true, ClassHandleSet::Overwrite);
 }
 
 inline void Compiler::LoopDsc::VERIFY_lpIterTree()
diff --git a/src/jit/copyprop.cpp b/src/jit/copyprop.cpp
index d0d173d519..4611e91b54 100644
--- a/src/jit/copyprop.cpp
+++ b/src/jit/copyprop.cpp
@@ -365,7 +365,7 @@ void Compiler::optBlockCopyProp(BasicBlock* block, LclNumToGenTreePtrStack* curS
                     stack = new (curSsaName->GetAllocator()) GenTreePtrStack(curSsaName->GetAllocator());
                 }
                 stack->Push(tree);
-                curSsaName->Set(lclNum, stack);
+                curSsaName->Set(lclNum, stack, LclNumToGenTreePtrStack::Overwrite);
             }
             // If we encounter first use of a param or this pointer add it as a live definition.
             // Since they are always live, do it only once.
diff --git a/src/jit/jithashtable.h b/src/jit/jithashtable.h
index d411a2b870..c68e4c8480 100644
--- a/src/jit/jithashtable.h
+++ b/src/jit/jithashtable.h
@@ -234,15 +234,25 @@ public:
     // Arguments:
     //    k - the key
     //    v - the value
+    //    kind - Normal, we are not allowed to overwrite
+    //           Overwrite, we are allowed to overwrite
+    //           currently only used by CHK/DBG builds in an assert.
     //
     // Return Value:
-    //    `true` if the key already exists, `false` otherwise.
+    //    `true` if the key exists and was overwritten,
+    //    `false` otherwise.
     //
     // Notes:
-    //    If the key already exists then its associated value is updated to
-    //    the new value.
+    //    If the key already exists and kind is Normal
+    //    this method will assert
     //
-    bool Set(Key k, Value v)
+    enum SetKind
+    {
+        None,
+        Overwrite
+    };
+
+    bool Set(Key k, Value v, SetKind kind = None)
     {
         CheckGrowth();
 
@@ -257,6 +267,7 @@ public:
         }
         if (pN != nullptr)
         {
+            assert(kind == Overwrite);
             pN->m_val = v;
             return true;
         }
diff --git a/src/jit/lclvars.cpp b/src/jit/lclvars.cpp
index b9cec032b2..a251534bd1 100644
--- a/src/jit/lclvars.cpp
+++ b/src/jit/lclvars.cpp
@@ -1949,7 +1949,7 @@ Compiler::lvaStructFieldInfo Compiler::StructPromotionHelper::GetFieldInfo(CORIN
             fieldInfo.fldType = compiler->getSIMDTypeForSize(simdSize);
             fieldInfo.fldSize = simdSize;
 #ifdef DEBUG
-            retypedFieldsMap.Set(fieldInfo.fldHnd, fieldInfo.fldType);
+            retypedFieldsMap.Set(fieldInfo.fldHnd, fieldInfo.fldType, RetypedAsScalarFieldsMap::Overwrite);
 #endif // DEBUG
         }
     }
@@ -2053,7 +2053,7 @@ bool Compiler::StructPromotionHelper::TryPromoteStructField(lvaStructFieldInfo&
     fieldInfo.fldType = fieldVarType;
     fieldInfo.fldSize = fieldSize;
 #ifdef DEBUG
-    retypedFieldsMap.Set(fieldInfo.fldHnd, fieldInfo.fldType);
+    retypedFieldsMap.Set(fieldInfo.fldHnd, fieldInfo.fldType, RetypedAsScalarFieldsMap::Overwrite);
 #endif // DEBUG
     return true;
 }
diff --git a/src/jit/morph.cpp b/src/jit/morph.cpp
index 48d97cfe4c..9d0193ba85 100644
--- a/src/jit/morph.cpp
+++ b/src/jit/morph.cpp
@@ -13640,8 +13640,15 @@ DONE_MORPHING_CHILDREN:
 
                 if (isZeroOffset)
                 {
+                    // The "op1" node might already be annotated with a zero-offset field sequence.
+                    FieldSeqNode* existingZeroOffsetFldSeq = nullptr;
+                    if (GetZeroOffsetFieldMap()->Lookup(op1, &existingZeroOffsetFldSeq))
+                    {
+                        // Append the zero field sequences
+                        zeroFieldSeq = GetFieldSeqStore()->Append(existingZeroOffsetFldSeq, zeroFieldSeq);
+                    }
                     // Transfer the annotation to the new GT_ADDR node.
-                    GetZeroOffsetFieldMap()->Set(op1, zeroFieldSeq);
+                    GetZeroOffsetFieldMap()->Set(op1, zeroFieldSeq, NodeToFieldSeqMap::Overwrite);
                 }
                 commaNode->gtOp.gtOp2 = op1;
                 // Originally, I gave all the comma nodes type "byref".  But the ADDR(IND(x)) == x transform
@@ -18743,7 +18750,16 @@ void Compiler::fgAddFieldSeqForZeroOffset(GenTree* op1, FieldSeqNode* fieldSeq)
 
         default:
             // Record in the general zero-offset map.
-            GetZeroOffsetFieldMap()->Set(op1, fieldSeq);
+
+            // The "op1" node might already be annotated with a zero-offset field sequence.
+            FieldSeqNode* existingZeroOffsetFldSeq = nullptr;
+            if (GetZeroOffsetFieldMap()->Lookup(op1, &existingZeroOffsetFldSeq))
+            {
+                // Append the zero field sequences
+                fieldSeq = GetFieldSeqStore()->Append(existingZeroOffsetFldSeq, fieldSeq);
+            }
+            // Set the new field sequence annotation for op1
+            GetZeroOffsetFieldMap()->Set(op1, fieldSeq, NodeToFieldSeqMap::Overwrite);
             break;
     }
 }
diff --git a/src/jit/optimizer.cpp b/src/jit/optimizer.cpp
index 9d9b515ec4..795625aaf2 100644
--- a/src/jit/optimizer.cpp
+++ b/src/jit/optimizer.cpp
@@ -3775,7 +3775,7 @@ void Compiler::optUnrollLoops()
                 {
                     BasicBlock* newBlock = insertAfter =
                         fgNewBBafter(block->bbJumpKind, insertAfter, /*extendRegion*/ true);
-                    blockMap.Set(block, newBlock);
+                    blockMap.Set(block, newBlock, BlockToBlockMap::Overwrite);
 
                     if (!BasicBlock::CloneBlockState(this, newBlock, block, lvar, lval))
                     {
diff --git a/src/jit/rangecheck.cpp b/src/jit/rangecheck.cpp
index b96374455c..8255afeecf 100644
--- a/src/jit/rangecheck.cpp
+++ b/src/jit/rangecheck.cpp
@@ -362,7 +362,7 @@ bool RangeCheck::IsMonotonicallyIncreasing(GenTree* expr, bool rejectNegativeCon
     JITDUMP("[RangeCheck::IsMonotonicallyIncreasing] [%06d]\n", Compiler::dspTreeID(expr));
 
     // Add hashtable entry for expr.
-    bool alreadyPresent = m_pSearchPath->Set(expr, nullptr);
+    bool alreadyPresent = !m_pSearchPath->Set(expr, nullptr, SearchPath::Overwrite);
     if (alreadyPresent)
     {
         return true;
@@ -1014,7 +1014,7 @@ bool RangeCheck::DoesOverflow(BasicBlock* block, GenTree* expr)
 bool RangeCheck::ComputeDoesOverflow(BasicBlock* block, GenTree* expr)
 {
     JITDUMP("Does overflow [%06d]?\n", Compiler::dspTreeID(expr));
-    m_pSearchPath->Set(expr, block);
+    m_pSearchPath->Set(expr, block, SearchPath::Overwrite);
 
     bool overflows = true;
 
@@ -1050,7 +1050,7 @@ bool RangeCheck::ComputeDoesOverflow(BasicBlock* block, GenTree* expr)
     {
         overflows = DoesPhiOverflow(block, expr);
     }
-    GetOverflowMap()->Set(expr, overflows);
+    GetOverflowMap()->Set(expr, overflows, OverflowMap::Overwrite);
     m_pSearchPath->Remove(expr);
     return overflows;
 }
@@ -1064,11 +1064,12 @@ bool RangeCheck::ComputeDoesOverflow(BasicBlock* block, GenTree* expr)
 // eg.: merge((0, dep), (dep, dep)) = (0, dep)
 Range RangeCheck::ComputeRange(BasicBlock* block, GenTree* expr, bool monotonic DEBUGARG(int indent))
 {
-    bool  newlyAdded = !m_pSearchPath->Set(expr, block);
+    bool  newlyAdded = !m_pSearchPath->Set(expr, block, SearchPath::Overwrite);
     Range range      = Limit(Limit::keUndef);
 
     ValueNum vn = m_pCompiler->vnStore->VNConservativeNormalValue(expr->gtVNPair);
-    // If newly added in the current search path, then reduce the budget.
+
+    // If we just added 'expr' in the current search path, then reduce the budget.
     if (newlyAdded)
     {
         // Assert that we are not re-entrant for a node which has been
@@ -1172,7 +1173,7 @@ Range RangeCheck::ComputeRange(BasicBlock* block, GenTree* expr, bool monotonic
         range = Range(Limit(Limit::keUnknown));
     }
 
-    GetRangeMap()->Set(expr, new (m_alloc) Range(range));
+    GetRangeMap()->Set(expr, new (m_alloc) Range(range), RangeMap::Overwrite);
     m_pSearchPath->Remove(expr);
     return range;
 }