From 8933c9c69780a3b236d48d883dcc6e463e3a0d5a Mon Sep 17 00:00:00 2001
From: Krzysztof Cwalina <kcwalina@microsoft.com>
Date: Tue, 5 Jun 2018 16:09:21 -0700
Subject: [PATCH] Added All Span APIs to Stream Polyfill (#2344)

* Bug fix in BufferWriter

* Added Stream Async Methods
---
 .../System.Memory.Polyfill.csproj             |  1 +
 src/System.Memory.Polyfill/System/Stream.cs   | 75 +++++++++++++++++++
 tests/System.Memory.Polyfill.Tests/Int32.cs   |  2 +-
 .../StreamTests.cs                            | 51 +++++++++++++
 .../System.Memory.Polyfill.Tests.csproj       |  5 +-
 5 files changed, 129 insertions(+), 5 deletions(-)

diff --git a/src/System.Memory.Polyfill/System.Memory.Polyfill.csproj b/src/System.Memory.Polyfill/System.Memory.Polyfill.csproj
index a33d3f88699..c0dd87be855 100644
--- a/src/System.Memory.Polyfill/System.Memory.Polyfill.csproj
+++ b/src/System.Memory.Polyfill/System.Memory.Polyfill.csproj
@@ -9,5 +9,6 @@
     <PackageReference Include="System.Memory" Version="$(SystemMemoryVersion)" />
     <PackageReference Include="System.Buffers" Version="$(SystemMemoryVersion)" />
     <PackageReference Include="System.Runtime.CompilerServices.Unsafe" Version="$(SystemCompilerServicesUnsafeVersion)" />
+    <PackageReference Include="System.Threading.Tasks.Extensions" Version="$(SystemMemoryVersion)" />
   </ItemGroup>
 </Project>
diff --git a/src/System.Memory.Polyfill/System/Stream.cs b/src/System.Memory.Polyfill/System/Stream.cs
index e1d046326e6..9fce3911d60 100644
--- a/src/System.Memory.Polyfill/System/Stream.cs
+++ b/src/System.Memory.Polyfill/System/Stream.cs
@@ -4,6 +4,9 @@
 
 using System.Buffers;
 using System.IO;
+using System.Runtime.InteropServices;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace System
 {
@@ -26,6 +29,78 @@ public static int Read(this Stream stream, Span<byte> buffer)
             {
                 if(pooled != null) ArrayPool<byte>.Shared.Return(pooled);
             }
+#endif
+        }
+
+        public static async ValueTask<int> ReadAsync(this Stream stream, Memory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+#if NETCOREAPP2_1
+            return await stream.ReadAsync(buffer, cancellationToken);
+#else
+            if (MemoryMarshal.TryGetArray(buffer, out ArraySegment<byte> array))
+            {
+                return await stream.ReadAsync(array.Array, array.Offset, array.Count, cancellationToken);
+            }
+            else
+            {
+                byte[] pooled = null;
+                try
+                {
+                    pooled = ArrayPool<byte>.Shared.Rent(buffer.Length);
+                    int read = await stream.ReadAsync(pooled, 0, pooled.Length, cancellationToken);
+                    pooled.AsSpan(0, read).CopyTo(buffer.Span);
+                    return read;
+                }
+                finally
+                {
+                    if (pooled != null) ArrayPool<byte>.Shared.Return(pooled);
+                }
+            }
+#endif
+        }
+
+        public static void Write(this Stream stream, ReadOnlySpan<byte> buffer)
+        {
+#if NETCOREAPP2_1
+            stream.Write(buffer);
+#else
+            byte[] pooled = null;
+            try
+            {
+                pooled = ArrayPool<byte>.Shared.Rent(buffer.Length);
+                buffer.CopyTo(pooled);
+                stream.Write(pooled, 0, pooled.Length);
+            }
+            finally
+            {
+                if(pooled != null) ArrayPool<byte>.Shared.Return(pooled);
+            }
+#endif
+        }
+
+        public static Task WriteAsync(this Stream stream, ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+#if NETCOREAPP2_1
+            return stream.WriteAsync(buffer, cancellationToken).AsTask();
+#else
+            if (MemoryMarshal.TryGetArray(buffer, out ArraySegment<byte> array))
+            {
+                return stream.WriteAsync(array.Array, array.Offset, array.Count, cancellationToken);
+            }
+            else
+            {
+                byte[] pooled = null;
+                try
+                {
+                    pooled = ArrayPool<byte>.Shared.Rent(buffer.Length);
+                    buffer.CopyTo(pooled);
+                    return stream.WriteAsync(pooled, 0, pooled.Length, cancellationToken);
+                }
+                finally
+                {
+                    if (pooled != null) ArrayPool<byte>.Shared.Return(pooled);
+                }
+            }
 #endif
         }
     }
