package pdf

import "core:bytes"
import "core:compress/zlib"
import "core:mem"
import "core:slice"
import "core:testing"

// Test vectors are produced by the reference implementation (zlib 1.2.x via
// CPython) so the expected output is never this decoder's own opinion, and the
// well-formed ones are additionally checked against core:compress/zlib at run
// time. The malformed ones are hand-assembled bit by bit.

@(private="file")
V_STORED_PLAIN :: "Hello, PDF world!"
@(private="file")
V_FIXED_PLAIN :: "abcabcabcabcabcabcabcabc"

// zlib.compress("") -- a final fixed block whose only symbol is end-of-block.
@(private="file")
v_empty := [?]byte{
    0x78, 0xda, 0x03, 0x00, 0x00, 0x00, 0x00, 0x01,
}

// level 0: a single stored block carrying V_STORED_PLAIN verbatim.
@(private="file")
v_stored := [?]byte{
    0x78, 0x01, 0x01, 0x11, 0x00, 0xee, 0xff, 0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x2c, 0x20, 0x50, 0x44,
    0x46, 0x20, 0x77, 0x6f, 0x72, 0x6c, 0x64, 0x21, 0x31, 0xf6, 0x05, 0x84,
}

// fixed-Huffman block with a length/distance pair ("abc" repeated 8 times).
@(private="file")
v_fixed := [?]byte{
    0x78, 0xda, 0x4b, 0x4c, 0x4a, 0x4e, 0xc4, 0x86, 0x00, 0x72, 0xe0, 0x09, 0x31,
}

// dynamic-Huffman block: 201 bytes of content-stream-like text.
@(private="file")
v_dynamic := [?]byte{
    0x78, 0xda, 0x2d, 0x8e, 0xbb, 0x15, 0xc4, 0x30, 0x08, 0x04, 0xf3, 0xab, 0x62, 0x4b, 0x10, 0x7f,
    0x54, 0x87, 0xdb, 0x38, 0xf7, 0x9f, 0x7a, 0x85, 0x1d, 0x10, 0xcc, 0x3c, 0x09, 0x66, 0x61, 0xe1,
    0xe6, 0x2c, 0xfc, 0x71, 0xfd, 0x04, 0x46, 0x2a, 0x88, 0x0c, 0x2a, 0x92, 0x28, 0x0e, 0xd5, 0x61,
    0xc3, 0x26, 0x2b, 0x5f, 0xd9, 0xb0, 0x43, 0xf4, 0x88, 0x86, 0xfb, 0x88, 0x80, 0x04, 0x85, 0x05,
    0x22, 0x46, 0x24, 0xa4, 0x29, 0x9c, 0x9b, 0x72, 0x44, 0x9d, 0xef, 0x14, 0x1b, 0x55, 0x23, 0x1a,
    0xea, 0x14, 0x91, 0xe8, 0x1e, 0xb1, 0xa1, 0x45, 0x91, 0x3c, 0xb6, 0xdf, 0xa8, 0x05, 0x3b, 0x91,
    0xb5, 0xd8, 0xf5, 0x75, 0x4e, 0xc2, 0x7d, 0x76, 0x88, 0xbe, 0xad, 0x0f, 0xb2, 0x55, 0x27, 0xe5,
}

