JIT: handle redefined enumerators in conditional escape analysis

src/coreclr/jit/objectalloc.cpp

@@ -3588,6 +3588,30 @@ void ObjectAllocator::CheckForGuardedAllocationOrCopy(BasicBlock* block,
                     // they are properly disjoint and things will work out just fine.
                     //
                     JITDUMP("Looks like enumerator var re-use (multiple defining GDVs)\n");
+
+                    // Since we are walking in RPO, all appearances assigned to
+                    // earlier candidates have already been seen. The partition
+                    // is unsafe if this definition may reach those appearances
+                    // through a backedge, or if it is on a sibling flow path.
+                    //
+                    for (CloneInfo* const previousInfo : CloneMap::ValueIteration(&m_CloneMap))
+                    {
+                        if (previousInfo->m_local != enumeratorLocal)
+                        {
+                            continue;
+                        }
+
+                        EnumeratorVar* previousEnumeratorVar = nullptr;
+                        bool const     hasDominatingDef =
+                            previousInfo->m_appearanceMap->Lookup(enumeratorLocal, &previousEnumeratorVar) &&
+                            (previousEnumeratorVar->m_def != nullptr) &&
+                            m_compiler->m_domTree->Dominates(previousEnumeratorVar->m_def->m_block, block);
+
+                        if (block->HasFlag(BBF_BACKWARD_JUMP) || !hasDominatingDef)
+                        {
+                            previousInfo->m_hasConflictingRedefinition = true;
+                        }
+                    }
                 }
 
                 // We will query this info if we see CALL(enumeratorLocal)
@@ -3977,6 +4001,12 @@ bool ObjectAllocator::CheckCanClone(CloneInfo* info)
     JITDUMP("** Seeing if we can clone to guarantee non-escape under V%02u\n", info->m_local);
     BasicBlock* const allocBlock = info->m_allocBlock;
 
+    if (info->m_hasConflictingRedefinition)
+    {
+        JITDUMP("V%02u has a later definition that may reach its guarded uses\n", info->m_local);
+        return false;
+    }
+
     // Cloning redirects the allocation block's sole outgoing edge to the fast path,
     // so the allocation block must be a block kind that has a single target.
     //

src/coreclr/jit/objectalloc.h

@@ -109,9 +109,10 @@ struct CloneInfo : public GuardInfo
     weight_t m_profileScale = 0.0;
 
     // Status of this candidate
-    bool m_checkedCanClone = false;
-    bool m_canClone        = false;
-    bool m_willClone       = false;
+    bool m_hasConflictingRedefinition = false;
+    bool m_checkedCanClone            = false;
+    bool m_canClone                   = false;
+    bool m_willClone                  = false;
 };
 
 struct StoreInfo

src/tests/JIT/opt/ObjectStackAllocation/Runtime_134605.cs

