Compresses DOOM WAD files using the LZSS algorthm from the Atari Jaguar port

WadCompressor.py

from io import BytesIO
import struct
import sys
from typing import BinaryIO, cast, override

# Based on LZSS.C by Haruhiko Okumura
# https://web.archive.org/web/19990203141013/http://oak.oakland.edu/pub/simtelnet/msdos/arcutils/lz_comp2.zip
# Compared to the original, the version used in DOOM is a bit different:
# - The flag bits are inverted (0 = literal, 1 = compressed)
# - The offset is encoded differently
#     - In the original, the first byte is the lower 8 bits, and the upper 4 bits of the second byte are the upper 4 bits
#     - Here, the first byte is the upper 8 bits, and the upper 4 bits of the second byte are the lower 4 bits
# - The offset is semantically different
#     - In the original, it's an absolute index in the ring buffer
#     - Here, it's a backward offset from the previously read byte
# - The range of possible lengths for the compressed chunks is different
#     - In the original, possible values are 3 to 18, encoded in 4 bits as 0 to 15
#     - Here, possible values are 2 to 16, encoded in 4 bits as 1 to 15
# - There is a terminating marker at the end of the encoded stream, which is a compressed chunk with length 0
# - The ring buffer is not initialized to any particular value, so trying to read from the parts that haven't been written yet is undefined behavior
#     - In the ZDoom code, it will literally read from uninitialized memory...

N: int = 4096
F: int = 16
THRESHOLD: int = 1

def getc(io: BinaryIO) -> int | None:
    r = io.read(1)
    if len(r) == 0:
        return None
    else:
        return r[0]

class EncodeBuffer:
    lson: list[int]
    rson: list[int]
    dad: list[int]

    text_buf: list[int]

    match_position: int
    match_length: int

    def __init__(self) -> None:
        self.lson = [0 for _ in range(N + 1)]
        self.rson = [0 for _ in range(N + 258)]
        self.dad = [0 for _ in range(N + 1)]

        for i in range(N + 1, (N + 257) + 1):
            self.rson[i] = N

        for i in range(N):
            self.dad[i] = N

        self.text_buf = [256 for _ in range(N + F - 1)]

        self.match_position = 0
        self.match_length = 0

    def insert_node(self, r: int) -> None:
        cmp: int = 1
        p: int = N + 1 + self.text_buf[r]

        self.rson[r] = N
        self.lson[r] = N
        self.match_length = 0

        while True:
            if cmp >= 0:
                if self.rson[p] != N:
                    p = self.rson[p]
                else:
                    self.rson[p] = r
                    self.dad[r] = p
                    return
            else:
                if self.lson[p] != N:
                    p = self.lson[p]
                else:
                    self.lson[p] = r
                    self.dad[r] = p
                    return

            i: int = 1
            while i < F:
                cmp = self.text_buf[r + i] - self.text_buf[p + i]
                if cmp != 0:
                    break
                i += 1

            if i > self.match_length:
                self.match_position = p
                self.match_length = i
                if i >= F:
                    break

        self.dad[r] = self.dad[p]
        self.lson[r] = self.lson[p]
        self.rson[r] = self.rson[p]

        self.dad[self.lson[p]] = r
        self.dad[self.rson[p]] = r

        if self.rson[self.dad[p]] == p:
            self.rson[self.dad[p]] = r
        else:
            self.lson[self.dad[p]] = r

        self.dad[p] = N

    def delete_node(self, p: int) -> None:
        q: int = 0

        if self.dad[p] == N:
            return

        if self.rson[p] == N:
            q = self.lson[p]
        elif self.lson[p] == N:
            q = self.rson[p]
        else:
            q = self.lson[p]
            if self.rson[q] != N:
                q = self.rson[q]
                while self.rson[q] != N:
                    q = self.rson[q]

                self.rson[self.dad[q]] = self.lson[q]
                self.dad[self.lson[q]] = self.dad[q]

                self.lson[q] = self.lson[p]
                self.dad[self.lson[p]] = q

            self.rson[q] = self.rson[p]
            self.dad[self.rson[p]] = q

        self.dad[q] = self.dad[p]

        if self.rson[self.dad[p]] == p:
            self.rson[self.dad[p]] = q
        else:
            self.lson[self.dad[p]] = q

        self.dad[p] = N