// a dynamic block with power-of-two symbol frequencies, so its literal tree
// is a 10-deep caterpillar and decoding it has to leave the 9-bit table.
// 1024 bytes out.
@(private="file")
v_slow := [?]byte{
    0x78, 0x01, 0x05, 0xc1, 0x81, 0x61, 0x04, 0x45, 0x10, 0x04, 0xb1, 0xd8, 0x00, 0xff, 0xed, 0x74,
    0x29, 0xff, 0x7c, 0x90, 0xfa, 0x5d, 0x0a, 0x0b, 0x78, 0x8b, 0xcd, 0x55, 0xb9, 0xca, 0x8b, 0xed,
    0x04, 0x78, 0x52, 0xa1, 0x6b, 0x9c, 0x19, 0xb3, 0x18, 0xe7, 0x86, 0x46, 0xde, 0xbd, 0x0d, 0x00,
    0x84, 0xfe, 0x70, 0xe4, 0xaa, 0x61, 0x79, 0xdb, 0xa9, 0x29, 0x87, 0xb0, 0xe1, 0x86, 0x2a, 0xa5,
    0xfe, 0xe0, 0x8f, 0x51, 0x76, 0x05, 0x95, 0xac, 0xf1, 0x79, 0xd1, 0x20, 0xc4, 0xac, 0x3a, 0x21,
    0x50, 0xfc, 0x3c, 0x0f, 0x77, 0xab, 0xfb, 0x6c, 0x0b, 0xfb, 0x48, 0x7a, 0x93, 0xdc, 0x7c, 0x4c,
    0x77, 0xa7, 0x98, 0x19, 0xe7, 0x77, 0xdf, 0x61, 0xb5, 0xf9, 0x8c, 0xc0, 0x0b, 0xe5, 0xba, 0x7b,
    0x56, 0x6c, 0x32, 0xee, 0x4e, 0xda, 0x91, 0x90, 0xb8, 0x6d, 0x81, 0x71, 0x27, 0xaf, 0xcb, 0x18,
    0x3f, 0xb1, 0x75, 0xea, 0xc1, 0x1b, 0x58, 0x9b, 0x08, 0x3b, 0x76, 0xca, 0xd0, 0x89, 0x16, 0x4d,
    0x7b, 0x9c, 0x46, 0x8c, 0x11, 0x1e, 0xb6, 0x1a, 0x77, 0xaf, 0xe8, 0x73, 0xed, 0xb4, 0x8a, 0x20,
    0x50, 0xe5, 0x37, 0xf4, 0x1f, 0xf4, 0x29, 0xf2, 0x51, 0x5a, 0x61, 0xa0, 0x3d, 0xba, 0xcf, 0x52,
    0x9f, 0x12, 0xab, 0x50, 0x01, 0x00, 0x08, 0x85, 0x16, 0xe1, 0x80, 0x2b, 0xfe, 0x09, 0x90, 0x9e,
    0xc5, 0xed, 0x59, 0x0f, 0x43, 0xec, 0x0b, 0x81, 0x85, 0x0a, 0xa3, 0xac, 0xa2, 0x50, 0x44, 0x3f,
    0x23, 0x2a, 0xe4, 0xa1, 0x83, 0xba, 0xc9, 0x64, 0xaa, 0x22, 0xa6, 0xd5, 0x53, 0xbb, 0xc0, 0x56,
    0x9e, 0x27, 0xc1, 0x5e, 0xdc, 0xb3, 0x62, 0xce, 0xa4, 0x23, 0x8b, 0x76, 0xbd, 0x82, 0x23, 0x94,
    0x39, 0xe6, 0xf5, 0xa3, 0xe2, 0x4b, 0xbe, 0xde, 0x64, 0x29, 0x75, 0xe1, 0x3f, 0xd4, 0x35, 0xf5,
    0x95, 0x31, 0x58, 0xcd, 0x66, 0x7f, 0xfe, 0xdd, 0x62, 0xc7, 0xc2, 0x6a, 0x7e, 0x9c, 0xa6, 0x6e,
    0x9e, 0x91, 0x0f, 0x20, 0xfe, 0x07, 0xbe, 0x54, 0x28, 0x11,
}

// 256 KiB of zeroes -- a ~946:1 expansion, close to what DEFLATE can do at
// all, and therefore the shape a decompression bomb takes.
@(private="file")
v_ratio := [?]byte{
    0x78, 0xda, 0xed, 0xc1, 0x31, 0x01, 0x00, 0x00, 0x00, 0xc2, 0xa0, 0xf5, 0x4f, 0xed, 0x6d, 0x07,
    0xa0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xde,
    0x00, 0x00, 0x3c, 0x00, 0x01,
}

// a bare fixed-Huffman block, no zlib wrapper at all: "Hi PDF".
@(private="file")
v_raw_fixed := [?]byte{
    0xf3, 0xc8, 0x54, 0x08, 0x70, 0x71, 0x03, 0x00,
}

// one literal, then a back-reference 7 bytes before a 1-byte output. The
// literal survives; the match cannot.
@(private="file")
v_bad_distance := [?]byte{
    0x73, 0x04, 0x52, 0x00,
}

// a stored block whose LEN and ones-complement NLEN disagree.
@(private="file")
v_stored_bad_nlen := [?]byte{
    0x78, 0x01, 0x01, 0x05, 0x00, 0x00, 0x00, 0x68, 0x65, 0x6c, 0x6c, 0x6f,
}

