solve edge case
This commit is contained in:
dennisarfan
parent
741d34a5e0
commit
0472bfe58e
15 changed files with 685 additions and 47 deletions
102
Infra/Buffers/PooledMemoryStream.cs
Normal file
102
Infra/Buffers/PooledMemoryStream.cs
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
using System.Buffers;
|
||||
|
||||
namespace StitchATon2.Infra.Buffers;
|
||||
|
||||
public class PooledMemoryStream : Stream
|
||||
{
|
||||
private byte[] _buffer;
|
||||
private int _length;
|
||||
private int _position;
|
||||
private readonly ArrayPool<byte> _pool;
|
||||
private bool _disposed;
|
||||
|
||||
public PooledMemoryStream(int initialCapacity = 1024, ArrayPool<byte>? pool = null)
|
||||
{
|
||||
_pool = pool ?? ArrayPool<byte>.Shared;
|
||||
_buffer = _pool.Rent(initialCapacity);
|
||||
}
|
||||
|
||||
public override bool CanRead => !_disposed;
|
||||
public override bool CanSeek => !_disposed;
|
||||
public override bool CanWrite => !_disposed;
|
||||
public override long Length => _length;
|
||||
public override long Position
|
||||
{
|
||||
get => _position;
|
||||
set
|
||||
{
|
||||
if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
|
||||
if (value < 0 || value > int.MaxValue) throw new ArgumentOutOfRangeException();
|
||||
_position = (int)value;
|
||||
}
|
||||
}
|
||||
|
||||
public byte[] GetBuffer() => _buffer;
|
||||
public ArraySegment<byte> GetWrittenSegment() => new(_buffer, 0, _length);
|
||||
|
||||
public override void Flush() { /* no-op */ }
|
||||
|
||||
public override int Read(byte[] buffer, int offset, int count)
|
||||
{
|
||||
if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
|
||||
int available = _length - _position;
|
||||
int toRead = Math.Min(count, available);
|
||||
Buffer.BlockCopy(_buffer, _position, buffer, offset, toRead);
|
||||
_position += toRead;
|
||||
return toRead;
|
||||
}
|
||||
|
||||
public override void Write(byte[] buffer, int offset, int count)
|
||||
{
|
||||
if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
|
||||
EnsureCapacity(_position + count);
|
||||
Buffer.BlockCopy(buffer, offset, _buffer, _position, count);
|
||||
_position += count;
|
||||
_length = Math.Max(_length, _position);
|
||||
}
|
||||
|
||||
public override long Seek(long offset, SeekOrigin origin)
|
||||
{
|
||||
if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
|
||||
int newPos = origin switch
|
||||
{
|
||||
SeekOrigin.Begin => (int)offset,
|
||||
SeekOrigin.Current => _position + (int)offset,
|
||||
SeekOrigin.End => _length + (int)offset,
|
||||
_ => throw new ArgumentOutOfRangeException()
|
||||
};
|
||||
if (newPos < 0) throw new IOException("Negative position");
|
||||
_position = newPos;
|
||||
return _position;
|
||||
}
|
||||
|
||||
public override void SetLength(long value)
|
||||
{
|
||||
if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
|
||||
if (value < 0 || value > int.MaxValue) throw new ArgumentOutOfRangeException();
|
||||
EnsureCapacity((int)value);
|
||||
_length = (int)value;
|
||||
if (_position > _length) _position = _length;
|
||||
}
|
||||
|
||||
private void EnsureCapacity(int size)
|
||||
{
|
||||
if (size <= _buffer.Length) return;
|
||||
int newSize = Math.Max(size, _buffer.Length * 2);
|
||||
byte[] newBuffer = _pool.Rent(newSize);
|
||||
Buffer.BlockCopy(_buffer, 0, newBuffer, 0, _length);
|
||||
_pool.Return(_buffer, clearArray: true);
|
||||
_buffer = newBuffer;
|
||||
}
|
||||
|
||||
protected override void Dispose(bool disposing)
|
||||
{
|
||||
if (!_disposed)
|
||||
{
|
||||
_pool.Return(_buffer, clearArray: true);
|
||||
_buffer = Array.Empty<byte>();
|
||||
_disposed = true;
|
||||
}
|
||||
base.Dispose(disposing);
|
||||
}
|
||||
}
|
||||
|
|
@ -15,6 +15,17 @@ public static class Crc32
|
|||
return ~crc;
|
||||
}
|
||||
|
||||
public static uint Compute(Stream stream, int count, uint initial = 0xFFFFFFFF)
|
||||
{
|
||||
uint crc = initial;
|
||||
while (count-- > 0)
|
||||
{
|
||||
crc = Table[(crc ^ stream.ReadByte()) & 0xFF] ^ (crc >> 8);
|
||||
}
|
||||
|
||||
return ~crc;
|
||||
}
|
||||
|
||||
private static uint[] GenerateTable()
|
||||
{
|
||||
const uint poly = 0xEDB88320;
|
||||
|
|
|
|||
|
|
@ -71,7 +71,7 @@ public class PngPipeEncoder : IDisposable
|
|||
_zlibStream.Write(buffer.Span.Slice(offset, FlushThreshold));
|
||||
await _zlibStream.FlushAsync(cancellationToken);
|
||||
offset += FlushThreshold;
|
||||
if(_outputPipe.UnflushedBytes >= PipeChunkThreshold)
|
||||
if(_memoryStream.Length >= BufferSize)
|
||||
await FlushAsync(cancellationToken);
|
||||
}
|
||||
|
||||
|
|
|
|||
166
Infra/Encoders/UnsafePngEncoder.cs
Normal file
166
Infra/Encoders/UnsafePngEncoder.cs
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
using System.Buffers;
|
||||
using System.Buffers.Binary;
|
||||
using System.IO.Compression;
|
||||
using System.IO.Pipelines;
|
||||
using StitchATon2.Infra.Buffers;
|
||||
|
||||
namespace StitchATon2.Infra.Encoders;
|
||||
|
||||
public class UnsafePngEncoder : IDisposable
|
||||
{
|
||||
private const int BufferSize = 8 * 1024;
|
||||
private const int FlushThreshold = 1024;
|
||||
private const int PipeChunkThreshold = 16 * 1024;
|
||||
|
||||
private readonly PipeWriter _outputPipe;
|
||||
private readonly int _width;
|
||||
private readonly int _height;
|
||||
|
||||
private MemoryHandle? _memoryHandle;
|
||||
private readonly RawPointerStream _memoryStream;// = new RawPointerStream();
|
||||
private ZLibStream? _zlibStream;// = new ZLibStream(_memoryStream, CompressionLevel.Optimal, leaveOpen: true);
|
||||
private bool _disposed;
|
||||
private bool _shouldFlush;
|
||||
|
||||
public UnsafePngEncoder(PipeWriter outputPipe, int width, int height)
|
||||
{
|
||||
_outputPipe = outputPipe;
|
||||
_width = width;
|
||||
_height = height;
|
||||
|
||||
_memoryStream = new RawPointerStream();
|
||||
}
|
||||
|
||||
~UnsafePngEncoder() => Dispose();
|
||||
|
||||
public void WriteHeader()
|
||||
{
|
||||
Span<byte> headerBytes = [
|
||||
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG Signature
|
||||
0x00, 0x00, 0x00, 0x0D, // Length
|
||||
|
||||
// IHDR chunk
|
||||
0x49, 0x48, 0x44, 0x52, // IHDR
|
||||
0x00, 0x00, 0x00, 0x00, // Reserve to write Width
|
||||
0x00, 0x00, 0x00, 0x00, // Reserve to write Height
|
||||
0x08, // Bit depth
|
||||
0x02, // Color type
|
||||
0x00, // Compression method
|
||||
0x00, // Filter method
|
||||
0x00, // Interlace method
|
||||
0x00, 0x00, 0x00, 0x00, // Reserve to write CRC-32
|
||||
];
|
||||
|
||||
BinaryPrimitives.WriteInt32BigEndian(headerBytes[16..], _width);
|
||||
BinaryPrimitives.WriteInt32BigEndian(headerBytes[20..], _height);
|
||||
var crc = Crc32.Compute(headerBytes.Slice(12, 17));
|
||||
|
||||
BinaryPrimitives.WriteUInt32BigEndian(headerBytes[29..], crc);
|
||||
|
||||
_outputPipe.Write(headerBytes);
|
||||
}
|
||||
|
||||
private unsafe void Initialize()
|
||||
{
|
||||
if (_memoryHandle == null)
|
||||
{
|
||||
var memory = _outputPipe.GetMemory(PipeChunkThreshold);
|
||||
var handle = memory.Pin();
|
||||
_memoryStream.Initialize((byte*)handle.Pointer, 0, memory.Length);
|
||||
_memoryHandle = handle;
|
||||
|
||||
_memoryStream.SetLength(8);
|
||||
_memoryStream.Position = 8;
|
||||
_zlibStream = new ZLibStream(_memoryStream, CompressionLevel.Optimal, true);
|
||||
}
|
||||
}
|
||||
|
||||
public async Task WriteDataAsync(IBuffer<byte> buffer, bool disposeBuffer = true, CancellationToken cancellationToken = default)
|
||||
{
|
||||
Initialize();
|
||||
_zlibStream!.Write([0]);
|
||||
|
||||
var offset = 0;
|
||||
while (buffer.Length - offset > FlushThreshold)
|
||||
{
|
||||
_zlibStream.Write(buffer.Span.Slice(offset, FlushThreshold));
|
||||
await _zlibStream.FlushAsync(cancellationToken);
|
||||
offset += FlushThreshold;
|
||||
if(_memoryStream.Length >= BufferSize)
|
||||
await FlushAsync(cancellationToken);
|
||||
}
|
||||
|
||||
if (buffer.Length > offset)
|
||||
{
|
||||
_zlibStream.Write(buffer.Span[offset..]);
|
||||
await _zlibStream.FlushAsync(cancellationToken);
|
||||
_shouldFlush = true;
|
||||
}
|
||||
|
||||
if(disposeBuffer) buffer.Dispose();
|
||||
}
|
||||
|
||||
private async Task FlushAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
await _zlibStream!.FlushAsync(cancellationToken);
|
||||
var dataSize = (int)(_memoryStream.Length - 8);
|
||||
|
||||
_memoryStream.Position = 0;
|
||||
Span<byte> buffer = stackalloc byte[4];
|
||||
BinaryPrimitives.WriteInt32BigEndian(buffer, dataSize);
|
||||
_memoryStream.Write(buffer);
|
||||
_memoryStream.Write("IDAT"u8);
|
||||
|
||||
_memoryStream.Position = 4;
|
||||
|
||||
// write Crc
|
||||
var crc = Crc32.Compute(_memoryStream, dataSize + 4);
|
||||
BinaryPrimitives.WriteUInt32BigEndian(buffer, crc);
|
||||
_memoryStream.Write(buffer);
|
||||
|
||||
_outputPipe.Advance((int)_memoryStream.Length);
|
||||
|
||||
await _memoryStream.DisposeAsync();
|
||||
_memoryHandle!.Value.Dispose();
|
||||
_memoryHandle = null;
|
||||
|
||||
_shouldFlush = false;
|
||||
}
|
||||
|
||||
public async Task WriteEndOfFileAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
if(_shouldFlush)
|
||||
await FlushAsync(cancellationToken);
|
||||
|
||||
Span<byte> endChunk = [
|
||||
0x00, 0x00, 0x00, 0x00, // Length
|
||||
0x49, 0x45, 0x4E, 0x44, // IEND
|
||||
0xAE, 0x42, 0x60, 0x82, // Crc
|
||||
];
|
||||
|
||||
_outputPipe.Write(endChunk);
|
||||
Dispose();
|
||||
}
|
||||
|
||||
public void Dispose()
|
||||
{
|
||||
if (!_disposed)
|
||||
{
|
||||
if (_memoryHandle != null)
|
||||
{
|
||||
_zlibStream!.Dispose();
|
||||
_memoryStream.Dispose();
|
||||
}
|
||||
_disposed = true;
|
||||
GC.SuppressFinalize(this);
|
||||
}
|
||||
}
|
||||
|
||||
private unsafe class RawPointerStream : UnmanagedMemoryStream
|
||||
{
|
||||
public void Initialize(byte* pointer, int length, int capacity)
|
||||
{
|
||||
Initialize(pointer, length, capacity, FileAccess.ReadWrite);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
using System.Buffers;
|
||||
using System.IO.MemoryMappedFiles;
|
||||
using System.Runtime.CompilerServices;
|
||||
using System.Runtime.InteropServices;
|
||||
|
|
@ -19,7 +18,7 @@ public class ImageIntegral : IDisposable
|
|||
private readonly int _width;
|
||||
private readonly int _height;
|
||||
|
||||
private IMemoryOwner<ManualResetEventSlim>? _rowLocks;
|
||||
private ManualResetEventSlim[]? _rowLocks;
|
||||
private MemoryMappedFile? _memoryMappedFile;
|
||||
private readonly Lock _lock = new();
|
||||
|
||||
|
|
@ -66,7 +65,7 @@ public class ImageIntegral : IDisposable
|
|||
}
|
||||
}
|
||||
|
||||
_rowLocks?.Memory.Span[row].Wait(cancellationToken);
|
||||
_rowLocks?[row].Wait(cancellationToken);
|
||||
}
|
||||
|
||||
private void Initialize(CancellationToken cancellationToken)
|
||||
|
|
@ -101,12 +100,11 @@ public class ImageIntegral : IDisposable
|
|||
// initialize resource gating, all rows is expected to be locked
|
||||
// if the backed file require to allocate, it should be safe to do this
|
||||
// asynchronously
|
||||
var rowLocks = MemoryAllocator.AllocateManaged<ManualResetEventSlim>(_height);
|
||||
var rowLocksSpan = rowLocks.Memory.Span;
|
||||
var rowLocks = new ManualResetEventSlim[_height];
|
||||
for (int i = 0; i < _height; i++)
|
||||
{
|
||||
var isOpen = i < header.ProcessedRows;
|
||||
rowLocksSpan[i] = new ManualResetEventSlim(isOpen);
|
||||
rowLocks[i] = new ManualResetEventSlim(isOpen);
|
||||
}
|
||||
|
||||
_rowLocks = rowLocks;
|
||||
|
|
@ -222,7 +220,7 @@ public class ImageIntegral : IDisposable
|
|||
view.DangerousWriteSpan(0, writeBuffer.Span, 0, _width);
|
||||
|
||||
writeBuffer.Dispose();
|
||||
_rowLocks!.Memory.Span[row].Set();
|
||||
_rowLocks![row].Set();
|
||||
Interlocked.Increment(ref _processedRows);
|
||||
|
||||
using (var view = AcquireHeaderView(MemoryMappedFileAccess.Write))
|
||||
|
|
@ -370,11 +368,8 @@ public class ImageIntegral : IDisposable
|
|||
if (_rowLocks is { } locks)
|
||||
{
|
||||
_rowLocks = null;
|
||||
var lockSpan = locks.Memory.Span;
|
||||
for(int i = 0; i < _height; i++)
|
||||
lockSpan[i].Dispose();
|
||||
|
||||
locks.Dispose();
|
||||
locks[i].Dispose();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -68,6 +68,12 @@ public record struct Int32Pixel
|
|||
return new Int32Pixel(a.R / b, a.G / b, a.B / b);
|
||||
}
|
||||
|
||||
[MethodImpl(MethodImplOptions.AggressiveInlining)]
|
||||
public static Int32Pixel operator /(Int32Pixel a, float b)
|
||||
{
|
||||
return new Int32Pixel((byte)(a.R / b), (byte)(a.G / b), (byte)(a.B / b));
|
||||
}
|
||||
|
||||
[MethodImpl(MethodImplOptions.AggressiveInlining)]
|
||||
public static explicit operator Int32Pixel(Rgb24 pixel)
|
||||
{
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue