Use next sibling in SyntaxNode.GetChildPosition() if available (#66876)

Author: cston

src/Compilers/CSharp/Test/Syntax/Syntax/SyntaxListTests.cs

@@ -6,9 +6,11 @@
 
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.CSharp.Test.Utilities;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Test.Utilities;
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
@@ -305,5 +307,92 @@ public void WithLotsOfChildrenTest()
                 }
             }
         }
+
+        [Theory]
+        [CombinatorialData]
+        public void EnumerateWithManyChildren_Forward(bool trailingSeparator)
+        {
+            const int n = 200000;
+            var builder = new StringBuilder();
+            builder.Append("int[] values = new[] { ");
+            for (int i = 0; i < n; i++) builder.Append("0, ");
+            if (!trailingSeparator) builder.Append("0 ");
+            builder.AppendLine("};");
+
+            var tree = CSharpSyntaxTree.ParseText(builder.ToString());
+            // Do not descend into InitializerExpressionSyntax since that will populate SeparatedWithManyChildren._children.
+            var node = tree.GetRoot().DescendantNodes().OfType<InitializerExpressionSyntax>().First();
+
+            foreach (var child in node.ChildNodesAndTokens())
+            {
+                _ = child.ToString();
+            }
+        }
+
+        // Tests should timeout when using SeparatedWithManyChildren.GetChildPosition()
+        // instead of GetChildPositionFromEnd().
+        [WorkItem(66475, "https://github.com/dotnet/roslyn/issues/66475")]
+        [Theory]
+        [CombinatorialData]
+        public void EnumerateWithManyChildren_Reverse(bool trailingSeparator)
+        {
+            const int n = 200000;
+            var builder = new StringBuilder();
+            builder.Append("int[] values = new[] { ");
+            for (int i = 0; i < n; i++) builder.Append("0, ");
+            if (!trailingSeparator) builder.Append("0 ");
+            builder.AppendLine("};");
+
+            var tree = CSharpSyntaxTree.ParseText(builder.ToString());
+            // Do not descend into InitializerExpressionSyntax since that will populate SeparatedWithManyChildren._children.
+            var node = tree.GetRoot().DescendantNodes().OfType<InitializerExpressionSyntax>().First();
+
+            foreach (var child in node.ChildNodesAndTokens().Reverse())
+            {
+                _ = child.ToString();
+            }
+        }
+
+        [Theory]
+        [InlineData("int[] values = new[] { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 };")]
+        [InlineData("int[] values = new[] { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, };")]
+        public void EnumerateWithManyChildren_Compare(string source)
+        {
+            CSharpSyntaxTree.ParseText(source).VerifyChildNodePositions();
+
+            var builder = ArrayBuilder<SyntaxNodeOrToken>.GetInstance();
+            foreach (var node in parseAndGetInitializer(source).ChildNodesAndTokens().Reverse())
+            {
+                builder.Add(node);
+            }
+            builder.ReverseContents();
+            var childNodes1 = builder.ToImmutableAndFree();
+
+            builder = ArrayBuilder<SyntaxNodeOrToken>.GetInstance();
+            foreach (var node in parseAndGetInitializer(source).ChildNodesAndTokens())
+            {
+                builder.Add(node);
+            }
+            var childNodes2 = builder.ToImmutableAndFree();
+
+            Assert.Equal(childNodes1.Length, childNodes2.Length);
+
+            for (int i = 0; i < childNodes1.Length; i++)
+            {
+                var child1 = childNodes1[i];
+                var child2 = childNodes2[i];
+                Assert.Equal(child1.Position, child2.Position);
+                Assert.Equal(child1.EndPosition, child2.EndPosition);
+                Assert.Equal(child1.Width, child2.Width);
+                Assert.Equal(child1.FullWidth, child2.FullWidth);
+            }
+
+            static InitializerExpressionSyntax parseAndGetInitializer(string source)
+            {
+                var tree = CSharpSyntaxTree.ParseText(source);
+                // Do not descend into InitializerExpressionSyntax since that will populate SeparatedWithManyChildren._children.
+                return tree.GetRoot().DescendantNodes().OfType<InitializerExpressionSyntax>().First();
+            }
+        }
     }
 }

