Fix skipping last segment logic in ROS.TryGet (#28266)

Author: pakrym

src/System.Memory/src/System/Buffers/ReadOnlySequence_helpers.cs

@@ -41,12 +41,18 @@ internal bool TryGetBuffer(in SequencePosition start, in SequencePosition end, o
                 Debug.Assert(startObject is ReadOnlySequenceSegment<T>);
                 ReadOnlySequenceSegment<T> startSegment = Unsafe.As<ReadOnlySequenceSegment<T>>(startObject);
 
+                data = startSegment.Memory;
+
                 if (startSegment != endObject)
                 {
-                    next = GetBufferCrossSegment(startIndex, endObject, ref startSegment, ref length);
-                }
+                    ReadOnlySequenceSegment<T> nextSegment = startSegment.Next;
 
-                data = startSegment.Memory;
+                    if (nextSegment == null && startSegment != endObject)
+                        ThrowHelper.ThrowInvalidOperationException_EndPositionNotReached();
+
+                    next = new SequencePosition(nextSegment, 0);
+                    length = data.Length - startIndex;
+                }
             }
             else if (type == SequenceType.Array)
             {
@@ -135,47 +141,6 @@ internal ReadOnlyMemory<T> GetFirstBuffer(in SequencePosition start, in Sequence
             return memory.Slice(startIndex, length);
         }
 
-        private static SequencePosition GetBufferCrossSegment(int startIndex, object endObject, ref ReadOnlySequenceSegment<T> startSegment, ref int length)
-        {
-            Debug.Assert(startSegment != null);
-
-            ReadOnlySequenceSegment<T> nextSegment = startSegment.Next;
-            int currentLength = startSegment.Memory.Length - startIndex;
-
-            while (currentLength == 0 && nextSegment != endObject && nextSegment != null)
-            {
-                // startSegment is at the end of a segment; move it to start of the next.
-                // However, skip any empty segments, else the caller will immedately run out of data and call back
-                startSegment = nextSegment;
-                nextSegment = nextSegment.Next;
-                currentLength = startSegment.Memory.Length;
-            }
-
-            length = currentLength;
-            while (nextSegment != null && nextSegment.Memory.Length == 0)
-            {
-                // nextSegment is an empty segment, skip until one is found with data.
-                if (nextSegment == endObject)
-                {
-                    // Reached the end and no segment with data found,
-                    // set nextSegment to null as if there were no following segments. 
-                    endObject = nextSegment = null;
-                    break;
-                }
-
-                nextSegment = nextSegment.Next;
-            }
-
-            if (nextSegment != null)
-            {
-                return new SequencePosition(nextSegment, 0);
-            }
-            else if (endObject != null)
-                ThrowHelper.ThrowInvalidOperationException_EndPositionNotReached();
-
-            return default;
-        }
-
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal SequencePosition Seek(in SequencePosition start, in SequencePosition end, int count) => Seek(start, end, (long)count);
 

src/System.Memory/tests/ReadOnlyBuffer/ReadOnlySequenceTests.Common.cs

@@ -163,153 +163,196 @@ public void SeekSkipsEmptySegments()
         }
 
         [Fact]
-        public void SeekEmptySkipDoesNotCrossPastEnd()
+        public void TryGetReturnsEmptySegments()
         {
-            var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
+            var bufferSegment1 = new BufferSegment<byte>(new byte[0]);
             BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[0]);
             BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[0]);
+            BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[0]);
+
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment3, 0);
+
+            var start = buffer.Start;
+            Assert.True(buffer.TryGet(ref start, out var memory));
+            Assert.Equal(0, memory.Length);
+            Assert.True(buffer.TryGet(ref start, out memory));
+            Assert.Equal(0, memory.Length);
+            Assert.True(buffer.TryGet(ref start, out memory));
+            Assert.Equal(0, memory.Length);
+            Assert.False(buffer.TryGet(ref start, out memory));
+        }
+
+        [Fact]
+        public void TryGetStopsAtEnd()
+        {
+            var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
+            BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[100]);
+            BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[100]);
             BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[100]);