@@ -0,0 +1,208 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using Xunit;
+
+public class Runtime_134605
+{
+    [Fact]
+    public static void Test()
+    {
+        int calls = 0;
+
+        for (int i = 0; i < 50; i++)
+        {
+            for (int j = 0; j < 1000; j++)
+            {
+                VersionedList<Item> list = new();
+                for (int k = 0; k < 6; k++)
+                {
+                    list.Add(new Item { Pending = (calls % 4 == 0) && (k == 2) });
+                }
+
+                calls++;
+                try
+                {
+                    Walk(list);
+                }
+                catch (InvalidOperationException)
+                {
+                    // Failures are checked after Tier1 compilation has settled.
+                }
+            }
+
+            Thread.Sleep(5);
+        }
+
+        Thread.Sleep(100);
+
+        int failures = 0;
+        for (int i = 0; i < 10_000; i++)
+        {
+            VersionedList<Item> list = new();
+            for (int k = 0; k < 6; k++)
+            {
+                list.Add(new Item { Pending = (calls % 4 == 0) && (k == 2) });
+            }
+
+            calls++;
+            try
+            {
+                Walk(list);
+            }
+            catch (InvalidOperationException)
+            {
+                failures++;
+            }
+        }
+
+        Assert.Equal(0, failures);
+
+        VersionedList<Item> first = CreateList();
+        VersionedList<Item> second = CreateList();
+        for (int i = 0; i < 50; i++)
+        {
+            for (int j = 0; j < 1000; j++)
+            {
+                Assert.Equal(12, CountDisjoint(first, second));
+            }
+
+            Thread.Sleep(5);
+        }
+
+        long allocatedBytesBefore = GC.GetAllocatedBytesForCurrentThread();
+        int count = CountDisjoint(first, second);
+        long allocatedBytesAfter = GC.GetAllocatedBytesForCurrentThread();
+        Assert.Equal(12, count);
+        Assert.Equal(0, allocatedBytesAfter - allocatedBytesBefore);
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static void Handle(VersionedList<Item> list, Item item)
+    {
+        item.Pending = false;
+        list.Add(new Item());
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static void Walk(VersionedList<Item> list)
+    {
+        IEnumerable<Item> items = list;
+        IEnumerator<Item> enumerator = items.GetEnumerator();
+        while (enumerator.MoveNext())
+        {
+            Item item = enumerator.Current;
+            if (item.Pending)
+            {
+                Handle(list, item);
+                enumerator = items.GetEnumerator();
+            }
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static int CountDisjoint(VersionedList<Item> first, VersionedList<Item> second)
+    {
+        int count = 0;
+        IEnumerable<Item> items = first;
+        IEnumerator<Item> enumerator = items.GetEnumerator();
+        while (enumerator.MoveNext())
+        {
+            count++;
+        }
+
+        items = second;
+        enumerator = items.GetEnumerator();
+        while (enumerator.MoveNext())
+        {
+            count++;
+        }
+
+        return count;
+    }
+
+    private static VersionedList<Item> CreateList()
+    {
+        VersionedList<Item> list = new();
+        for (int i = 0; i < 6; i++)
+        {
+            list.Add(new Item());
+        }
+
+        return list;
+    }
+
+    private sealed class VersionedList<T> : Collection<T>, IEnumerable<T>
+    {
+        private int _version;
+
+        protected override void InsertItem(int index, T item)
+        {
+            _version++;
+            base.InsertItem(index, item);
+        }
+
+        public new Enumerator GetEnumerator() => new(this);
+
+        IEnumerator<T> IEnumerable<T>.GetEnumerator() => GetEnumerator();
+
+        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
+#if STRUCT_ENUMERATOR
+        public struct Enumerator : IEnumerator<T>
+#else
+        public sealed class Enumerator : IEnumerator<T>
+#endif
+        {
+            private readonly VersionedList<T> _list;
+            private readonly int _expectedVersion;
+            private int _cursor;
+            private T _current;
+
+            internal Enumerator(VersionedList<T> list)
+            {
+                _list = list;
+                _expectedVersion = list._version;
+                _cursor = 0;
+                _current = default!;
+            }
+
+            public T Current => _current;
+
+            object? IEnumerator.Current => _current;
+
+            public bool MoveNext()
+            {
+                if (_cursor == _list.Count)
+                {
+                    return false;
+                }
+
+                if (_list._version != _expectedVersion)
+                {
+                    throw new InvalidOperationException();
+                }
+
+                _current = _list[_cursor++];
+                return true;
+            }
+
+            public void Reset() => throw new NotSupportedException();
+
+            public void Dispose()
+            {
+            }
+        }
+    }
+
+    private sealed class Item
+    {
+        public bool Pending;
+    }
+}

src/tests/JIT/opt/ObjectStackAllocation/Runtime_134605.csproj

@@ -0,0 +1,15 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <CLRTestPriority>1</CLRTestPriority>
+    <RequiresProcessIsolation>true</RequiresProcessIsolation>
+    <DebugType>None</DebugType>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+
+    <CLRTestEnvironmentVariable Include="DOTNET_TieredCompilation" Value="1" />
+    <CLRTestEnvironmentVariable Include="DOTNET_TieredPGO" Value="1" />
+    <CLRTestEnvironmentVariable Include="DOTNET_TC_QuickJitForLoops" Value="1" />
+  </ItemGroup>
+</Project>

src/tests/JIT/opt/ObjectStackAllocation/Runtime_134605_Struct.csproj

@@ -0,0 +1,16 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <CLRTestPriority>1</CLRTestPriority>
+    <RequiresProcessIsolation>true</RequiresProcessIsolation>
+    <DebugType>None</DebugType>
+    <Optimize>True</Optimize>
+    <DefineConstants>$(DefineConstants);STRUCT_ENUMERATOR</DefineConstants>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="Runtime_134605.cs" />
+
+    <CLRTestEnvironmentVariable Include="DOTNET_TieredCompilation" Value="1" />
+    <CLRTestEnvironmentVariable Include="DOTNET_TieredPGO" Value="1" />
+    <CLRTestEnvironmentVariable Include="DOTNET_TC_QuickJitForLoops" Value="1" />
+  </ItemGroup>
+</Project>