src/Compilers/Core/Portable/Syntax/SyntaxList.SeparatedWithManyChildren.cs

@@ -2,13 +2,11 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System;
-
 namespace Microsoft.CodeAnalysis.Syntax
 {
     internal partial class SyntaxList
     {
-        internal class SeparatedWithManyChildren : SyntaxList
+        internal sealed class SeparatedWithManyChildren : SyntaxList
         {
             private readonly ArrayElement<SyntaxNode?>[] _children;
 
@@ -39,6 +37,25 @@ internal SeparatedWithManyChildren(InternalSyntax.SyntaxList green, SyntaxNode?
 
                 return _children[i >> 1].Value;
             }
+
+            internal override int GetChildPosition(int index)
+            {
+                // If the previous sibling (ignoring separator) is not cached, but the next sibling
+                // (ignoring separator) is cached, use the next sibling to determine position.
+                int valueIndex = (index & 1) != 0 ? index - 1 : index;
+                // The check for valueIndex >= Green.SlotCount - 2 ignores the last item because the last item
+                // is a separator and separators are not cached. In those cases, when the index represents
+                // the last or next to last item, we still want to calculate the position from the end of
+                // the list rather than the start.
+                if (valueIndex > 1
+                    && GetCachedSlot(valueIndex - 2) is null
+                    && (valueIndex >= Green.SlotCount - 2 || GetCachedSlot(valueIndex + 2) is { }))
+                {
+                    return GetChildPositionFromEnd(index);
+                }
+
+                return base.GetChildPosition(index);
+            }
         }
     }
 }

src/Compilers/Core/Portable/Syntax/SyntaxList.SeparatedWithManyWeakChildren.cs