// a stored block that promises 32 bytes and supplies 5.
@(private="file")
v_stored_short := [?]byte{
    0x78, 0x01, 0x01, 0x20, 0x00, 0xdf, 0xff, 0x68, 0x65, 0x6c, 0x6c, 0x6f,
}

// BTYPE 3 is reserved; nothing here can be decoded.
@(private="file")
v_btype3 := [?]byte{
    0x07, 0x07, 0x07, 0x07,
}

// a dynamic block whose code-length alphabet assigns four 1-bit codes, which
// no prefix code can satisfy.
@(private="file")
v_oversubscribed := [?]byte{
    0x05, 0x00, 0x92, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
}

// raw DEFLATE flushed but never finished, so no block ever sets BFINAL --
// what a writer that crashed mid-stream leaves behind. 68 bytes out.
@(private="file")
v_no_final := [?]byte{
    0xf2, 0x48, 0xcd, 0xc9, 0xc9, 0xd7, 0x51, 0x08, 0x70, 0x71, 0x53, 0x28, 0xcf, 0x2f, 0xca, 0x49,
    0x51, 0xf4, 0x20, 0x5d, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff,
}

// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------

// zlib_reference inflates with core:compress/zlib, which is strict and
// therefore only usable on the well-formed vectors. Caller frees.
@(private="file")
zlib_reference :: proc(t: ^testing.T, src: []byte) -> []byte {
    buf: bytes.Buffer
    err := zlib.inflate(src, &buf)
    if err != nil {
        testing.expectf(t, false, "core:compress/zlib rejected a vector it should accept: %v", err)
        bytes.buffer_destroy(&buf)
        return nil
    }
    return bytes.buffer_to_bytes(&buf)
}

@(private="file")
expect_decodes_to :: proc(t: ^testing.T, src: []byte, want: string, label: string) {
    got, err := flate_decode(src)
    defer delete(got)
    testing.expectf(t, err == .None, "%s: %v", label, error_string(err))
    testing.expectf(t, string(got) == want, "%s: got %q, want %q", label, string(got), want)
}

// ---------------------------------------------------------------------------
// the four block shapes
// ---------------------------------------------------------------------------

@(test)
test_flate_empty_stream :: proc(t: ^testing.T) {
    // An empty stream is not damage: it inflates to nothing and reports .None.
    out, err := flate_decode(v_empty[:])
    defer delete(out)
    testing.expectf(t, err == .None, "empty stream: %v", error_string(err))
    testing.expect_value(t, len(out), 0)
}

@(test)
test_flate_stored_block :: proc(t: ^testing.T) {
    expect_decodes_to(t, v_stored[:], V_STORED_PLAIN, "stored block")
}

@(test)
test_flate_fixed_huffman :: proc(t: ^testing.T) {
    expect_decodes_to(t, v_fixed[:], V_FIXED_PLAIN, "fixed Huffman")
}

@(test)
test_flate_raw_fixed_huffman :: proc(t: ^testing.T) {
    expect_decodes_to(t, v_raw_fixed[:], "Hi PDF", "bare fixed-Huffman block")
}

@(test)
test_flate_dynamic_huffman :: proc(t: ^testing.T) {
    want := zlib_reference(t, v_dynamic[:])
    defer delete(want)
    got, err := flate_decode(v_dynamic[:])
    defer delete(got)
    testing.expectf(t, err == .None, "dynamic Huffman: %v", error_string(err))
    testing.expect_value(t, len(got), 201)
    testing.expect(t, slice.equal(got, want), "dynamic Huffman output disagrees with core:compress/zlib")
}

@(test)
test_flate_long_codes_take_the_slow_path :: proc(t: ^testing.T) {
    // v_slow's literal tree is 10 deep, so at least one symbol per repetition
    // misses the 9-bit first-level table.
    want := zlib_reference(t, v_slow[:])
    defer delete(want)
    got, err := flate_decode(v_slow[:])
    defer delete(got)
    testing.expectf(t, err == .None, "long codes: %v", error_string(err))
    testing.expect_value(t, len(got), 1024)
    testing.expect(t, slice.equal(got, want), "long-code output disagrees with core:compress/zlib")
}

