acdream/src/AcDream.Content/Pak/PakWriter.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;

namespace AcDream.Content.Pak;

/// <summary>
/// Streams pak blobs to disk: header placeholder -> 64-byte-aligned blobs ->
/// TOC (sorted by key) -> seek back and finalize the header once
/// tocOffset/tocCount are known. This lets the writer stream blobs without
/// knowing the final count up front (plan: "TOC last").
/// </summary>
public sealed class PakWriter : IDisposable {
    private const int Alignment = 64;

    private readonly FileStream _stream;
    private readonly PakHeader _headerTemplate;
    private readonly List<PakTocEntry> _tocEntries = new();
    private readonly HashSet<ulong> _seenKeys = new();
    private readonly Dictionary<ulong, PakTocEntry> _receiptByKey = new();
    private bool _finished;

    public PakWriter(string path, PakHeader headerTemplate) {
        _stream = new FileStream(path, FileMode.Create, FileAccess.ReadWrite, FileShare.None);
        _headerTemplate = headerTemplate;
        // The format version is the WRITER's identity, not caller input —
        // stamp it unconditionally so a caller-populated header template can
        // never produce a pak claiming a version this code doesn't write
        // (the default-0 footgun; PakReader rejects any version but
        // PakFormat.CurrentFormatVersion).
        _headerTemplate.FormatVersion = PakFormat.CurrentFormatVersion;
        _headerTemplate.BakeToolVersion = PakFormat.CurrentBakeToolVersion;

        // Write a header placeholder now; finalized in Finish() once TocOffset/TocCount are known.
        _headerTemplate.WriteTo(_stream);
        PadToAlignment();
    }

    /// <summary>
    /// Adds one blob (an <see cref="ObjectMeshData"/> serialized via
    /// <see cref="ObjectMeshDataSerializer"/>) under <paramref name="key"/>.
    /// Blobs are written immediately at the current (64-byte-aligned) stream
    /// position; the writer pads after each blob so the NEXT blob also
    /// starts aligned.
    /// </summary>
    public void AddBlob(ulong key, ObjectMeshData data) {
        ThrowIfFinished();
        ReserveKey(key);

        using var ms = new MemoryStream();
        ObjectMeshDataSerializer.Write(data, ms);
        var bytes = ms.ToArray();
        AddReservedBlob(key, bytes);
    }

    /// <summary>
    /// Adds an already serialized immutable payload. This is the package seam
    /// used by typed non-render assets such as flat collision records.
    /// </summary>
    public void AddBlob(ulong key, ReadOnlySpan<byte> bytes) {
        ThrowIfFinished();
        ReserveKey(key);
        AddReservedBlob(key, bytes);
    }

    private void AddReservedBlob(ulong key, ReadOnlySpan<byte> bytes) {
        if ((ulong)bytes.Length > uint.MaxValue) {
            throw new ArgumentException(
                "one pak payload cannot exceed UInt32.MaxValue bytes",
                nameof(bytes));
        }

        long offset = _stream.Position;
        _stream.Write(bytes);
        PadToAlignment();

        var receipt = new PakTocEntry {
            Key = key,
            Offset = (ulong)offset,
            Length = checked((uint)bytes.Length),
            Crc32 = Content.Pak.Crc32.Compute(bytes),
        };
        _tocEntries.Add(receipt);
        _receiptByKey.Add(key, receipt);
    }

    private void ReserveKey(ulong key) {
        if (!_seenKeys.Add(key)) {
            throw new ArgumentException($"duplicate pak key 0x{key:X16}", nameof(key));
        }
    }

    /// <summary>
    /// Adds a second independently addressable key for an existing physical
    /// blob. The alias receives a normal TOC row with the source blob's exact
    /// offset, length, and CRC; no payload bytes are written again.
    /// </summary>
    public void AddAlias(ulong aliasKey, ulong existingKey) {
        ThrowIfFinished();
        if (aliasKey == existingKey) {
            throw new ArgumentException("an alias key must differ from its source key", nameof(aliasKey));
        }
        if (!_receiptByKey.TryGetValue(existingKey, out var source)) {
            throw new KeyNotFoundException($"source pak key 0x{existingKey:X16} has not been written");
        }
        if (!_seenKeys.Add(aliasKey)) {
            throw new ArgumentException($"duplicate pak key 0x{aliasKey:X16}", nameof(aliasKey));
        }

        var alias = source;
        alias.Key = aliasKey;
        _tocEntries.Add(alias);
        _receiptByKey.Add(aliasKey, alias);
    }

    /// <summary>Writes the sorted TOC and finalizes the header. Must be called exactly once.</summary>
    public void Finish() {
        ThrowIfFinished();
        _finished = true;

        var sorted = _tocEntries.OrderBy(e => e.Key).ToList();
        ulong tocOffset = (ulong)_stream.Position;
        foreach (var entry in sorted) {
            entry.WriteTo(_stream);
        }

        var finalHeader = _headerTemplate;
        finalHeader.TocOffset = tocOffset;
        finalHeader.TocCount = (uint)sorted.Count;

        _stream.Position = 0;
        finalHeader.WriteTo(_stream);
        _stream.Flush();
    }

    private void PadToAlignment() {
        long pos = _stream.Position;
        long remainder = pos % Alignment;
        if (remainder == 0) return;
        long pad = Alignment - remainder;
        Span<byte> zeros = stackalloc byte[(int)pad];
        zeros.Clear();
        _stream.Write(zeros);
    }

    private void ThrowIfFinished() {
        if (_finished) throw new InvalidOperationException("PakWriter.Finish() has already been called.");
    }

    public void Dispose() {
        // Best-effort finalize so a forgotten Finish() call doesn't leave a
        // half-written file with a placeholder header claiming 0 entries
        // while blobs are on disk. try/finally, NOT a bare sequence: Finish()
        // performs real I/O (TOC write, header seek-back, flush) and can
        // throw for environmental reasons (disk full, I/O error) — the
        // stream must ALWAYS be closed so the file handle is never leaked
        // mid-unwind (review finding 5 on the a5926ebc cleanup, which had
        // removed exactly this guarantee).
        try {
            if (!_finished) {
                Finish();
            }
        }
        finally {
            _stream.Dispose();
        }
    }
}