JIT: enhance escape analysis to understand Span<T> capture

Implement a very simplistic "field sensitive" analysis for `Span<T>`
where we model the span as simply its byref field.

If the span is only consumed locally, and the array does not otherwise
escape, then the array does not escape.

This is a subset of dotnet#112543 that does not try and reason interprocedurally.

Contributes to dotnet#104936 / dotnet#108913

Author: AndyAyersMS

src/coreclr/jit/compiler.h

@@ -1048,6 +1048,13 @@ class LclVarDsc
         m_layout = layout;
     }
 
+    // Change the layout to one that may not be compatible.
+    void ChangeLayout(ClassLayout* layout)
+    {
+        assert(varTypeIsStruct(lvType));
+        m_layout = layout;
+    }
+
     // Grow the size of a block layout local.
     void GrowBlockLayout(ClassLayout* layout)
     {

src/coreclr/jit/objectalloc.cpp

@@ -233,7 +233,7 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph()
                 lclEscapes = false;
                 m_allocator->CheckForGuardedAllocationOrCopy(m_block, m_stmt, use, lclNum);
             }
-            else if (tree->OperIs(GT_LCL_VAR) && tree->TypeIs(TYP_REF, TYP_BYREF, TYP_I_IMPL))
+            else if (tree->OperIs(GT_LCL_VAR, GT_LCL_ADDR) && tree->TypeIs(TYP_REF, TYP_BYREF, TYP_I_IMPL, TYP_STRUCT))
             {
                 assert(tree == m_ancestors.Top());
                 if (!m_allocator->CanLclVarEscapeViaParentStack(&m_ancestors, lclNum, m_block))
@@ -265,7 +265,7 @@ void ObjectAllocator::MarkEscapingVarsAndBuildConnGraph()
     {
         var_types type = comp->lvaTable[lclNum].TypeGet();
 
-        if (type == TYP_REF || genActualType(type) == TYP_I_IMPL || type == TYP_BYREF)
+        if (type == TYP_REF || genActualType(type) == TYP_I_IMPL || type == TYP_BYREF || type == TYP_STRUCT)
         {
             m_ConnGraphAdjacencyMatrix[lclNum] = BitVecOps::MakeEmpty(&m_bitVecTraits);
 
@@ -1023,10 +1023,13 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent
     assert(parentStack != nullptr);
     int parentIndex = 1;
 
-    bool keepChecking                  = true;
-    bool canLclVarEscapeViaParentStack = true;
-    bool isCopy                        = true;
-    bool isEnumeratorLocal             = comp->lvaGetDesc(lclNum)->lvIsEnumerator;
+    LclVarDsc* const lclDsc = comp->lvaGetDesc(lclNum);
+
+    bool       keepChecking                  = true;
+    bool       canLclVarEscapeViaParentStack = true;
+    bool       isCopy                        = true;
+    bool const isEnumeratorLocal             = lclDsc->lvIsEnumerator;
+    bool const isSpanLocal                   = lclDsc->IsSpan();
 
     while (keepChecking)
     {
@@ -1093,11 +1096,21 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent
             case GT_ADD:
             case GT_SUB:
             case GT_FIELD_ADDR:
-                // Check whether the local escapes via its grandparent.
+                // Check whether the local escapes higher up
                 ++parentIndex;
                 keepChecking = true;
                 break;
 
+            case GT_LCL_ADDR:
+                if (isSpanLocal)
+                {
+                    // Check whether the local escapes higher up
+                    ++parentIndex;
+                    keepChecking = true;
+                    JITDUMP("... i'm good thanks\n");
+                }
+                break;
+
             case GT_BOX:
                 isCopy = wasCopy;
                 ++parentIndex;
@@ -1119,11 +1132,43 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent
             case GT_STOREIND:
             case GT_STORE_BLK:
             case GT_BLK:
-                if (tree != parent->AsIndir()->Addr())
+            {
+                GenTree* const addr = parent->AsIndir()->Addr();
+                if (tree != addr)
                 {
-                    // TODO-ObjectStackAllocation: track stores to fields.
+                    JITDUMP("... tree != addr\n");
+
+                    // Is this an array element address store to (the pointer) field of a span?
+                    // (note we can't yet handle cases where a span captures an object)
+                    //
+                    if (parent->OperIs(GT_STOREIND) && addr->OperIs(GT_FIELD_ADDR) && tree->OperIs(GT_INDEX_ADDR))
+                    {
+                        // Todo: mark the span pointer field addr like we mark IsSpanLength?
+                        // (for now we don't worry which field we store to)
+                        //
+                        GenTree* const base = addr->AsOp()->gtGetOp1();
+
+                        if (base->OperIs(GT_LCL_ADDR))
+                        {
+                            unsigned const   dstLclNum = base->AsLclVarCommon()->GetLclNum();
+                            LclVarDsc* const dstLclDsc = comp->lvaGetDesc(dstLclNum);
+
+                            if (dstLclDsc->IsSpan())
+                            {
+                                JITDUMP("... span ptr store\n");
+                                // Add an edge to the connection graph.
+                                AddConnGraphEdge(dstLclNum, lclNum);
+                                canLclVarEscapeViaParentStack = false;
+                            }
+                        }
+                    }
                     break;
                 }
+                else
+                {
+                    JITDUMP("... tree == addr\n");
+                }
+            }
                 FALLTHROUGH;
             case GT_IND:
                 // Address of the field/ind is not taken so the local doesn't escape.
@@ -1132,20 +1177,20 @@ bool ObjectAllocator::CanLclVarEscapeViaParentStack(ArrayStack<GenTree*>* parent
 
             case GT_CALL:
             {
-                GenTreeCall* const asCall = parent->AsCall();
+                GenTreeCall* const call = parent->AsCall();
 
-                if (asCall->IsHelperCall())
+                if (call->IsHelperCall())
                 {
                     canLclVarEscapeViaParentStack =
-                        !Compiler::s_helperCallProperties.IsNoEscape(comp->eeGetHelperNum(asCall->gtCallMethHnd));
+                        !Compiler::s_helperCallProperties.IsNoEscape(comp->eeGetHelperNum(call->gtCallMethHnd));
                 }
-                else if (asCall->IsSpecialIntrinsic())
+                else if (call->IsSpecialIntrinsic())
                 {
                     // Some known special intrinsics don't escape. At this moment, only the ones accepting byrefs
                     // are supported. In order to support more intrinsics accepting objects, we need extra work
                     // on the VM side which is not ready for that yet.
                     //
-                    switch (comp->lookupNamedIntrinsic(asCall->gtCallMethHnd))
+                    switch (comp->lookupNamedIntrinsic(call->gtCallMethHnd))
                     {
                         case NI_System_SpanHelpers_ClearWithoutReferences:
                         case NI_System_SpanHelpers_Fill:
@@ -1246,6 +1291,7 @@ void ObjectAllocator::UpdateAncestorTypes(GenTree* tree, ArrayStack<GenTree*>* p
             case GT_SUB:
             case GT_FIELD_ADDR:
             case GT_INDEX_ADDR:
+            case GT_LCL_ADDR:
                 if (parent->TypeGet() == TYP_REF)
                 {
                     parent->ChangeType(newType);
@@ -1283,17 +1329,18 @@ void ObjectAllocator::UpdateAncestorTypes(GenTree* tree, ArrayStack<GenTree*>* p
             case GT_STOREIND:
             case GT_STORE_BLK:
             case GT_BLK:
-                assert(tree == parent->AsIndir()->Addr());
-
-                // The new target could be *not* on the heap.
-                parent->gtFlags &= ~GTF_IND_TGT_HEAP;
-
-                if (newType != TYP_BYREF)
+                if (tree == parent->AsIndir()->Addr())
                 {
-                    // This indicates that a write barrier is not needed when writing
-                    // to this field/indirection since the address is not pointing to the heap.
-                    // It's either null or points to inside a stack-allocated object.
-                    parent->gtFlags |= GTF_IND_TGT_NOT_HEAP;
+                    // The new target could be *not* on the heap.
+                    parent->gtFlags &= ~GTF_IND_TGT_HEAP;
+
+                    if (newType != TYP_BYREF)
+                    {
+                        // This indicates that a write barrier is not needed when writing
+                        // to this field/indirection since the address is not pointing to the heap.
+                        // It's either null or points to inside a stack-allocated object.
+                        parent->gtFlags |= GTF_IND_TGT_NOT_HEAP;
+                    }
                 }
                 break;
 
@@ -1354,10 +1401,7 @@ void ObjectAllocator::RewriteUses()
             if ((lclNum < BitVecTraits::GetSize(&m_allocator->m_bitVecTraits)) &&
                 m_allocator->MayLclVarPointToStack(lclNum))
             {
-                // Analysis does not handle indirect access to pointer locals.
-                assert(tree->OperIsScalarLocal());
-
-                var_types newType;
+                var_types newType = TYP_UNDEF;
                 if (m_allocator->m_HeapLocalToStackLocalMap.TryGetValue(lclNum, &newLclNum))
                 {
                     assert(tree->OperIs(GT_LCL_VAR)); // Must be a use.
@@ -1374,12 +1418,33 @@ void ObjectAllocator::RewriteUses()
                     }
                 }
 
-                if (lclVarDsc->lvType != newType)
+                // For local structs, retype the GC fields.
+                //
+                if (lclVarDsc->lvType == TYP_STRUCT)
+                {
+                    // We should only see spans here.
+                    //
+                    assert(lclVarDsc->IsSpan());
+                    ClassLayout* const layout    = lclVarDsc->GetLayout();
+                    ClassLayout*       newLayout = nullptr;
+
+                    if (newType == TYP_I_IMPL)
+                    {
+                        // No GC refs remain, so now a block layout
+                        newLayout = m_compiler->typGetBlkLayout(layout->GetSize());
+                        JITDUMP("Changing layout of span V%02u to block\n", lclNum);
+                        lclVarDsc->ChangeLayout(newLayout);
+                    }
+                }
+                // For locals, retype the local
+                //
+                else if (lclVarDsc->lvType != newType)
                 {
                     JITDUMP("Changing the type of V%02u from %s to %s\n", lclNum, varTypeName(lclVarDsc->lvType),
                             varTypeName(newType));
                     lclVarDsc->lvType = newType;
                 }
+
                 m_allocator->UpdateAncestorTypes(tree, &m_ancestors, newType);
 
                 if (newLclNum != BAD_VAR_NUM)