@@ -8,7 +8,7 @@ namespace Microsoft.CodeAnalysis.Syntax
 {
     internal partial class SyntaxList
     {
-        internal class SeparatedWithManyWeakChildren : SyntaxList
+        internal sealed class SeparatedWithManyWeakChildren : SyntaxList
         {
             private readonly ArrayElement<WeakReference<SyntaxNode>?>[] _children;
 

src/Compilers/Core/Portable/Syntax/SyntaxList.WithManyChildren.cs

@@ -8,7 +8,7 @@ namespace Microsoft.CodeAnalysis.Syntax
 {
     internal partial class SyntaxList
     {
-        internal class WithManyChildren : SyntaxList
+        internal sealed class WithManyChildren : SyntaxList
         {
             private readonly ArrayElement<SyntaxNode?>[] _children;
 

src/Compilers/Core/Portable/Syntax/SyntaxList.WithManyWeakChildren.cs

@@ -9,7 +9,7 @@ namespace Microsoft.CodeAnalysis.Syntax
 {
     internal partial class SyntaxList
     {
-        internal class WithManyWeakChildren : SyntaxList
+        internal sealed class WithManyWeakChildren : SyntaxList
         {
             private readonly ArrayElement<WeakReference<SyntaxNode>?>[] _children;
 

src/Compilers/Core/Portable/Syntax/SyntaxList.WithThreeChildren.cs

@@ -8,7 +8,7 @@ namespace Microsoft.CodeAnalysis.Syntax
 {
     internal partial class SyntaxList
     {
-        internal class WithThreeChildren : SyntaxList
+        internal sealed class WithThreeChildren : SyntaxList
         {
             private SyntaxNode? _child0;
             private SyntaxNode? _child1;

src/Compilers/Core/Portable/Syntax/SyntaxList.WithTwoChildren.cs

@@ -10,7 +10,7 @@ namespace Microsoft.CodeAnalysis.Syntax
 {
     internal partial class SyntaxList
     {
-        internal class WithTwoChildren : SyntaxList
+        internal sealed class WithTwoChildren : SyntaxList
         {
             private SyntaxNode? _child0;
             private SyntaxNode? _child1;

src/Compilers/Core/Portable/Syntax/SyntaxNode.cs

@@ -612,6 +612,11 @@ internal int GetChildIndex(int slot)
         /// </summary>
         internal virtual int GetChildPosition(int index)
         {
+            if (this.GetCachedSlot(index) is { } node)
+            {
+                return node.Position;
+            }
+
             int offset = 0;
             var green = this.Green;
             while (index > 0)
@@ -632,6 +637,36 @@ internal virtual int GetChildPosition(int index)
             return this.Position + offset;
         }
 
+        // Similar to GetChildPosition() but calculating based on the positions of
+        // following siblings rather than previous siblings.
+        internal int GetChildPositionFromEnd(int index)
+        {
+            if (this.GetCachedSlot(index) is { } node)
+            {
+                return node.Position;
+            }
+
+            var green = this.Green;
+            int offset = green.GetSlot(index)?.FullWidth ?? 0;
+            int slotCount = green.SlotCount;
+            while (index < slotCount - 1)
+            {
+                index++;
+                var nextSibling = this.GetCachedSlot(index);
+                if (nextSibling != null)
+                {
+                    return nextSibling.Position - offset;
+                }
+                var greenChild = green.GetSlot(index);
+                if (greenChild != null)
+                {
+                    offset += greenChild.FullWidth;
+                }
+            }
+
+            return this.EndPosition - offset;
+        }
+
         public Location GetLocation()
         {
             return this.SyntaxTree.GetLocation(this.Span);

src/Compilers/Test/Core/Compilation/CompilationExtensions.cs

@@ -27,6 +27,7 @@
 using Roslyn.Test.Utilities;
 using Roslyn.Utilities;
 using Xunit;
+using SeparatedWithManyChildren = Microsoft.CodeAnalysis.Syntax.SyntaxList.SeparatedWithManyChildren;
 
 namespace Microsoft.CodeAnalysis.Test.Utilities
 {
@@ -315,6 +316,8 @@ void checkTimeout()
                         }
                     }
                 }
+
+                tree.VerifyChildNodePositions();
             }
 
             var explicitNodeMap = new Dictionary<SyntaxNode, IOperation>();
@@ -381,6 +384,72 @@ void checkControlFlowGraph(IOperation root, ISymbol associatedSymbol)
             }
         }
 
+        internal static void VerifyChildNodePositions(this SyntaxTree tree)
+        {
+            var nodes = tree.GetRoot().DescendantNodesAndSelf();
+            foreach (var node in nodes)
+            {
+                var childNodesAndTokens = node.ChildNodesAndTokens();
+                if (childNodesAndTokens.Node is { } container)
+                {
+                    for (int i = 0; i < childNodesAndTokens.Count; i++)
+                    {
+                        if (container.GetNodeSlot(i) is SeparatedWithManyChildren separatedList)
+                        {
+                            verifyPositions(separatedList);
+                        }
+                    }
+                }
+            }
+
+            static void verifyPositions(SeparatedWithManyChildren separatedList)
+            {
+                var green = (Microsoft.CodeAnalysis.Syntax.InternalSyntax.SyntaxList)separatedList.Green;
+
+                // Calculate positions from start, using existing cache.
+                int[] positions = getPositionsFromStart(separatedList);
+
+                // Calculate positions from end, using existing cache.
+                AssertEx.Equal(positions, getPositionsFromEnd(separatedList));
+
+                // Avoid testing without caches if the number of children is large.
+                if (separatedList.SlotCount > 100)
+                {
+                    return;
+                }
+
+                // Calculate positions from start, with empty cache.
+                AssertEx.Equal(positions, getPositionsFromStart(new SeparatedWithManyChildren(green, null, separatedList.Position)));
+
+                // Calculate positions from end, with empty cache.
+                AssertEx.Equal(positions, getPositionsFromEnd(new SeparatedWithManyChildren(green, null, separatedList.Position)));
+            }
+
+            // Calculate positions from start, using any existing cache of red nodes on separated list.
+            static int[] getPositionsFromStart(SeparatedWithManyChildren separatedList)
+            {
+                int n = separatedList.SlotCount;
+                var positions = new int[n];
+                for (int i = 0; i < n; i++)
+                {
+                    positions[i] = separatedList.GetChildPosition(i);
+                }
+                return positions;
+            }
+
+            // Calculate positions from end, using any existing cache of red nodes on separated list.
+            static int[] getPositionsFromEnd(SeparatedWithManyChildren separatedList)
+            {
+                int n = separatedList.SlotCount;
+                var positions = new int[n];
+                for (int i = n - 1; i >= 0; i--)
+                {
+                    positions[i] = separatedList.GetChildPositionFromEnd(i);
+                }
+                return positions;
+            }
+        }
+
         /// <summary>
         /// The reference assembly System.Runtime.InteropServices.WindowsRuntime was removed in net5.0. This builds
         /// up <see cref="CompilationReference"/> which contains all of the well known types that were used from that

src/Compilers/VisualBasic/Test/Syntax/Syntax/SyntaxListTests.vb

@@ -3,6 +3,7 @@
 ' See the LICENSE file in the project root for more information.
 
 Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
+Imports Roslyn.Test.Utilities
 
 Namespace Microsoft.CodeAnalysis.VisualBasic.UnitTests
     Public Class SyntaxListTests
@@ -241,5 +242,54 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.UnitTests
             Next
         End Sub
 
+        <Fact>
+        Public Sub EnumerateWithManyChildren_Forward()
+            Const n = 200000
+            Dim builder As New System.Text.StringBuilder()
+            builder.AppendLine("Module M")
+            builder.AppendLine("    Sub Main")
+            builder.Append("        Dim values As Integer() = {")
+            For i = 0 To n - 1
+                builder.Append("0, ")
+            Next
+            builder.AppendLine("}")
+            builder.AppendLine("    End Sub")
+            builder.AppendLine("End Module")
+
+            Dim tree = VisualBasicSyntaxTree.ParseText(builder.ToString())
+            ' Do not descend into CollectionInitializerSyntax since that will populate SeparatedWithManyChildren._children.
+            Dim node = tree.GetRoot().DescendantNodes().OfType(Of CollectionInitializerSyntax)().First()
+
+            For Each child In node.ChildNodesAndTokens()
+                child.ToString()
+            Next
+        End Sub
+
+        ' Tests should timeout when using SeparatedWithManyChildren.GetChildPosition()
+        ' instead of GetChildPositionFromEnd().
+        <WorkItem(66475, "https://github.com/dotnet/roslyn/issues/66475")>
+        <Fact>
+        Public Sub EnumerateWithManyChildren_Reverse()
+            Const n = 200000
+            Dim builder As New System.Text.StringBuilder()
+            builder.AppendLine("Module M")
+            builder.AppendLine("    Sub Main")
+            builder.Append("        Dim values As Integer() = {")
+            For i = 0 To n - 1
+                builder.Append("0, ")
+            Next
+            builder.AppendLine("}")
+            builder.AppendLine("    End Sub")
+            builder.AppendLine("End Module")
+
+            Dim tree = VisualBasicSyntaxTree.ParseText(builder.ToString())
+            ' Do not descend into CollectionInitializerSyntax since that will populate SeparatedWithManyChildren._children.
+            Dim node = tree.GetRoot().DescendantNodes().OfType(Of CollectionInitializerSyntax)().First()
+
+            For Each child In node.ChildNodesAndTokens().Reverse()
+                child.ToString()
+            Next
+        End Sub
+
     End Class
 End Namespace