@(test)
test_flate_agrees_with_zlib :: proc(t: ^testing.T) {
    cases := [][]byte{v_empty[:], v_stored[:], v_fixed[:], v_dynamic[:], v_slow[:], v_ratio[:]}
    for src, i in cases {
        want := zlib_reference(t, src)
        defer delete(want)
        got, err := flate_decode(src)
        defer delete(got)
        testing.expectf(t, err == .None, "case %d: %v", i, error_string(err))
        testing.expectf(t, len(got) == len(want), "case %d: %d bytes, zlib says %d", i, len(got), len(want))
        testing.expectf(t, slice.equal(got, want), "case %d: bytes differ from core:compress/zlib", i)
    }
}

@(test)
test_flate_high_ratio_stream :: proc(t: ^testing.T) {
    // ~946:1. A real stream may legitimately expand this much, so the bomb
    // guard must not be tighter than what DEFLATE can actually produce.
    out, err := flate_decode(v_ratio[:])
    defer delete(out)
    testing.expectf(t, err == .None, "high ratio: %v", error_string(err))
    testing.expect_value(t, len(out), 256 * 1024)
    all_zero := true
    for b in out {
        if b != 0 {
            all_zero = false
            break
        }
    }
    testing.expect(t, all_zero, "high-ratio stream should decode to zeroes")
}

// ---------------------------------------------------------------------------
// the wrapper is where producers go wrong
// ---------------------------------------------------------------------------

@(test)
test_flate_raw_deflate_without_wrapper :: proc(t: ^testing.T) {
    // Some producers write the deflate stream with no zlib header at all.
    expect_decodes_to(t, v_fixed[2:], V_FIXED_PLAIN, "header stripped")
}

@(test)
test_flate_corrupt_header :: proc(t: ^testing.T) {
    src: [len(v_fixed)]byte
    copy(src[:], v_fixed[:])
    src[0] = 0x00 // CM is no longer 8, so FCHECK cannot rescue it either
    expect_decodes_to(t, src[:], V_FIXED_PLAIN, "corrupt CMF")

    copy(src[:], v_fixed[:])
    src[1] ~= 0x01 // FCHECK fails
    expect_decodes_to(t, src[:], V_FIXED_PLAIN, "corrupt FLG")
}

@(test)
test_flate_acrobat_off_by_one :: proc(t: ^testing.T) {
    // A known Acrobat bug drops the first byte of the header.
    expect_decodes_to(t, v_fixed[1:], V_FIXED_PLAIN, "one header byte lost")
}

@(test)
test_flate_leading_whitespace :: proc(t: ^testing.T) {
    src: [4 + len(v_dynamic)]byte
    copy(src[:], []byte{0x0d, 0x0a, 0x20, 0x09})
    copy(src[4:], v_dynamic[:])
    want := zlib_reference(t, v_dynamic[:])
    defer delete(want)
    got, err := flate_decode(src[:])
    defer delete(got)
    testing.expectf(t, err == .None, "leading whitespace: %v", error_string(err))
    testing.expect(t, slice.equal(got, want), "leading whitespace changed the output")
}

@(test)
test_flate_leading_junk :: proc(t: ^testing.T) {
    // Not whitespace, so the header has to be found by scanning.
    src: [7 + len(v_dynamic)]byte
    copy(src[:], transmute([]byte)string("stream\n"))
    copy(src[7:], v_dynamic[:])
    want := zlib_reference(t, v_dynamic[:])
    defer delete(want)
    got, err := flate_decode(src[:])
    defer delete(got)
    testing.expectf(t, err == .None, "leading junk: %v", error_string(err))
    testing.expect(t, slice.equal(got, want), "leading junk changed the output")
}

@(test)
test_flate_trailing_garbage :: proc(t: ^testing.T) {
    // Everything after the final block is somebody else's problem.
    src: [len(v_fixed) + 24]byte
    copy(src[:], v_fixed[:])
    copy(src[len(v_fixed):], transmute([]byte)string("\nendstream\nendobj\n%%EOF\n"))
    expect_decodes_to(t, src[:], V_FIXED_PLAIN, "trailing garbage")
}

// ---------------------------------------------------------------------------
// damage
// ---------------------------------------------------------------------------

