add ResumableTokenStream. fixes #48

benjamin-hodgson · Jul 24, 2020 · 851dc1c · 851dc1c
1 parent 4141044
commit 851dc1c
Show file tree

Hide file tree

Showing 9 changed files with 120 additions and 19 deletions.
diff --git a/Pidgin.Tests/ResumableTokenStreamTests.cs b/Pidgin.Tests/ResumableTokenStreamTests.cs
@@ -0,0 +1,37 @@
+using System;
+using System.IO;
+using Pidgin.TokenStreams;
+using Xunit;
+
+namespace Pidgin.Tests
+{
+    public class ReusableTokenStreamTests
+    {
+        [Fact]
+        public void TestResume()
+        {
+            var input = "aaabb";
+            var stream = new ResumableTokenStream<char>(new ReaderTokenStream(new StringReader(input)));
+
+            // consume two 'a's, reject the third one
+            var chunk = new char[3].AsSpan();
+            stream.Read(chunk);
+            stream.OnParserEnd(chunk.Slice(2));
+
+            stream.Read(chunk);
+            Assert.Equal("abb", chunk.ToString());
+        }
+
+        [Fact]
+        public void TestReturnMultipleChunks()
+        {
+            var stream = new ResumableTokenStream<char>(new ReaderTokenStream(new StringReader("")));
+            stream.OnParserEnd("aa");
+            stream.OnParserEnd("bb");
+
+            var chunk = new char[4].AsSpan();
+            stream.Read(chunk);
+            Assert.Equal("bbaa", chunk.ToString());
+        }
+    }
+}
diff --git a/Pidgin/ParseState.cs b/Pidgin/ParseState.cs
@@ -201,7 +201,7 @@ private void Buffer(int readAhead)
                 _bufferStartLocation += keepFrom;
                 _currentIndex = keepSeenLength;
                 _bufferedCount = keepLength;
-                _bufferedCount += _stream!.ReadInto(_buffer.AsSpan().Slice(_bufferedCount, amountToRead));
+                _bufferedCount += _stream!.Read(_buffer.AsSpan().Slice(_bufferedCount, amountToRead));
             }
         }
 
@@ -248,10 +248,10 @@ public void Dispose()
         {
             if (_buffer != null)
             {
+                _stream!.OnParserEnd(_buffer.AsSpan().Slice(_currentIndex, _bufferedCount - _currentIndex));
                 _arrayPool!.Return(_buffer, _needsClear);
                 _buffer = null;
             }
-            _stream?.Dispose();
             _bookmarks.Dispose();
         }
     }

diff --git a/Pidgin/TokenStreams/EnumeratorTokenStream.cs b/Pidgin/TokenStreams/EnumeratorTokenStream.cs
@@ -14,7 +14,7 @@ public EnumeratorTokenStream(IEnumerator<TToken> input)
         }
 
 
-        public int ReadInto(Span<TToken> buffer)
+        public int Read(Span<TToken> buffer)
         {
             for (var i = 0; i < buffer.Length; i++)
             {
@@ -27,7 +27,5 @@ public int ReadInto(Span<TToken> buffer)
             }
             return buffer.Length;
         }
-
-        public void Dispose() { }
     }
 }
diff --git a/Pidgin/TokenStreams/ITokenStream.cs b/Pidgin/TokenStreams/ITokenStream.cs
@@ -2,9 +2,10 @@
 
 namespace Pidgin.TokenStreams
 {
-    internal interface ITokenStream<TToken> : IDisposable
+    internal interface ITokenStream<TToken>
     {
         int ChunkSizeHint => 1024;
-        int ReadInto(Span<TToken> buffer);
+        int Read(Span<TToken> buffer);
+        void OnParserEnd(ReadOnlySpan<TToken> unconsumed) { }
     }
 }
diff --git a/Pidgin/TokenStreams/ListTokenStream.cs b/Pidgin/TokenStreams/ListTokenStream.cs
@@ -16,7 +16,7 @@ public ListTokenStream(IList<TToken> value)
             _input = value;
         }
 
-        public int ReadInto(Span<TToken> buffer)
+        public int Read(Span<TToken> buffer)
         {
             var actualLength = Math.Min(_input.Count - _index, buffer.Length);
             for (var i = 0; i < actualLength; i++)
@@ -26,7 +26,5 @@ public int ReadInto(Span<TToken> buffer)
             }
             return actualLength;
         }
-
-        public void Dispose() { }
     }
 }