+            BufferSegment<byte> bufferSegment5 = bufferSegment4.Append(new byte[100]);
+
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment3, 100);
+
+            var start = buffer.Start;
+            Assert.True(buffer.TryGet(ref start, out var memory));
+            Assert.Equal(100, memory.Length);
+            Assert.True(buffer.TryGet(ref start, out memory));
+            Assert.Equal(100, memory.Length);
+            Assert.True(buffer.TryGet(ref start, out memory));
+            Assert.Equal(100, memory.Length);
+            Assert.False(buffer.TryGet(ref start, out memory));
+        }
 
-            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
+        [Fact]
+        public void TryGetStopsAtEndWhenEndIsLastByteOfFull()
+        {
+            var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
+            BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[100]);
 
-            SequencePosition c1 = buffer.GetPosition(100);
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment1, 100);
 
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment2, c1.GetObject());
+            var start = buffer.Start;
+            Assert.True(buffer.TryGet(ref start, out var memory));
+            Assert.Equal(100, memory.Length);
+            Assert.False(buffer.TryGet(ref start, out memory));
+        }
 
-            c1 = buffer.GetPosition(100, buffer.Start);
+        [Fact]
+        public void TryGetStopsAtEndWhenEndIsFirstByteOfFull()
+        {
+            var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
+            BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[100]);
 
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment2, c1.GetObject());
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
 
-            // Go out of bounds for segment
-            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(150, buffer.Start));
-            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(250, buffer.Start));
+            var start = buffer.Start;
+            Assert.True(buffer.TryGet(ref start, out var memory));
+            Assert.Equal(100, memory.Length);
+            Assert.True(buffer.TryGet(ref start, out memory));
+            Assert.Equal(0, memory.Length);
+            Assert.False(buffer.TryGet(ref start, out memory));
         }
 
         [Fact]
-        public void SeekEmptySkipDoesNotCrossPastEndWithExtraChainedBlocks()
+        public void TryGetStopsAtEndWhenEndIsFirstByteOfEmpty()
         {
             var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
             BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[100]);
-            BufferSegment<byte> bufferSegment5 = bufferSegment4.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment6 = bufferSegment5.Append(new byte[100]);
 
             var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
 
-            SequencePosition c1 = buffer.GetPosition(100);
+            var start = buffer.Start;
+            Assert.True(buffer.TryGet(ref start, out var memory));
+            Assert.Equal(100, memory.Length);
+            Assert.True(buffer.TryGet(ref start, out memory));
+            Assert.Equal(0, memory.Length);
+            Assert.False(buffer.TryGet(ref start, out memory));
+        }
 
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment2, c1.GetObject());
+        [Fact]
+        public void EnumerableStopsAtEndWhenEndIsLastByteOfFull()
+        {
+            var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
+            BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[100]);
 
-            c1 = buffer.GetPosition(100, buffer.Start);
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment1, 100);
 
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment2, c1.GetObject());
+            List<int> sizes = new List<int>();
+            foreach (var memory in buffer)
+            {
+                sizes.Add(memory.Length);
+            }
 
-            // Go out of bounds for segment
-            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(150, buffer.Start));
-            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(250, buffer.Start));
+            Assert.Equal(1, sizes.Count);
+            Assert.Equal(new [] { 100 }, sizes);
         }
 
         [Fact]
-        public void TryGetSkipsEmptyForNextUsingGetPositionWithOffset()
+        public void EnumerableStopsAtEndWhenEndIsFirstByteOfFull()
         {
             var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
-            BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[100]);
-
-            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment4, 100);
-            SequencePosition c1 = buffer.GetPosition(0);
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment1, c1.GetObject());
-
-            ReadOnlyMemory<byte> data;
-            Assert.True(buffer.TryGet(ref c1, out data, advance: true));
-
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment4, c1.GetObject());
+            BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[100]);
 
-            Assert.True(buffer.TryGet(ref c1, out data, advance: true));
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
 
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(null, c1.GetObject());
+            List<int> sizes = new List<int>();
+            foreach (var memory in buffer)
+            {
+                sizes.Add(memory.Length);
+            }
 