diff --git a/tests/System.Memory.Polyfill.Tests/Int32.cs b/tests/System.Memory.Polyfill.Tests/Int32.cs
index 59275764edc..81336866031 100644
--- a/tests/System.Memory.Polyfill.Tests/Int32.cs
+++ b/tests/System.Memory.Polyfill.Tests/Int32.cs
@@ -11,7 +11,7 @@ public class Int32Tests
         [Fact]
         public void Int32TryParse()
         {
-            ReadOnlySpan<char> span = int.MaxValue.ToString().ToCharArray().AsSpan();
+            ReadOnlySpan<char> span = int.MaxValue.ToString().AsSpan();
             Assert.True(Int32Polyfill.TryParse(span, out int value));
             Assert.Equal(int.MaxValue, value);
         }
diff --git a/tests/System.Memory.Polyfill.Tests/StreamTests.cs b/tests/System.Memory.Polyfill.Tests/StreamTests.cs
index 67a4f156890..b06faa0f4b9 100644
--- a/tests/System.Memory.Polyfill.Tests/StreamTests.cs
+++ b/tests/System.Memory.Polyfill.Tests/StreamTests.cs
@@ -24,5 +24,56 @@ public void StreamRead()
                 Assert.Equal(i, span[i]);
             }
         }
+
+        [Fact]
+        public async void StreamReadAsync()
+        {
+            var buffer = new byte[100];
+            for (int i = 0; i < buffer.Length; i++) buffer[i] = (byte)i;
+            var stream = new MemoryStream(buffer);
+
+            buffer = new byte[100];
+            var memory = new Memory<byte>(buffer);
+            int read = await stream.ReadAsync(memory);
+
+            Assert.Equal(buffer.Length, read);
+            for (int i = 0; i < buffer.Length; i++)
+            {
+                Assert.Equal(i, buffer[i]);
+            }
+        }
+
+        [Fact]
+        public void StreamWrite()
+        {
+            var span = new Span<byte>(new byte[100]);
+            for (int i = 0; i < span.Length; i++) span[i] = (byte)i;
+
+            var buffer = new byte[1000];
+            var stream = new MemoryStream(buffer);
+
+            stream.Write(span);
+            for (int i = 0; i < span.Length; i++)
+            {
+                Assert.Equal(i, buffer[i]);
+            }
+        }
+
+        [Fact]
+        public async void StreamWriteAsync()
+        {
+            var array = new byte[100];
+            var memory = new Memory<byte>(array);
+            for (int i = 0; i < array.Length; i++) array[i] = (byte)i;
+
+            var buffer = new byte[100];
+            var stream = new MemoryStream(buffer);
+            await stream.WriteAsync(memory);
+
+            for (int i = 0; i < buffer.Length; i++)
+            {
+                Assert.Equal(i, buffer[i]);
+            }
+        }
     }
 }
diff --git a/tests/System.Memory.Polyfill.Tests/System.Memory.Polyfill.Tests.csproj b/tests/System.Memory.Polyfill.Tests/System.Memory.Polyfill.Tests.csproj
index 68de1e22c13..6642f6f0f86 100644
--- a/tests/System.Memory.Polyfill.Tests/System.Memory.Polyfill.Tests.csproj
+++ b/tests/System.Memory.Polyfill.Tests/System.Memory.Polyfill.Tests.csproj
@@ -1,10 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk" ToolsVersion="15.0" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
   <Import Project="..\..\tools\common.props" />
   <PropertyGroup>
-    <TargetFrameworks>netcoreapp2.1</TargetFrameworks>
-  </PropertyGroup>
-  <PropertyGroup Condition="'$(TargetsWindows)' == 'true'">
-    <TargetFrameworks>net46;netcoreapp2.1</TargetFrameworks>
+    <TargetFrameworks>netcoreapp2.0;netcoreapp2.1</TargetFrameworks>
   </PropertyGroup>
   <ItemGroup>
     <PackageReference Include="xunit" Version="2.3.1" />