JIT: Add Statement::m_treeListEnd (#81031)

Before this change statements have only one field for the "tree list"
even though bidirectional iteration is supported after nodes have been
linked. This worked fine before the locals tree list existed since the
"root node" was always the last node.

This does not work for the locals tree list since the "root node" is not
part of the list, meaning that we need somewhere else to store the
'last' node. Before this PR the assumption was that the root node was
_never_ part of the locals linked list, so we could use the
gtNext/gtPrev fields of the root node. However, the added test case
shows that in fact it is possible we end up with a top level local.

This PR fixes the problem by adding a Statement::m_treeListEnd field
that can keep the last node of the locals linked list.  While this takes
some memory, it seems like the most maintainable way to resolve the
problem. I experimented with making the linked list circular or by
allocating a new stand-in node when necessary but eventually I ended up
here.

Fix #81018

Author: jakobbotsch

src/coreclr/jit/compiler.cpp

@@ -4609,15 +4609,15 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     //
     lvaRefCountState = RCS_EARLY;
 
-    // Figure out what locals are address-taken.
-    //
-    DoPhase(this, PHASE_STR_ADRLCL, &Compiler::fgMarkAddressExposedLocals);
-
     if (opts.OptimizationEnabled())
     {
         fgNodeThreading = NodeThreading::AllLocals;
     }
 
+    // Figure out what locals are address-taken.
+    //
+    DoPhase(this, PHASE_STR_ADRLCL, &Compiler::fgMarkAddressExposedLocals);
+
     // Do an early pass of liveness for forward sub and morph. This data is
     // valid until after morph.
     //

src/coreclr/jit/compiler.h

@@ -4531,8 +4531,8 @@ class Compiler
     // doubly linked lists during certain phases of the compilation.
     // - Local morph threads all locals to be used for early liveness and
     //   forward sub when optimizing. This is kept valid until after forward sub.
-    //   The first local is kept in Statement::GetRootNode()->gtNext and the last
-    //   local in Statement::GetRootNode()->gtPrev. fgSequenceLocals can be used
+    //   The first local is kept in Statement::GetTreeList() and the last
+    //   local in Statement::GetTreeListEnd(). fgSequenceLocals can be used
     //   to (re-)sequence a statement into this form, and
     //   Statement::LocalsTreeList for range-based iteration. The order must
     //   match tree order.

src/coreclr/jit/fgdiagnostic.cpp

@@ -3412,15 +3412,25 @@ void Compiler::fgDebugCheckLinkedLocals()
     {
         for (Statement* stmt : block->Statements())
         {
-            GenTree* first = stmt->GetRootNode()->gtNext;
+            GenTree* first = stmt->GetTreeList();
             CheckDoublyLinkedList<GenTree, &GenTree::gtPrev, &GenTree::gtNext>(first);
 
             seq.Sequence(stmt);
 
             ArrayStack<GenTree*>* expected = seq.GetSequence();
 
-            bool success   = true;
-            int  nodeIndex = 0;
+            bool success = true;
+
+            if (expected->Height() > 0)
+            {
+                success &= (stmt->GetTreeList() == expected->Bottom(0)) && (stmt->GetTreeListEnd() == expected->Top(0));
+            }
+            else
+            {
+                success &= (stmt->GetTreeList() == nullptr) && (stmt->GetTreeListEnd() == nullptr);
+            }
+
+            int nodeIndex = 0;
             for (GenTree* cur = first; cur != nullptr; cur = cur->gtNext)
             {
                 success &= cur->OperIsLocal() || cur->OperIsLocalAddr();

src/coreclr/jit/gentree.cpp

@@ -562,7 +562,7 @@ void GenTree::DumpNodeSizes(FILE* fp)
 //
 LocalsGenTreeList::iterator LocalsGenTreeList::begin() const
 {
-    GenTree* first = m_stmt->GetRootNode()->gtNext;
+    GenTree* first = m_stmt->GetTreeList();
     assert((first == nullptr) || first->OperIsLocal() || first->OperIsLocalAddr());
     return iterator(static_cast<GenTreeLclVarCommon*>(first));
 }
@@ -580,8 +580,8 @@ GenTree** LocalsGenTreeList::GetForwardEdge(GenTreeLclVarCommon* node)
 {
     if (node->gtPrev == nullptr)
     {
-        assert(m_stmt->GetRootNode()->gtNext == node);
-        return &m_stmt->GetRootNode()->gtNext;
+        assert(m_stmt->GetTreeList() == node);
+        return m_stmt->GetTreeListPointer();
     }
     else
     {
@@ -603,8 +603,8 @@ GenTree** LocalsGenTreeList::GetBackwardEdge(GenTreeLclVarCommon* node)
 {
     if (node->gtNext == nullptr)
     {
-        assert(m_stmt->GetRootNode()->gtPrev == node);
-        return &m_stmt->GetRootNode()->gtPrev;
+        assert(m_stmt->GetTreeListEnd() == node);
+        return m_stmt->GetTreeListEndPointer();
     }
     else
     {

src/coreclr/jit/gentree.h

@@ -7449,6 +7449,7 @@ struct Statement
     Statement(GenTree* expr DEBUGARG(unsigned stmtID))
         : m_rootNode(expr)
         , m_treeList(nullptr)
+        , m_treeListEnd(nullptr)
         , m_next(nullptr)
         , m_prev(nullptr)
 #ifdef DEBUG
@@ -7478,11 +7479,31 @@ struct Statement
         return m_treeList;
     }
 
+    GenTree** GetTreeListPointer()
+    {
+        return &m_treeList;
+    }
+
     void SetTreeList(GenTree* treeHead)
     {
         m_treeList = treeHead;
     }
 
+    GenTree* GetTreeListEnd() const
+    {
+        return m_treeListEnd;
+    }
+
+    GenTree** GetTreeListEndPointer()
+    {
+        return &m_treeListEnd;
+    }
+
+    void SetTreeListEnd(GenTree* end)
+    {
+        m_treeListEnd = end;
+    }
+
     GenTreeList       TreeList() const;
     LocalsGenTreeList LocalsTreeList();
 
@@ -7559,6 +7580,12 @@ struct Statement
     // The value is `nullptr` until we have set the sequencing of the nodes.
     GenTree* m_treeList;
 
+    // The tree list tail. Only valid when locals are linked (fgNodeThreading
+    // == AllLocals), in which case this is the last local.
+    // When all nodes are linked (fgNodeThreading == AllTrees), m_rootNode
+    // should be considered the last node.
+    GenTree* m_treeListEnd;
+
     // The statement nodes are doubly-linked. The first statement node in a block points
     // to the last node in the block via its `m_prev` link. Note that the last statement node
     // does not point to the first: it has `m_next == nullptr`; that is, the list is not fully circular.

src/coreclr/jit/lclmorph.cpp

@@ -5,7 +5,6 @@
 
 class LocalSequencer final : public GenTreeVisitor<LocalSequencer>
 {
-    GenTree* m_rootNode;
     GenTree* m_prevNode;
 
 public:
@@ -15,7 +14,7 @@ class LocalSequencer final : public GenTreeVisitor<LocalSequencer>
         UseExecutionOrder = true,
     };
 
-    LocalSequencer(Compiler* comp) : GenTreeVisitor(comp), m_rootNode(nullptr), m_prevNode(nullptr)
+    LocalSequencer(Compiler* comp) : GenTreeVisitor(comp), m_prevNode(nullptr)
     {
     }
 
@@ -30,12 +29,10 @@ class LocalSequencer final : public GenTreeVisitor<LocalSequencer>
     {
         // We use the root node as a 'sentinel' node that will keep the head
         // and tail of the sequenced list.
-        m_rootNode = stmt->GetRootNode();
-        assert(!m_rootNode->OperIsLocal() && !m_rootNode->OperIsLocalAddr());
-
-        m_rootNode->gtPrev = nullptr;
-        m_rootNode->gtNext = nullptr;
-        m_prevNode         = m_rootNode;
+        GenTree* rootNode = stmt->GetRootNode();
+        rootNode->gtPrev  = nullptr;
+        rootNode->gtNext  = nullptr;
+        m_prevNode        = rootNode;
     }
 
     //-------------------------------------------------------------------
@@ -47,26 +44,35 @@ class LocalSequencer final : public GenTreeVisitor<LocalSequencer>
     //
     void Finish(Statement* stmt)
     {
-        assert(stmt->GetRootNode() == m_rootNode);
+        GenTree* rootNode = stmt->GetRootNode();
+
+        GenTree* firstNode = rootNode->gtNext;
+        GenTree* lastNode  = m_prevNode;
 
-        GenTree* firstNode = m_rootNode->gtNext;
         if (firstNode == nullptr)
         {
-            assert(m_rootNode->gtPrev == nullptr);
+            lastNode = nullptr;
         }
         else
         {
-            GenTree* lastNode = m_prevNode;
-
-            // We only sequence leaf nodes that we shouldn't see as standalone
-            // statements here.
-            assert(m_rootNode != firstNode);
-            assert((m_rootNode->gtPrev == nullptr) && (lastNode->gtNext == nullptr));
+            // In the rare case that the root node becomes part of the linked
+            // list (i.e. top level local) we get a circular linked list here.
+            if (firstNode == rootNode)
+            {
+                assert(firstNode == lastNode);
+                lastNode->gtNext = nullptr;
+            }
+            else
+            {
+                assert(lastNode->gtNext == nullptr);
+                assert(lastNode->OperIsLocal() || lastNode->OperIsLocalAddr());
+            }
 
-            assert(lastNode->OperIsLocal() || lastNode->OperIsLocalAddr());
-            firstNode->gtPrev  = nullptr;
-            m_rootNode->gtPrev = lastNode;
+            firstNode->gtPrev = nullptr;
         }
+
+        stmt->SetTreeList(firstNode);
+        stmt->SetTreeListEnd(lastNode);
     }
 
     fgWalkResult PostOrderVisit(GenTree** use, GenTree* user)
@@ -202,7 +208,7 @@ class LocalSequencer final : public GenTreeVisitor<LocalSequencer>
 //
 void Compiler::fgSequenceLocals(Statement* stmt)
 {
-    assert((fgNodeThreading == NodeThreading::AllLocals) || (mostRecentlyActivePhase == PHASE_STR_ADRLCL));
+    assert(fgNodeThreading == NodeThreading::AllLocals);
     LocalSequencer seq(this);
     seq.Sequence(stmt);
 }

src/coreclr/jit/liveness.cpp

@@ -498,7 +498,7 @@ void Compiler::fgPerBlockLocalVarLiveness()
                     {
                         // Assigned local should be the very last local.
                         assert((dst == nullptr) ||
-                               ((stmt->GetRootNode()->gtPrev == dst) && ((dst->gtFlags & GTF_VAR_DEF) != 0)));
+                               ((stmt->GetTreeListEnd() == dst) && ((dst->gtFlags & GTF_VAR_DEF) != 0)));
 
                         // Conservatively ignore defs that may be conditional
                         // but would otherwise still interfere with the
@@ -1825,7 +1825,7 @@ GenTree* Compiler::fgTryRemoveDeadStoreEarly(Statement* stmt, GenTreeLclVarCommo
 
     JITDUMP("Store [%06u] is dead", dspTreeID(stmt->GetRootNode()));
     // The def ought to be the last thing.
-    assert(stmt->GetRootNode()->gtPrev == cur);
+    assert(stmt->GetTreeListEnd() == cur);
 
     GenTree* sideEffects = nullptr;
     gtExtractSideEffList(stmt->GetRootNode()->gtGetOp2(), &sideEffects);
@@ -1844,7 +1844,7 @@ GenTree* Compiler::fgTryRemoveDeadStoreEarly(Statement* stmt, GenTreeLclVarCommo
         DISPTREE(sideEffects);
         JITDUMP("\n");
         // continue at tail of the side effects
-        return stmt->GetRootNode()->gtPrev;
+        return stmt->GetTreeListEnd();
     }
 }
 
@@ -2820,7 +2820,7 @@ void Compiler::fgInterBlockLocalVarLiveness()
 
                 if (qmark != nullptr)
                 {
-                    for (GenTree* cur = stmt->GetRootNode()->gtPrev; cur != nullptr;)
+                    for (GenTree* cur = stmt->GetTreeListEnd(); cur != nullptr;)
                     {
                         assert(cur->OperIsLocal() || cur->OperIsLocalAddr());
                         bool isDef = ((cur->gtFlags & GTF_VAR_DEF) != 0) && ((cur->gtFlags & GTF_VAR_USEASG) == 0);
@@ -2846,7 +2846,7 @@ void Compiler::fgInterBlockLocalVarLiveness()
                 }
                 else
                 {
-                    for (GenTree* cur = stmt->GetRootNode()->gtPrev; cur != nullptr;)
+                    for (GenTree* cur = stmt->GetTreeListEnd(); cur != nullptr;)
                     {
                         assert(cur->OperIsLocal() || cur->OperIsLocalAddr());
                         if (!fgComputeLifeLocal(life, keepAliveVars, cur))

src/tests/JIT/Regression/JitBlue/Runtime_81018/Runtime_81018.cs

@@ -0,0 +1,36 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+// Generated by Fuzzlyn v1.5 on 2023-01-22 16:00:16
+// Run on Arm64 Windows
+// Seed: 17286164302317655577
+// Reduced from 117.6 KiB to 0.4 KiB in 00:02:13
+// Hits JIT assert in Release:
+// Assertion failed '!m_rootNode->OperIsLocal() && !m_rootNode->OperIsLocalAddr()' in 'Program:Main(Fuzzlyn.ExecutionServer.IRuntime)' during 'Morph - Structs/AddrExp' (IL size 83; hash 0xade6b36b; FullOpts)
+// 
+//     File: D:\a\_work\1\s\src\coreclr\jit\lclmorph.cpp Line: 34
+// 
+public interface I0
+{
+}
+
+public struct S0 : I0
+{
+    public sbyte F0;
+    public S0 M17(I0 arg0, ulong arg1)
+    {
+        return this;
+    }
+}
+
+public class Runtime_81018
+{
+    public static ulong s_2;
+    public static int Main()
+    {
+        var vr6 = new S0();
+        var vr7 = new S0();
+        new S0().M17(new S0().M17(vr7, 0).M17(vr6, s_2), s_2);
+        return 100;
+    }
+}

src/tests/JIT/Regression/JitBlue/Runtime_81018/Runtime_81018.csproj

@@ -0,0 +1,9 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>