-            Assert.False(buffer.TryGet(ref c1, out data, advance: true));
+            Assert.Equal(2, sizes.Count);
+            Assert.Equal(new [] { 100, 0 }, sizes);
         }
 
         [Fact]
-        public void TryGetSkipsEmptyForNextUsingGetPositionWithOffsetAndOrigin()
+        public void EnumerableStopsAtEndWhenEndIsFirstByteOfEmpty()
         {
             var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
             BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[0]);
-            BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[100]);
-
-            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment4, 100);
-            SequencePosition c1 = buffer.GetPosition(0, buffer.Start);
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment1, c1.GetObject());
-
-            ReadOnlyMemory<byte> data;
-            Assert.True(buffer.TryGet(ref c1, out data, advance: true));
-
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment4, c1.GetObject());
 
-            Assert.True(buffer.TryGet(ref c1, out data, advance: true));
+            var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
 
-            Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(null, c1.GetObject());
+            List<int> sizes = new List<int>();
+            foreach (var memory in buffer)
+            {
+                sizes.Add(memory.Length);
+            }
 
-            Assert.False(buffer.TryGet(ref c1, out data, advance: true));
+            Assert.Equal(2, sizes.Count);
+            Assert.Equal(new [] { 100, 0 }, sizes);
         }
 
         [Fact]
-        public void TryGetSkipDoesNotCrossPastEndUsingGetPositionWithOffset()
+        public void SeekEmptySkipDoesNotCrossPastEnd()
         {
             var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
             BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[0]);
             BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[0]);
             BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[100]);
 
             var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
-            SequencePosition c1 = buffer.GetPosition(0);
+
+            SequencePosition c1 = buffer.GetPosition(100);
+
             Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment1, c1.GetObject());
+            Assert.Equal(bufferSegment2, c1.GetObject());
 
-            ReadOnlyMemory<byte> data;
-            Assert.True(buffer.TryGet(ref c1, out data, advance: true));
+            c1 = buffer.GetPosition(100, buffer.Start);
 
             Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(null, c1.GetObject());
+            Assert.Equal(bufferSegment2, c1.GetObject());
 
-            Assert.False(buffer.TryGet(ref c1, out data, advance: true));
+            // Go out of bounds for segment
+            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(150, buffer.Start));
+            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(250, buffer.Start));
         }
 
         [Fact]
-        public void TryGetSkipDoesNotCrossPastEndUsingGetPositionWithOffsetAndOrigin()
+        public void SeekEmptySkipDoesNotCrossPastEndWithExtraChainedBlocks()
         {
             var bufferSegment1 = new BufferSegment<byte>(new byte[100]);
             BufferSegment<byte> bufferSegment2 = bufferSegment1.Append(new byte[0]);
             BufferSegment<byte> bufferSegment3 = bufferSegment2.Append(new byte[0]);
             BufferSegment<byte> bufferSegment4 = bufferSegment3.Append(new byte[100]);
+            BufferSegment<byte> bufferSegment5 = bufferSegment4.Append(new byte[0]);
+            BufferSegment<byte> bufferSegment6 = bufferSegment5.Append(new byte[100]);
 
             var buffer = new ReadOnlySequence<byte>(bufferSegment1, 0, bufferSegment2, 0);
-            SequencePosition c1 = buffer.GetPosition(0, buffer.Start);
+
+            SequencePosition c1 = buffer.GetPosition(100);
+
             Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(bufferSegment1, c1.GetObject());
+            Assert.Equal(bufferSegment2, c1.GetObject());
 
-            ReadOnlyMemory<byte> data;
-            Assert.True(buffer.TryGet(ref c1, out data, advance: true));
+            c1 = buffer.GetPosition(100, buffer.Start);
 
             Assert.Equal(0, c1.GetInteger());
-            Assert.Equal(null, c1.GetObject());
+            Assert.Equal(bufferSegment2, c1.GetObject());
 
-            Assert.False(buffer.TryGet(ref c1, out data, advance: true));
+            // Go out of bounds for segment
+            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(150, buffer.Start));
+            Assert.Throws<ArgumentOutOfRangeException>(() => c1 = buffer.GetPosition(250, buffer.Start));
         }
 
         [Fact]