def encode(input: BinaryIO, output: BinaryIO) -> tuple[int, int]:
    textsize: int = 0
    codesize: int = 0

    buffer: EncodeBuffer = EncodeBuffer()

    code_buf = [0 for _ in range(17)]

    code_buf_ptr: int = 1
    mask: int = 1

    s: int = 0
    r: int = N - F

    last_match_length: int = 0

    c: int | None = None

    len: int = 0
    while len < F:
        c = getc(input)
        if c is None:
            break
        buffer.text_buf[r + len] = c
        len += 1

    textsize = len
    if len == 0:
        output.write(b"\x01\x00\x00")
        return (0, 3)

    for i in range(1, F + 1):
        buffer.insert_node(r - i)

    buffer.insert_node(r)

    while len > 0:
        if buffer.match_length > len:
            buffer.match_length = len

        if buffer.match_length <= THRESHOLD:
            buffer.match_length = 1
            code_buf[code_buf_ptr] = buffer.text_buf[r]
            code_buf_ptr += 1
        else:
            code_buf[0] |= mask

            match_offset: int = (r - buffer.match_position - 1) % N

            code_buf[code_buf_ptr] = (match_offset >> 4) & 0xff
            code_buf_ptr += 1

            code_buf[code_buf_ptr] = (((match_offset & 0x0f) << 4) | (buffer.match_length - THRESHOLD)) & 0xff
            code_buf_ptr += 1

        mask = (mask << 1) & 0xff
        if mask == 0:
            output.write(bytes(code_buf[:code_buf_ptr]))
            codesize += code_buf_ptr
            code_buf[0] = 0
            code_buf_ptr = 1
            mask = 1

        last_match_length = buffer.match_length

        i: int = 0
        while i < last_match_length:
            c = getc(input)
            if c is None:
                break

            buffer.delete_node(s)
            buffer.text_buf[s] = c
            if s < F - 1:
                buffer.text_buf[s + N] = c

            s = (s + 1) & (N - 1)
            r = (r + 1) & (N - 1)

            buffer.insert_node(r)

            i += 1

        textsize += i

        while i < last_match_length:
            buffer.delete_node(s)
            s = (s + 1) & (N - 1)
            r = (r + 1) & (N - 1)
            len -= 1
            if len > 0:
                buffer.insert_node(r)

            i += 1

    code_buf[0] |= mask

    code_buf[code_buf_ptr] = 0
    code_buf_ptr += 1

    code_buf[code_buf_ptr] = 0
    code_buf_ptr += 1

    output.write(bytes(code_buf[:code_buf_ptr]))
    codesize += code_buf_ptr

    return (textsize, codesize)

def decode(input: BinaryIO, output: BinaryIO) -> bool:
    text_buf: list[int] = [0 for _ in range(N + F - 1)]

    i: int = 0
    j: int = 0

    r: int = N - F

    c: int | None = None

    flags: int = 0

    while True:
        flags >>= 1
        if (flags & 256) == 0:
            c = getc(input)
            if c is None:
                return False
            flags = c | 0xff00

        if not (flags & 1):
            c = getc(input)
            if c is None:
                return False
            output.write(c.to_bytes())
            text_buf[r] = c
            r = (r + 1) & (N - 1)
        else:
            c = getc(input)
            if c is None:
                return False
            i = c

            c = getc(input)
            if c is None:
                return False
            j = c

            i = (i << 4) | (j >> 4)
            j = (j & 0x0f) + THRESHOLD

            if j == 1:
                return True

            start_write_pos: int = r

            for k in range(j):
                c = text_buf[(start_write_pos - i - 1 + k) % N]
                output.write(c.to_bytes())
                text_buf[r] = c
                r = (r + 1) & (N - 1)

class LimitedIO(BinaryIO):
    inner: BinaryIO
    offset: int
    length: int

    def __init__(self, inner: BinaryIO, length: int) -> None:
        self.inner = inner
        self.offset = inner.tell()
        self.length = length

    @property
    @override
    def mode(self) -> str:
        return self.inner.mode

    @property
    @override
    def name(self) -> str:
        return self.inner.name

    @override
    def close(self) -> None:
        self.inner.close()

    @property
    @override
    def closed(self) -> bool:
        return self.inner.closed

    @override
    def fileno(self) -> int:
        return self.inner.fileno()

    @override
    def flush(self) -> None:
        self.inner.flush()

    @override
    def isatty(self) -> bool:
        return self.inner.isatty()

    @override
    def read(self, n: int = -1) -> bytes:
        length_left = self.length + self.offset - self.inner.tell()

        if n < 0:
            n = length_left
        else:
            n = min(n, length_left)

        return self.inner.read(n)


    @override
    def readable(self) -> bool:
        return self.inner.readable()

    @override
    def readline(self, limit: int = -1) -> bytes:
        raise NotImplementedError()

    @override
    def readlines(self, hint: int = -1) -> list[bytes]:
        raise NotImplementedError()

    @override
    def seek(self, offset: int, whence: int = 0) -> int:
        end = self.offset + self.length

        if whence == 1:
            offset += self.inner.tell()
        elif whence == 2:
            offset = end - offset

        offset = max(self.offset, min(self.offset + offset, end))

        self.inner.seek(offset)

        return offset

    @override
    def seekable(self) -> bool:
        return self.inner.seekable()

    @override
    def tell(self) -> int:
        return self.inner.tell() - self.offset

    @override
    def truncate(self, size: int | None = None) -> int:
        raise NotImplementedError()

    @override
    def writable(self) -> bool:
        return False

    @override
    def writelines(self, lines) -> None:  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
        raise NotImplementedError()

    @override
    def __exit__(self, type, value, traceback) -> None:  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
        pass

    @override
    def write(self, s) -> int:  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
        raise NotImplementedError()

    @override
    def __enter__(self) -> BinaryIO:
        return self

