Fix System.ArgumentOutOfRangeException in FindLeafNodeAndPartner (#28382)

src/EditorFeatures/CSharpTest/EditAndContinue/SyntaxUtilitiesTests.cs

@@ -2,6 +2,7 @@
 
 using System.Linq;
 using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Xunit;
 using SyntaxUtilities = Microsoft.CodeAnalysis.CSharp.EditAndContinue.SyntaxUtilities;
 
@@ -59,5 +60,96 @@ static void Main(string[] args)
 ";
             VerifySyntaxMap(source1, source2);
         }
+
+        [Fact]
+        public void FindLeafNodeAndPartner1()
+        {
+            var leftRoot = SyntaxFactory.ParseSyntaxTree(@"
+using System;
+
+class C
+{
+    public void M()
+    {
+        if (0 == 1)
+        {
+            Console.WriteLine(0);
+        }
+    }
+}
+").GetRoot();
+            var leftPosition = leftRoot.DescendantNodes().OfType<LiteralExpressionSyntax>().ElementAt(2).SpanStart; // 0 within Console.WriteLine(0)
+            var rightRoot = SyntaxFactory.ParseSyntaxTree(@"
+using System;
+
+class C
+{
+    public void M()
+    {
+        if (0 == 1)
+        {
+            if (2 == 3)
+            {
+                Console.WriteLine(0);
+            }
+        }
+    }
+}
+").GetRoot();
+
+            SyntaxUtilities.FindLeafNodeAndPartner(leftRoot, leftPosition, rightRoot, out SyntaxNode leftNode, out SyntaxNode rightNodeOpt);
+            Assert.Equal("0", leftNode.ToString());
+            Assert.Null(rightNodeOpt);
+        }
+
+        [Fact]
+        public void FindLeafNodeAndPartner2()
+        {
+            // Check that the method does not fail even if the index of the child (4) 
+            // is greater than the count of children on the corresponding (from the upper side) node (3).
+            var leftRoot = SyntaxFactory.ParseSyntaxTree(@"
+using System;
+
+class C
+{
+    public void M()
+    {
+        if (0 == 1)
+        {
+            Console.WriteLine(0);
+            Console.WriteLine(1);
+            Console.WriteLine(2);
+            Console.WriteLine(3);
+        }
+    }
+}
+").GetRoot();
+
+            var leftPosition = leftRoot.DescendantNodes().OfType<LiteralExpressionSyntax>().ElementAt(5).SpanStart; // 3 within Console.WriteLine(3)
+            var rightRoot = SyntaxFactory.ParseSyntaxTree(@"
+using System;
+
+class C
+{
+    public void M()
+    {
+        if (0 == 1)
+        {
+            if (2 == 3)
+            {
+                Console.WriteLine(0);
+                Console.WriteLine(1);
+                Console.WriteLine(2);
+                Console.WriteLine(3);
+            }
+        }
+    }
+}
+").GetRoot();
+
+            SyntaxUtilities.FindLeafNodeAndPartner(leftRoot, leftPosition, rightRoot, out SyntaxNode leftNode, out SyntaxNode rightNodeOpt);
+            Assert.Equal("3", leftNode.ToString());
+            Assert.Null(rightNodeOpt);
+        }
     }
 }

src/EditorFeatures/VisualBasicTest/EditAndContinue/SyntaxUtilitiesTests.vb

 ' Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
 Imports SyntaxUtilities = Microsoft.CodeAnalysis.VisualBasic.EditAndContinue.SyntaxUtilities
 
 Namespace Microsoft.CodeAnalysis.Editor.VisualBasic.UnitTests.EditAndContinue
@@ -47,5 +48,84 @@ End Class
 "
             VerifySyntaxMap(source1, source2)
         End Sub
+
+        <Fact>
+        Public Sub FindLeafNodeAndPartner1()
+            Dim leftRoot = SyntaxFactory.ParseSyntaxTree("
+Imports System;
+
+Class C
+    Public Sub M()
+        If 0 = 1 Then
+            Console.WriteLine(0)
+        End If
+    End Sub
+End Class
+").GetRoot()
+            Dim leftPosition = leftRoot.DescendantNodes().OfType(Of LiteralExpressionSyntax).ElementAt(2).SpanStart '0 within Console.WriteLine(0)
+            Dim rightRoot = SyntaxFactory.ParseSyntaxTree("
+Imports System;
+
+Class C
+    Public Sub M()
+        If 0 = 1 Then
+            If 2 = 3 Then
+                Console.WriteLine(0)
+            End If
+        End If
+    End Sub
+End Class
+").GetRoot()
+
+            Dim leftNode As SyntaxNode = Nothing
+            Dim rightNodeOpt As SyntaxNode = Nothing
+            SyntaxUtilities.FindLeafNodeAndPartner(leftRoot, leftPosition, rightRoot, leftNode, rightNodeOpt)
+            Assert.Equal("0", leftNode.ToString())
+            Assert.Null(rightNodeOpt)
+        End Sub
+
+        <Fact>
+        Public Sub FindLeafNodeAndPartner2()
+            ' Check that the method does Not fail even if the index of the child (4) 
+            ' is greater than the count of children on the corresponding (from the upper side) node (3).
+            Dim leftRoot = SyntaxFactory.ParseSyntaxTree("
+Imports System;
+
+Class C
+    Public Sub M()
+        If 0 = 1 Then
+            Console.WriteLine(0)
+            Console.WriteLine(1)
+            Console.WriteLine(2)
+            Console.WriteLine(3)
+        End If
+    End Sub
+End Class
+").GetRoot()
+
+            Dim leftPosition = leftRoot.DescendantNodes().OfType(Of LiteralExpressionSyntax).ElementAt(5).SpanStart '3 within Console.WriteLine(3)
+            Dim rightRoot = SyntaxFactory.ParseSyntaxTree("
+Imports System;
+
+Class C
+    Public Sub M()
+        If 0 = 1 Then
+            If 2 = 3 Then
+                Console.WriteLine(0)
+                Console.WriteLine(1)
+                Console.WriteLine(2)
+                Console.WriteLine(3)
+            End If
+        End If
+    End Sub
+End Class
+").GetRoot()
+
+            Dim leftNode As SyntaxNode = Nothing
+            Dim rightNodeOpt As SyntaxNode = Nothing
+            SyntaxUtilities.FindLeafNodeAndPartner(leftRoot, leftPosition, rightRoot, leftNode, rightNodeOpt)
+            Assert.Equal("3", leftNode.ToString())
+            Assert.Null(rightNodeOpt)
+        End Sub
     End Class
 End Namespace

src/Features/CSharp/Portable/EditAndContinue/SyntaxUtilities.cs

@@ -114,20 +114,36 @@ public static void AssertIsBody(SyntaxNode syntax, bool allowLambda)
             Debug.Assert(false);
         }
 
-        public static void FindLeafNodeAndPartner(SyntaxNode leftRoot, int leftPosition, SyntaxNode rightRoot, out SyntaxNode leftNode, out SyntaxNode rightNode)
+        public static void FindLeafNodeAndPartner(SyntaxNode leftRoot, int leftPosition, SyntaxNode rightRoot, out SyntaxNode leftNode, out SyntaxNode rightNodeOpt)
         {
             leftNode = leftRoot;
-            rightNode = rightRoot;
+            rightNodeOpt = rightRoot;
             while (true)
             {
-                Debug.Assert(leftNode.RawKind == rightNode.RawKind);
+                if (rightNodeOpt != null && leftNode.RawKind != rightNodeOpt.RawKind)
+                {
+                    rightNodeOpt = null;
+                }
+
                 var leftChild = leftNode.ChildThatContainsPosition(leftPosition, out var childIndex);
                 if (leftChild.IsToken)
                 {
                     return;
                 }
 
-                rightNode = rightNode.ChildNodesAndTokens()[childIndex].AsNode();
+                if (rightNodeOpt != null)
+                {
+                    var rightNodeChildNodesAndTokens = rightNodeOpt.ChildNodesAndTokens();
+                    if (childIndex >= 0 && childIndex < rightNodeChildNodesAndTokens.Count)
+                    {
+                        rightNodeOpt = rightNodeChildNodesAndTokens[childIndex].AsNode();
+                    }
+                    else
+                    {
+                        rightNodeOpt = null;
+                    }
+                }
+
                 leftNode = leftChild.AsNode();
             }
         }

src/Features/VisualBasic/Portable/EditAndContinue/SyntaxUtilities.vb

@@ -67,18 +67,29 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.EditAndContinue
                                           leftPosition As Integer,
                                           rightRoot As SyntaxNode,
                                           <Out> ByRef leftNode As SyntaxNode,
-                                          <Out> ByRef rightNode As SyntaxNode)
+                                          <Out> ByRef rightNodeOpt As SyntaxNode)
             leftNode = leftRoot
-            rightNode = rightRoot
+            rightNodeOpt = rightRoot
             While True
-                Debug.Assert(leftNode.RawKind = rightNode.RawKind)
+                If rightNodeOpt IsNot Nothing AndAlso leftNode.RawKind <> rightNodeOpt.RawKind Then
+                    rightNodeOpt = Nothing
+                End If
+
                 Dim childIndex As Integer = 0
                 Dim leftChild = leftNode.ChildThatContainsPosition(leftPosition, childIndex)
                 If leftChild.IsToken Then
                     Return
                 End If
 
-                rightNode = rightNode.ChildNodesAndTokens()(childIndex).AsNode()
+                If rightNodeOpt IsNot Nothing Then
+                    Dim rightNodeChildNodesAndTokens = rightNodeOpt.ChildNodesAndTokens()
+                    If childIndex >= 0 AndAlso childIndex < rightNodeChildNodesAndTokens.Count Then
+                        rightNodeOpt = rightNodeChildNodesAndTokens(childIndex).AsNode()
+                    Else
+                        rightNodeOpt = Nothing
+                    End If
+                End If
+
                 leftNode = leftChild.AsNode()
             End While
         End Sub