[WIP] Fix LargeArrayBuilder.CopyTo returning incorrect end-of-copy position

src/Common/src/System/Collections/Generic/LargeArrayBuilder.cs

@@ -46,6 +46,20 @@ internal CopyPosition(int row, int column)
         /// Gets a string suitable for display in the debugger.
         /// </summary>
         private string DebuggerDisplay => $"[{Row}, {Column}]";
+
+        /// <summary>
+        /// If this position is at the end of the current buffer, returns the position
+        /// at the start of the next buffer. Otherwise, returns this position.
+        /// </summary>
+        /// <param name="endColumn">The length of the current buffer.</param>
+        public CopyPosition Normalize(int endColumn)
+        {
+            Debug.Assert(Column <= endColumn);
+
+            return Column == endColumn ?
+                new CopyPosition(Row + 1, 0) :
+                this;
+        }
     }
 
     /// <summary>
@@ -175,7 +189,7 @@ public void CopyTo(T[] array, int arrayIndex, int count)
             for (int i = 0; count > 0; i++)
             {
                 // Find the buffer we're copying from.
-                T[] buffer = GetBuffer(index: i);
+                T[] buffer = GetBuffer(i);
 
                 // Copy until we satisfy count, or we reach the end of the buffer.
                 int toCopy = Math.Min(count, buffer.Length);
@@ -197,6 +211,24 @@ public void CopyTo(T[] array, int arrayIndex, int count)
         /// <returns>The position in this builder that was copied up to.</returns>
         public CopyPosition CopyTo(CopyPosition position, T[] array, int arrayIndex, int count)
         {
+            int CopyToCore(T[] sourceBuffer, int sourceIndex)
+            {
+                Debug.Assert(sourceBuffer.Length > sourceIndex);
+
+                // Copy until we satisfy `count` or reach the end of the current buffer.
+                int copyCount = Math.Min(sourceBuffer.Length - sourceIndex, count);
+
+                if (copyCount > 0)
+                {
+                    Array.Copy(sourceBuffer, sourceIndex, array, arrayIndex, copyCount);
+
+                    arrayIndex += copyCount;
+                    count -= copyCount;
+                }
+
+                return copyCount;
+            }
+
             Debug.Assert(arrayIndex >= 0);
             Debug.Assert(count >= 0 && count <= Count);
             Debug.Assert(array?.Length - arrayIndex >= count);
@@ -207,34 +239,30 @@ public CopyPosition CopyTo(CopyPosition position, T[] array, int arrayIndex, int
             /*
              * Visual representation:
              * 
-             *       C0   C1   C2 ..  C31 ..   C63
-             * R0:  [0]  [1]  [2] .. [31]
-             * R1: [32] [33] [34] .. [63]
-             * R2: [64] [65] [66] .. [95] .. [127]
+             *       C0   C1   C2 ..   C7 ..  C15
+             * R0:  [0]  [1]  [2] ..  [7]
+             * R1:  [8]  [9] [10] .. [15]
+             * R2: [16] [17] [18] .. [23] .. [31]
              */
 
             int row = position.Row;
             int column = position.Column;
 
-            for (; count > 0; row++, column = 0)
+            T[] buffer = GetBuffer(row);
+            int copied = CopyToCore(buffer, column);
+            if (count == 0)
             {
-                T[] buffer = GetBuffer(index: row);
-
-                // During this iteration, copy until we satisfy `count` or reach the
-                // end of the current buffer.
-                int copyCount = Math.Min(buffer.Length, count);
-
-                if (copyCount > 0)
-                {
-                    Array.Copy(buffer, column, array, arrayIndex, copyCount);
+                return new CopyPosition(row, column + copied).Normalize(buffer.Length);
+            }
 
-                    arrayIndex += copyCount;
-                    count -= copyCount;
-                    column += copyCount;
-                }
+            do
+            {
+                buffer = GetBuffer(++row);
+                copied = CopyToCore(buffer, 0);
             }
+            while (count > 0);
 
-            return new CopyPosition(row: row, column: column);
+            return new CopyPosition(row, copied).Normalize(buffer.Length);
         }
 
         /// <summary>

src/System.Linq/tests/ConcatTests.cs

@@ -421,5 +421,18 @@ public void GetEnumerableOfConcatCollectionChainFollowedByEnumerableNodeShouldBe
                 Assert.Equal(0xf00, en.Current);
             }
         }
+
+        [Fact]
+        public void CollectionInterleavedWithLazyEnumerables_ToArray_Regression()
+        {
+            var results = new TestEnumerable<int>(new[] { 1 })
+                .Concat(new[] { 2 })
+                .Concat(new TestEnumerable<int>(new[] { 3 }))
+                // Do not omit this ToArray()! There was a previous bug where the ToArray() implementation
+                // was incorrect, while iterating with foreach produced the correct results.
+                .ToArray();
+
+            Assert.Equal(new[] { 1, 2, 3 }, results);
+        }
     }
 }

src/System.Linq/tests/SelectManyTests.cs

@@ -477,5 +477,19 @@ public void ThrowOverflowExceptionOnConstituentLargeCounts(int[] counts)
             IEnumerable<int> iterator = counts.SelectMany(c => Enumerable.Range(1, c));
             Assert.Throws<OverflowException>(() => iterator.Count());
         }
+
+        [Fact]
+        public void CollectionInterleavedWithLazyEnumerables_ToArray_Regression()
+        {
+            var results = new[] { 4, 5, 6 }
+                .SelectMany(i => i == 5 ?
+                    (IEnumerable<int>)new List<int>() { i } :
+                    new TestEnumerable<int>(new int[] { i }))
+                // Do not omit this ToArray()! There was a previous bug where the ToArray() implementation
+                // was incorrect, while iterating with foreach produced the correct results.
+                .ToArray();
+
+            Assert.Equal(new[] { 4, 5, 6 }, results);
+        }
     }
 }