@(test)
test_flate_missing_final_block :: proc(t: ^testing.T) {
    // No block ever sets BFINAL, so the decoder simply runs out of input; the
    // bytes it already produced are still good.
    out, err := flate_decode(v_no_final[:])
    defer delete(out)
    testing.expectf(t, err == .None, "no final block: %v", error_string(err))
    want := "Hello, PDF world!Hello, PDF world!Hello, PDF world!Hello, PDF world!"
    testing.expectf(t, string(out) == want, "no final block: got %q", string(out))
}

@(test)
test_flate_stored_block_truncated :: proc(t: ^testing.T) {
    // LEN says 32, five bytes are present. Keep the five.
    out, err := flate_decode(v_stored_short[:])
    defer delete(out)
    testing.expectf(t, err == .None, "short stored block: %v", error_string(err))
    testing.expectf(t, string(out) == "hello", "short stored block: got %q", string(out))
}

@(test)
test_flate_stored_block_bad_nlen :: proc(t: ^testing.T) {
    // LEN and ~LEN disagreeing means these are not really block-header bytes,
    // so nothing is recoverable and nothing is invented.
    out, err := flate_decode(v_stored_bad_nlen[:])
    defer delete(out)
    testing.expectf(t, err == .Bad_Flate_Data, "bad NLEN: %v", error_string(err))
    testing.expect_value(t, len(out), 0)
}

@(test)
test_flate_distance_before_start :: proc(t: ^testing.T) {
    // A match reaching behind the first output byte stops the block; the
    // literal decoded before it is still returned.
    out, err := flate_decode(v_bad_distance[:])
    defer delete(out)
    testing.expectf(t, err == .None, "bad distance: %v", error_string(err))
    testing.expectf(t, string(out) == "A", "bad distance: got %q", string(out))
}

@(test)
test_flate_nothing_recoverable :: proc(t: ^testing.T) {
    empty: []byte
    cases := [][]byte{
        empty,
        transmute([]byte)string("   \r\n\t"),
        v_btype3[:],
    }
    for src, i in cases {
        out, err := flate_decode(src)
        defer delete(out)
        testing.expectf(t, err == .Bad_Flate_Data, "case %d: %v", i, error_string(err))
        testing.expect_value(t, len(out), 0)
    }
}

@(test)
test_flate_every_truncation :: proc(t: ^testing.T) {
    // Cutting a stream short must never crash, must never produce bytes the
    // whole stream would not have produced, and must never lose ground as more
    // input arrives.
    streams := [][]byte{v_dynamic[:], v_slow[:], v_fixed[:], v_stored[:], v_ratio[:]}
    for src in streams {
        full, ferr := flate_decode(src)
        testing.expect(t, ferr == .None, "full stream should decode")
        prev := 0
        for n in 1 ..= len(src) {
            got, _ := flate_decode(src[:n])
            testing.expectf(t, len(got) <= len(full),
                "truncating to %d/%d produced %d bytes, more than the whole stream's %d",
                n, len(src), len(got), len(full))
            if len(got) <= len(full) {
                testing.expectf(t, slice.equal(got, full[:len(got)]),
                    "truncating to %d/%d produced bytes the whole stream does not", n, len(src))
            }
            testing.expectf(t, len(got) >= prev,
                "truncating to %d/%d produced %d bytes, fewer than %d bytes of input gave (%d)",
                n, len(src), len(got), n - 1, prev)
            prev = len(got)
            delete(got)
        }
        testing.expect_value(t, prev, len(full))
        delete(full)
    }
}

@(test)
test_flate_bit_flips :: proc(t: ^testing.T) {
    // Every single-bit corruption of a dynamic stream: no crash, no hang, no
    // allocation the caller cannot free.
    src: [len(v_dynamic)]byte
    limit := flate_output_limit(len(v_dynamic))
    for i in 0 ..< len(v_dynamic) {
        for bit in u32(0) ..< 8 {
            copy(src[:], v_dynamic[:])
            src[i] ~= byte(1 << bit)
            out, err := flate_decode(src[:])
            testing.expectf(t, len(out) <= limit, "bit flip %d/%d exceeded the output cap", i, bit)
            testing.expect(t, err == .None || err == .Bad_Flate_Data, "unexpected error kind")
            testing.expect(t, err != .Bad_Flate_Data || len(out) == 0, "an error must come with no output")
            delete(out)
        }
    }
}

