Extract Consume from CopyTo

Author: benaadams

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/MemoryPoolIterator.cs

@@ -870,11 +870,16 @@ public int GetLength(MemoryPoolIterator end)
 
         public MemoryPoolIterator CopyTo(byte[] array, int offset, int count, out int actual)
         {
-            if (IsDefault)
+            // Note: Ensure for any changes in this function Consume is changed to match
+            if (count == 0 || IsDefault)
             {
                 actual = 0;
                 return this;
             }
+            if (array == null)
+            {
+                return Consume(count, out actual);
+            }
 
             var block = _block;
             var index = _index;
@@ -890,27 +895,18 @@ public MemoryPoolIterator CopyTo(byte[] array, int offset, int count, out int ac
                 if (remaining <= following)
                 {
                     actual = count;
-                    if (array != null)
-                    {
-                        Buffer.BlockCopy(block.Array, index, array, offset, remaining);
-                    }
+                    Buffer.BlockCopy(block.Array, index, array, offset, remaining);
                     return new MemoryPoolIterator(block, index + remaining);
                 }
                 else if (wasLastBlock)
                 {
                     actual = count - remaining + following;
-                    if (array != null)
-                    {
-                        Buffer.BlockCopy(block.Array, index, array, offset, following);
-                    }
+                    Buffer.BlockCopy(block.Array, index, array, offset, following);
                     return new MemoryPoolIterator(block, index + following);
                 }
                 else
                 {
-                    if (array != null)
-                    {
-                        Buffer.BlockCopy(block.Array, index, array, offset, following);
-                    }
+                    Buffer.BlockCopy(block.Array, index, array, offset, following);
                     offset += following;
                     remaining -= following;
                     block = block.Next;
@@ -919,6 +915,34 @@ public MemoryPoolIterator CopyTo(byte[] array, int offset, int count, out int ac
             }
         }
 
+        private MemoryPoolIterator Consume(int count, out int actual)
+        {
+            var block = _block;
+            var index = _index;
+            var remaining = count;
+            while (true)
+            {
+                var wasLastBlock = block.Next == null;
+                var following = block.End - index;
+                if (remaining <= following)
+                {
+                    actual = count;
+                    return new MemoryPoolIterator(block, index + remaining);
+                }
+                else if (wasLastBlock)
+                {
+                    actual = count - remaining + following;
+                    return new MemoryPoolIterator(block, index + following);
+                }
+                else
+                {
+                    remaining -= following;
+                    block = block.Next;
+                    index = block.Start;
+                }
+            }
+        }
+
         public void CopyFrom(byte[] data)
         {
             CopyFrom(data, 0, data.Length);