Skip to content

Commit

Permalink
Feedback updates and massively expand write tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@JimBobSquarePants
JimBobSquarePants committed Nov 4, 2024
1 parent b74d2e4 commit 6301662
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 117 deletions.
148 changes: 49 additions & 99 deletions src/ImageSharp/IO/ChunkedMemoryStream.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
// Licensed under the Six Labors Split License.

using System.Buffers;
using System.Collections;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using SixLabors.ImageSharp.Memory;
Expand All @@ -14,15 +13,13 @@ namespace SixLabors.ImageSharp.IO;
/// Chunks are allocated by the <see cref="MemoryAllocator"/> assigned via the constructor
/// and is designed to take advantage of buffer pooling when available.
/// </summary>
public class ChunkedMemoryStream : Stream
internal sealed class ChunkedMemoryStream : Stream
{
private readonly MemoryChunkBuffer memoryChunkBuffer;
private readonly byte[] singleByteBuffer = new byte[1];

private long length;
private long position;
private int currentChunk;
private int currentChunkIndex;
private int bufferIndex;
private int chunkIndex;
private bool isDisposed;

/// <summary>
Expand Down Expand Up @@ -95,21 +92,13 @@ public override void SetLength(long value)
/// <inheritdoc/>
public override int ReadByte()
{
this.EnsureNotDisposed();
if (this.position >= this.length)
{
return -1;
}

_ = this.Read(this.singleByteBuffer, 0, 1);
return MemoryMarshal.GetReference<byte>(this.singleByteBuffer);
Unsafe.SkipInit(out byte b);
return this.Read(MemoryMarshal.CreateSpan(ref b, 1)) == 1 ? b : -1;
}

/// <inheritdoc/>
public override int Read(byte[] buffer, int offset, int count)
{
this.EnsureNotDisposed();

Guard.NotNull(buffer, nameof(buffer));
Guard.MustBeGreaterThanOrEqualTo(offset, 0, nameof(offset));
Guard.MustBeGreaterThanOrEqualTo(count, 0, nameof(count));
Expand All @@ -135,39 +124,34 @@ public override int Read(Span<byte> buffer)
return 0;
}

if (remaining > count)
{
remaining = count;
}

int bytesToRead = (int)remaining;
int bytesToRead = count;
int bytesRead = 0;
while (bytesToRead != 0 && this.currentChunk != this.memoryChunkBuffer.Length)
while (bytesToRead > 0 && this.bufferIndex != this.memoryChunkBuffer.Length)
{
bool moveToNextChunk = false;
MemoryChunk chunk = this.memoryChunkBuffer[this.currentChunk];
MemoryChunk chunk = this.memoryChunkBuffer[this.bufferIndex];
int n = bytesToRead;
int remainingBytesInCurrentChunk = chunk.Length - this.currentChunkIndex;
int remainingBytesInCurrentChunk = chunk.Length - this.chunkIndex;
if (n >= remainingBytesInCurrentChunk)
{
n = remainingBytesInCurrentChunk;
moveToNextChunk = true;
}

// Read n bytes from the current chunk
chunk.Buffer.Memory.Span.Slice(this.currentChunkIndex, n).CopyTo(buffer.Slice(offset, n));
chunk.Buffer.Memory.Span.Slice(this.chunkIndex, n).CopyTo(buffer.Slice(offset, n));
bytesToRead -= n;
offset += n;
bytesRead += n;

if (moveToNextChunk)
{
this.currentChunkIndex = 0;
this.currentChunk++;
this.chunkIndex = 0;
this.bufferIndex++;
}
else
{
this.currentChunkIndex += n;
this.chunkIndex += n;
}
}

Expand All @@ -177,11 +161,7 @@ public override int Read(Span<byte> buffer)

/// <inheritdoc/>
public override void WriteByte(byte value)
{
this.EnsureNotDisposed();
MemoryMarshal.Write(this.singleByteBuffer, ref value);
this.Write(this.singleByteBuffer, 0, 1);
}
=> this.Write(MemoryMarshal.CreateSpan(ref value, 1));

/// <inheritdoc/>
public override void Write(byte[] buffer, int offset, int count)
Expand Down Expand Up @@ -213,39 +193,34 @@ public override void Write(ReadOnlySpan<byte> buffer)
remaining = this.memoryChunkBuffer.Length - this.position;
}

if (remaining > count)
{
remaining = count;
}

int bytesToWrite = (int)remaining;
int bytesToWrite = count;
int bytesWritten = 0;
while (bytesToWrite != 0 && this.currentChunk != this.memoryChunkBuffer.Length)
while (bytesToWrite > 0 && this.bufferIndex != this.memoryChunkBuffer.Length)
{
bool moveToNextChunk = false;
MemoryChunk chunk = this.memoryChunkBuffer[this.currentChunk];
MemoryChunk chunk = this.memoryChunkBuffer[this.bufferIndex];
int n = bytesToWrite;
int remainingBytesInCurrentChunk = chunk.Length - this.currentChunkIndex;
int remainingBytesInCurrentChunk = chunk.Length - this.chunkIndex;
if (n >= remainingBytesInCurrentChunk)
{
n = remainingBytesInCurrentChunk;
moveToNextChunk = true;
}

// Write n bytes to the current chunk
buffer.Slice(offset, n).CopyTo(chunk.Buffer.Slice(this.currentChunkIndex, n));
buffer.Slice(offset, n).CopyTo(chunk.Buffer.Slice(this.chunkIndex, n));
bytesToWrite -= n;
offset += n;
bytesWritten += n;

if (moveToNextChunk)
{
this.currentChunkIndex = 0;
this.currentChunk++;
this.chunkIndex = 0;
this.bufferIndex++;
}
else
{
this.currentChunkIndex += n;
this.chunkIndex += n;
}
}