@(test)
test_flate_malformed_never_crashes :: proc(t: ^testing.T) {
    cases := [][]byte{
        v_oversubscribed[:],
        v_btype3[:],
        v_stored_bad_nlen[:],
        v_stored_short[:],
        v_bad_distance[:],
        {0x78},
        {0x78, 0x9c},
        {0xff},
        {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
        {0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
    }
    for src, i in cases {
        out, err := flate_decode(src)
        testing.expectf(t, err == .None || err == .Bad_Flate_Data, "case %d: %v", i, error_string(err))
        delete(out)
    }
}

// ---------------------------------------------------------------------------
// the pieces, on their own
// ---------------------------------------------------------------------------

@(test)
test_flate_output_limit :: proc(t: ^testing.T) {
    // Small streams get a floor so that a legitimate tiny stream is never
    // clipped; large ones are bounded by what DEFLATE can actually expand to;
    // everything is bounded absolutely.
    testing.expect_value(t, flate_output_limit(0), FLATE_MIN_OUTPUT)
    testing.expect_value(t, flate_output_limit(1), FLATE_MIN_OUTPUT)
    testing.expect_value(t, flate_output_limit(1024), 1024 * FLATE_MAX_RATIO)
    testing.expect_value(t, flate_output_limit(1 << 30), FLATE_MAX_OUTPUT)
    testing.expect_value(t, flate_output_limit(max(int) / 2), FLATE_MAX_OUTPUT)
    testing.expect(t, FLATE_MAX_RATIO > 1032,
        "DEFLATE can expand about 1032:1, so a tighter cap would reject valid streams")
}

@(test)
test_flate_inflate_stops_at_the_limit :: proc(t: ^testing.T) {
    // The bomb guard, exercised directly: 256 KiB of output, 100 bytes allowed.
    out, done := flate_inflate(v_ratio[2:], 100)
    defer delete(out)
    testing.expect_value(t, len(out), 100)
    testing.expect(t, !done, "a stream cut off by the output cap has not finished")
}

@(test)
test_flate_header_len :: proc(t: ^testing.T) {
    Case :: struct {
        input:  []byte,
        n:      int,
        ok:     bool,
        why:    string,
    }
    cases := []Case{
        {{0x78, 0x01}, 2, true,  "level 0/1"},
        {{0x78, 0x9c}, 2, true,  "default level"},
        {{0x78, 0xda}, 2, true,  "level 9"},
        {{0x68, 0xde}, 2, true,  "8K window"},
        {{0x78, 0x9d}, 0, false, "FCHECK is wrong"},
        {{0x77, 0x9c}, 0, false, "CM is not DEFLATE"},
        {{0x88, 0x9c}, 0, false, "CINFO asks for a 64K window"},
        {{0x78}, 0, false, "one byte is not a header"},
        {{}, 0, false, "no bytes at all"},
        {{0x78, 0xbb, 0x01, 0x02, 0x03, 0x04}, 6, true, "FDICT adds a 4-byte identifier"},
        {{0x78, 0xbb, 0x01}, 0, false, "FDICT set but the identifier is cut off"},
    }
    for c in cases {
        n, ok := flate_header_len(c.input)
        testing.expectf(t, ok == c.ok && n == c.n, "%s: got (%d, %v), want (%d, %v)", c.why, n, ok, c.n, c.ok)
    }
}

@(test)
test_flate_build_rejects_oversubscribed :: proc(t: ^testing.T) {
    h := new(Flate_Huffman)
    defer free(h)

    // Four one-bit codes: no prefix code can hold them.
    testing.expect(t, !flate_build(h, []u8{1, 1, 1, 1}), "over-subscribed code accepted")
    // A length past 15 bits cannot be expressed.
    testing.expect(t, !flate_build(h, []u8{16, 1}), "a 16-bit code length was accepted")
    // Incomplete codes are legal and do occur: a distance alphabet with a
    // single symbol is the usual example.
    testing.expect(t, flate_build(h, []u8{1}), "single-symbol code rejected")
    // An alphabet with nothing in it builds; decoding against it just fails.
    testing.expect(t, flate_build(h, []u8{0, 0, 0}), "empty alphabet rejected")
    testing.expect(t, flate_build(h, []u8{1, 1}), "complete two-symbol code rejected")
}

@(test)
test_flate_build_fixed_tables :: proc(t: ^testing.T) {
    st := new(Flate_State)
    defer free(st)
    testing.expect(t, flate_fixed_tables(st), "the fixed tables must always build")

    // RFC 1951 3.2.6: literal 0 is 00110000, so the stream bits are 0,0,0,0,1,1,0,0
    // and the 9-bit peek that selects it is 0b000110000 reversed = 0x00C.
    st.br = {data = []byte{0x0c, 0x00}}
    sym, ok := flate_sym(&st.br, &st.lit)
    testing.expect(t, ok, "fixed literal 0 should decode")
    testing.expect_value(t, sym, u16(0))

    // End-of-block is seven zero bits.
    st.br = {data = []byte{0x00, 0x00}}
    sym2, ok2 := flate_sym(&st.br, &st.lit)
    testing.expect(t, ok2, "end-of-block should decode")
    testing.expect_value(t, sym2, u16(256))
}

@(test)
test_flate_reverse :: proc(t: ^testing.T) {
    testing.expect_value(t, flate_reverse(0b1, 1), u32(0b1))
    testing.expect_value(t, flate_reverse(0b10, 2), u32(0b01))
    testing.expect_value(t, flate_reverse(0b00110000, 8), u32(0b00001100))
    testing.expect_value(t, flate_reverse(0, 9), u32(0))
    testing.expect_value(t, flate_reverse(0b111111111, 9), u32(0b111111111))
}

@(test)
test_flate_bit_reader_runs_dry :: proc(t: ^testing.T) {
    br := Flate_Bits{data = []byte{0xff}}
    v, ok := flate_get(&br, 8)
    testing.expect(t, ok, "eight bits are available")
    testing.expect_value(t, v, u32(0xff))
    _, ok2 := flate_get(&br, 1)
    testing.expect(t, !ok2, "the ninth bit is not")
    v3, ok3 := flate_get(&br, 0)
    testing.expect(t, ok3 && v3 == 0, "asking for no bits always succeeds")
}
// three blocks in one stream: content, a sync-flush boundary, then the final
// block. "first half, second half."
@(private="file")
v_multi_block := [?]byte{
    0x78, 0xda, 0x4a, 0xcb, 0x2c, 0x2a, 0x2e, 0x51, 0xc8, 0x48, 0xcc, 0x49, 0xd3, 0x51, 0x00, 0x00,
    0x00, 0x00, 0xff, 0xff, 0x2b, 0x4e, 0x4d, 0xce, 0xcf, 0x4b, 0x01, 0x73, 0xf4, 0x00, 0x6e, 0x7b,
    0x08, 0x95,
}

// a fixed block using distance symbol 30, which RFC 1951 never assigns.
@(private="file")
v_bad_dist_code := [?]byte{
    0xab, 0xa8, 0x04, 0xbe, 0x2a, 0x00,
}

// a fixed block using literal/length symbol 286, likewise unassigned.
@(private="file")
v_bad_len_code := [?]byte{
    0xab, 0xa8, 0x1c, 0x03, 0x00,
}

// FDICT is set, so a 4-byte dictionary identifier follows the header and the
// data cannot be decoded without a dictionary we do not have.
@(private="file")
v_fdict := [?]byte{
    0x78, 0x20, 0x11, 0x22, 0x33, 0x44, 0xcb, 0x4b, 0x4d, 0x4d, 0x29, 0x56, 0x48, 0x54, 0x48, 0xc9,
    0x4c, 0x2e, 0xc9, 0xcc, 0xcf, 0x4b, 0x2c, 0xaa, 0x04, 0x00, 0x3f, 0x3b, 0x06, 0xe7,
}

@(test)
test_flate_multiple_blocks :: proc(t: ^testing.T) {
    expect_decodes_to(t, v_multi_block[:], "first half, second half.", "three blocks")
}

@(test)
test_flate_preset_dictionary_header :: proc(t: ^testing.T) {
    // FDICT pushes the data four bytes further along. Nothing else changes.
    expect_decodes_to(t, v_fdict[:], "needs a dictionary", "FDICT header")
}

@(test)
test_flate_unassigned_symbols :: proc(t: ^testing.T) {
    // Symbols 286/287 and distance symbols 30/31 exist in the alphabets but
    // have no meaning; hitting one ends the block, keeping what came before.
    out1, err1 := flate_decode(v_bad_dist_code[:])
    defer delete(out1)
    testing.expectf(t, err1 == .None, "distance symbol 30: %v", error_string(err1))
    testing.expectf(t, string(out1) == "xy", "distance symbol 30: got %q", string(out1))

    out2, err2 := flate_decode(v_bad_len_code[:])
    defer delete(out2)
    testing.expectf(t, err2 == .None, "length symbol 286: %v", error_string(err2))
    testing.expectf(t, string(out2) == "xy", "length symbol 286: got %q", string(out2))
}

@(test)
test_flate_random_bytes :: proc(t: ^testing.T) {
    // Nothing here is a valid stream; the point is that a few thousand
    // arbitrary byte sequences neither crash, hang, nor leak.
    seed := u64(0x9e3779b97f4a7c15)
    src: [64]byte
    for round in 0 ..< 2000 {
        n := 1 + round % len(src)
        for i in 0 ..< n {
            seed = seed * 6364136223846793005 + 1442695040888963407
            src[i] = byte(seed >> 33)
        }
        out, err := flate_decode(src[:n])
        testing.expect(t, err == .None || err == .Bad_Flate_Data, "unexpected error kind")
        testing.expect(t, len(out) <= flate_output_limit(n), "output cap ignored")
        delete(out)
    }
}

@(test)
test_flate_honours_the_allocator :: proc(t: ^testing.T) {
    // An arena that cannot free: every attempt, and the result, have to come
    // out of the buffer the caller supplied and nowhere else.
    backing := make([]byte, 256 * 1024)
    defer delete(backing)
    arena: mem.Arena
    mem.arena_init(&arena, backing)

    want := zlib_reference(t, v_dynamic[:])
    defer delete(want)

    out, err := flate_decode(v_dynamic[:], mem.arena_allocator(&arena))
    testing.expectf(t, err == .None, "arena decode: %v", error_string(err))
    testing.expect(t, slice.equal(out, want), "arena decode produced different bytes")
    testing.expect(t, arena.offset > 0, "nothing was taken from the arena")
}

// Rejected attempts are the hazard peculiar to this codec. flate_decode reads
// the same bytes several ways -- a header at the front, a header a little way
// in, raw DEFLATE, Acrobat's dropped first byte -- and each attempt allocates
// an output buffer before it can know whether the guess will yield anything.
// Dropping the losers with `delete` is right against the heap and does nothing
// whatever against a Document arena, which answers .Free with
// .Mode_Not_Implemented; the buffers would simply sit there until `close`.
// apply_filter is what keeps that from happening, by decoding into scratch and
// copying only the winner across, so this test is on the boundary rather than
// on flate_decode itself -- flate_decode is allowed its guesses, it just may
// not make them with the Document's memory.
@(test)
test_flate_speculation_stays_out_of_the_document_arena :: proc(t: ^testing.T) {
    doc := filter_test_doc()
    defer filter_test_doc_free(doc)

    // A well-formed zlib header over bytes that are not DEFLATE at all, so
    // every candidate offset is tried and every one of them comes back empty.
    broken := [?]byte{0x78, 0x9c, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07}
    before := document_memory_used(doc)
    out, err := apply_filter(doc, "FlateDecode", broken[:], nil, doc.allocator)
    testing.expect_value(t, err, Error.Bad_Flate_Data)
    testing.expect_value(t, len(out), 0)
    testing.expectf(t, document_memory_used(doc) == before,
        "%d bytes stranded in the arena by eight bytes that decoded to nothing",
        document_memory_used(doc) - before)

    // And a stream that does decode, but only after the wrong guesses have
    // been made and thrown away: no zlib wrapper, so the offset is a guess.
    before = document_memory_used(doc)
    out2, err2 := apply_filter(doc, "FlateDecode", v_raw_fixed[:], nil, doc.allocator)
    testing.expect_value(t, err2, Error.None)
    testing.expect_value(t, string(out2), "Hi PDF")
    growth := document_memory_used(doc) - before
    testing.expectf(t, growth <= len(out2) + FT_ARENA_SLACK,
        "the arena grew by %d bytes for %d bytes of output", growth, len(out2))
}