diff --git a/Pidgin/TokenStreams/ReadOnlyListTokenStream.cs b/Pidgin/TokenStreams/ReadOnlyListTokenStream.cs
@@ -16,7 +16,7 @@ public ReadOnlyListTokenStream(IReadOnlyList<TToken> value)
             _input = value;
         }
 
-        public int ReadInto(Span<TToken> buffer)
+        public int Read(Span<TToken> buffer)
         {
             var actualLength = Math.Min(_input.Count - _index, buffer.Length);
             for (var i = 0; i < actualLength; i++)
@@ -26,7 +26,5 @@ public int ReadInto(Span<TToken> buffer)
             }
             return actualLength;
         }
-
-        public void Dispose() { }
     }
 }
diff --git a/Pidgin/TokenStreams/ReaderTokenStream.cs b/Pidgin/TokenStreams/ReaderTokenStream.cs
@@ -14,8 +14,6 @@ public ReaderTokenStream(TextReader input)
             _input = input;
         }
 
-        public int ReadInto(Span<char> buffer) => _input.Read(buffer);
-
-        public void Dispose() { }
+        public int Read(Span<char> buffer) => _input.Read(buffer);
     }
 }
diff --git a/Pidgin/TokenStreams/ResumableTokenStream.cs b/Pidgin/TokenStreams/ResumableTokenStream.cs
@@ -0,0 +1,73 @@
+using System;
+using System.Buffers;
+using System.Runtime.CompilerServices;
+
+namespace Pidgin.TokenStreams
+{
+    internal class ResumableTokenStream<TToken> : ITokenStream<TToken>, IDisposable
+    {
+        private static readonly bool _needsClear = RuntimeHelpers.IsReferenceOrContainsReferences<TToken>();
+        private readonly ArrayPool<TToken> _pool;
+        private readonly ITokenStream<TToken> _next;
+        private TToken[]? _buffer = null;
+        private int _bufferStart = 0;  // amount of empty space at left-hand end of _buffer, aka index of first value
+
+        public ResumableTokenStream(ITokenStream<TToken> next, ArrayPool<TToken>? pool = null)
+        {
+            if (next == null)
+            {
+                throw new ArgumentNullException(nameof(next));
+            }
+            _next = next;
+            _pool = pool ?? ArrayPool<TToken>.Shared;
+        }
+
+        public int Read(Span<TToken> buffer)
+        {
+            var bufferedCount = 0;
+            if (_buffer != null && _bufferStart < _buffer.Length)
+            {
+                bufferedCount = Math.Min(_buffer.Length - _bufferStart, buffer.Length);
+                _buffer.AsSpan().Slice(_bufferStart, bufferedCount).CopyTo(buffer);
+                _bufferStart += bufferedCount;
+            }
+            return bufferedCount + _next.Read(buffer.Slice(bufferedCount));
+        }
+
+        public void OnParserEnd(ReadOnlySpan<TToken> unconsumed)
+        {
+            if (unconsumed.Length == 0)
+            {
+                return;
+            }
+            if (_buffer == null)
+            {
+                _buffer = _pool.Rent(unconsumed.Length);
+                _bufferStart = _buffer.Length;
+            }
+            if (_bufferStart < unconsumed.Length)
+            {
+                var bufferedCount = _buffer.Length - _bufferStart;
+                var newBuffer = _pool.Rent(bufferedCount + unconsumed.Length);
+                var newBufferStart = newBuffer.Length - bufferedCount;
+
+                Array.Copy(_buffer, _bufferStart, newBuffer, newBufferStart, bufferedCount);
+
+                _pool.Return(_buffer, _needsClear);
+                _buffer = newBuffer;
+                _bufferStart = newBufferStart;
+            }
+            _bufferStart -= unconsumed.Length;
+            unconsumed.CopyTo(_buffer.AsSpan().Slice(_bufferStart));
+        }
+
+        public void Dispose()
+        {
+            if (_buffer != null)
+            {
+                _pool.Return(_buffer, _needsClear);
+                _bufferStart = 0;
+            }
+        }
+    }
+}
diff --git a/Pidgin/TokenStreams/StreamTokenStream.cs b/Pidgin/TokenStreams/StreamTokenStream.cs
@@ -14,8 +14,6 @@ public StreamTokenStream(Stream input)
             _input = input;
         }
 
-        public int ReadInto(Span<byte> buffer) => _input.Read(buffer);
-
-        public void Dispose() { }
+        public int Read(Span<byte> buffer) => _input.Read(buffer);
     }
 }