dennisarfan/StitchATon2
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 2 Packages Wiki Activity Actions

solve edge case

Browse source
This commit is contained in:
dennisarfan 2025-08-01 09:51:39 +07:00
parent 741d34a5e0
commit 0472bfe58e
15 changed files with 685 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

11
Infra/Encoders/Crc32.cs
View file

@ -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;

2
Infra/Encoders/PngPipeEncoder.cs
View file

@ -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
View 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);
}
}
}