Expand Down Expand Up @@ -275,31 +250,31 @@ public void WriteTo(Stream stream)

int bytesToRead = (int)remaining;
int bytesRead = 0;
while (bytesToRead != 0 && this.currentChunk != this.memoryChunkBuffer.Length)
while (bytesToRead > 0 && this.bufferIndex != this.memoryChunkBuffer.Length)
{
bool moveToNextChunk = false;
MemoryChunk chunk = this.memoryChunkBuffer[this.currentChunk];
MemoryChunk chunk = this.memoryChunkBuffer[this.bufferIndex];
int n = bytesToRead;
int remainingBytesInCurrentChunk = chunk.Length - this.currentChunkIndex;
int remainingBytesInCurrentChunk = chunk.Length - this.chunkIndex;
if (n >= remainingBytesInCurrentChunk)
{
n = remainingBytesInCurrentChunk;
moveToNextChunk = true;
}

// Read n bytes from the current chunk
stream.Write(chunk.Buffer.Memory.Span.Slice(this.currentChunkIndex, n));
stream.Write(chunk.Buffer.Memory.Span.Slice(this.chunkIndex, n));
bytesToRead -= n;
bytesRead += n;

if (moveToNextChunk)
{
this.currentChunkIndex = 0;
this.currentChunk++;
this.chunkIndex = 0;
this.bufferIndex++;
}
else
{
this.currentChunkIndex += n;
this.chunkIndex += n;
}
}

Expand Down Expand Up @@ -338,8 +313,8 @@ protected override void Dispose(bool disposing)
this.memoryChunkBuffer.Dispose();
}

this.currentChunk = 0;
this.currentChunkIndex = 0;
this.bufferIndex = 0;
this.chunkIndex = 0;
this.position = 0;
this.length = 0;
}
Expand All @@ -366,8 +341,8 @@ private void SetPosition(long value)
// If the new position is greater than the length of the stream, set the position to the end of the stream
if (offset > 0 && offset >= this.memoryChunkBuffer.Length)
{
this.currentChunk = this.memoryChunkBuffer.ChunkCount - 1;
this.currentChunkIndex = this.memoryChunkBuffer[this.currentChunk].Length - 1;
this.bufferIndex = this.memoryChunkBuffer.ChunkCount - 1;
this.chunkIndex = this.memoryChunkBuffer[this.bufferIndex].Length - 1;
return;
}

Expand All @@ -386,10 +361,10 @@ private void SetPosition(long value)
currentChunkIndex++;
}

this.currentChunk = currentChunkIndex;
this.bufferIndex = currentChunkIndex;

// Safe to cast here as we know the offset is less than the chunk length.
this.currentChunkIndex = (int)offset;
this.chunkIndex = (int)offset;
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
Expand All @@ -404,7 +379,7 @@ private void EnsureNotDisposed()
[MethodImpl(MethodImplOptions.NoInlining)]
private static void ThrowDisposed() => throw new ObjectDisposedException(nameof(ChunkedMemoryStream), "The stream is closed.");

private sealed class MemoryChunkBuffer : IEnumerable<MemoryChunk>, IDisposable
private sealed class MemoryChunkBuffer : IDisposable
{
private readonly List<MemoryChunk> memoryChunks = new();
private readonly MemoryAllocator allocator;
Expand Down Expand Up @@ -439,15 +414,19 @@ public void Expand()

public void Dispose()
{
this.Dispose(true);
GC.SuppressFinalize(this);
}
if (!this.isDisposed)
{
foreach (MemoryChunk chunk in this.memoryChunks)
{
chunk.Dispose();
}

public IEnumerator<MemoryChunk> GetEnumerator()
=> ((IEnumerable<MemoryChunk>)this.memoryChunks).GetEnumerator();
this.memoryChunks.Clear();

IEnumerator IEnumerable.GetEnumerator()
=> ((IEnumerable)this.memoryChunks).GetEnumerator();
this.Length = 0;
this.isDisposed = true;
}
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static int GetChunkSize(int i)
Expand All @@ -459,25 +438,6 @@ private static int GetChunkSize(int i)
const int b4M = 1 << 22;
return i < 16 ? b128K * (1 << (int)((uint)i / 4)) : b4M;
}

private void Dispose(bool disposing)
{
if (!this.isDisposed)
{
if (disposing)
{
foreach (MemoryChunk chunk in this.memoryChunks)
{
chunk.Dispose();
}

this.memoryChunks.Clear();
}

this.Length = 0;
this.isDisposed = true;
}
}
}

private sealed class MemoryChunk : IDisposable
Expand All @@ -490,23 +450,13 @@ private sealed class MemoryChunk : IDisposable

public int Length { get; init; }

private void Dispose(bool disposing)
public void Dispose()
{
if (!this.isDisposed)
{
if (disposing)
{
this.Buffer.Dispose();
}

this.Buffer.Dispose();
this.isDisposed = true;
}
}

public void Dispose()
{
this.Dispose(disposing: true);
GC.SuppressFinalize(this);
}
}
}
Loading

0 comments on commit 6301662

Please sign in to comment.