class WadReader:
    _file: BinaryIO
    _lumps: list[tuple[str, int, int, bool]]
    _iwad: bool

    def __init__(self, file_obj: BinaryIO) -> None:
        self._file = file_obj

        signature = file_obj.read(4)
        if signature == b"IWAD":
            self._iwad = True
        elif signature == b"PWAD":
            self._iwad = False
        else:
            raise ValueError(f"Invalid WAD file: expected PWAD or IWAD signature, got {signature!r}")

        lump_count = cast(int, struct.unpack("<I", file_obj.read(4))[0])
        directory_offset = cast(int, struct.unpack("<I", file_obj.read(4))[0])

        file_obj.seek(directory_offset)

        self._lumps = []

        for _ in range(lump_count):
            offset = cast(int, struct.unpack("<I", file_obj.read(4))[0])
            length = cast(int, struct.unpack("<I", file_obj.read(4))[0])
            name_bytes = file_obj.read(8)

            compressed = bool(name_bytes[0] & 0x80)

            name = name_bytes.decode("ascii").rstrip("\x00")

            self._lumps.append((name, offset, length, compressed))

    def get_lump(self, index: int) -> BinaryIO:
        if index < 0 or index >= len(self._lumps):
            raise IndexError(f"Lump index {index} out of range (0-{len(self._lumps) - 1})")

        _, offset, length, compressed = self._lumps[index]

        self._file.seek(offset)

        if compressed:
            compressed_data = LimitedIO(self._file, length)

            output_stream = BytesIO()
            decode(compressed_data, output_stream)

            output_stream.seek(0)
            return output_stream
        else:
            return LimitedIO(self._file, length)

    def get_lump_by_name(self, name: str) -> BinaryIO:

        for i, (lump_name, _, _, _) in enumerate(self._lumps):
            if lump_name == name:
                return self.get_lump(i)

        raise ValueError(f"Lump '{name}' not found in WAD file")

    @property
    def lump_count(self) -> int:
        return len(self.lumps)

    @property
    def lumps(self) -> list[tuple[str, int, int, bool]]:
        return self._lumps

    @property
    def iwad(self) -> bool:
        return self._iwad

class WadWriter:
    _file: BinaryIO
    _lumps: list[tuple[str, int, int, bool]]

    def __init__(self, file_obj: BinaryIO, iwad: bool = False) -> None:
        self._file = file_obj
        self._lumps = []

        if iwad:
            self._file.write(b"IWAD")
        else:
            self._file.write(b"PWAD")
        self._file.write(b"\x00\x00\x00\x00")
        self._file.write(b"\x00\x00\x00\x00")

    def add_lump(self, data_file_obj: BinaryIO | str | None, name: str, compressed: bool = False) -> None:
        if len(name) > 8:
            raise ValueError(f"Lump name must be at most 8 characters, got {len(name)}")

        for char in name:
            if ord(char) > 127:
                raise ValueError(f"Lump name contains non-ASCII character: {char}")

        lump_offset = self._file.tell()

        data: BinaryIO

        if isinstance(data_file_obj, str):
            data = BytesIO(data_file_obj.encode("ascii"))
        elif data_file_obj is None:
            data = BytesIO()
        else:
            data = data_file_obj

        lump_length: int

        if compressed:
            lump_length, _ = encode(data, self._file)
        else:
            while True:
                chunk = data.read(8192)
                if not chunk:
                    break
                self._file.write(chunk)
            lump_length = self._file.tell() - lump_offset

        self._lumps.append((name, lump_offset, lump_length, compressed))

    def close_wad(self) -> None:
        dir_offset = self._file.tell()

        self._file.seek(4)
        self._file.write(struct.pack("<I", len(self._lumps)))
        self._file.write(struct.pack("<I", dir_offset))

        self._file.seek(0, 2)

        for name, offset, length, compressed in self._lumps:
            padded_name = name.encode("ascii").ljust(8, b"\x00")
            self._file.write(struct.pack("<II", offset, length))

            if compressed:
                self._file.write((padded_name[0] | 0x80).to_bytes())
                self._file.write(padded_name[1:])
            else:
                self._file.write(padded_name)

in_path = sys.argv[1]
out_path = sys.argv[2]

with open(in_path, "rb") as in_file:
    with open(out_path, "wb") as out_file:
        reader = WadReader(in_file)
        writer = WadWriter(out_file, reader.iwad)

        wasCompress = False

        for i, lump in enumerate(reader.lumps):
            name = lump[0]
            print(name)

            # Add other lumps to compress here
            compress = name in ["TEXTMAP", "THINGS", "LINEDEFS", "VERTEXES", "SEGS", "SSECTORS", "NODES", "SECTORS", "REJECT", "BLOCKMAP", "SCRIPTS", "ZSCRIPT"]

            if compress:
                print("    Compressing...")

            writer.add_lump(reader.get_lump(i), name, compress)
            wasCompress = compress

        if wasCompress:
            writer.add_lump(None, "ENDOFWAD")

        writer.close_wad()