summaryrefslogtreecommitdiff
path: root/tools/buildtools/makeyar.py
blob: d8d18d16accc02782408dd6cfbf03f03049ac149 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
#!/usr/bin/env python3

# SPDX-FileCopyrightText: © 2023-2024 ZeldaRET
# SPDX-License-Identifier: MIT

# Program to generate compressed yar (Yaz0 ARchive) files.
#
# The program expects an .o elf file and outputs a raw yar binary file and a
# "symbols" elf.
#
# A yar file consists of multiple Yaz0 files compressed individually. The
# archive begins with a header of non-fixed size, which describes the
# location of each individual Yaz0 block within the archive itself. This
# header is followed by each Yaz0 file.
#
# The first word (a 4 byte group) of the header indicates the size in bytes of
# the header itself (also describes the offset of the first Yaz0 block). The
# rest of the header consists of words describing the offsets of each Yaz0
# block relative to the end of the header, because the first Yaz0
# block is omitted from the offsets in the header.
#
# Each Yaz0 block is 0xFF-padded to a multiple of 0x10 in size.
#
# The entire archive is 0-padded to a multiple of 0x10 in size.
#
# The program works by compressing each .data symbol in the input elf file as
# its own Yaz0 compressed file, appending them in order for the generated
# archive. Other elf sections are ignored for the resulting yar file.
#
# The program also outputs an elf file that's identical to the elf input,
# but with its .data section zero'ed out completely. This "symbols" elf can be
# used for referencing each symbol as the whole file were completely
# uncompressed.

from __future__ import annotations

import argparse
import dataclasses
from pathlib import Path
import struct
import crunch64

from elftools.elf.elffile import ELFFile
from elftools.elf.sections import SymbolTableSection


def write_word_as_bytes(buff: bytearray, offset: int, word: int):
    struct.pack_into(f">I", buff, offset, word)


@dataclasses.dataclass
class Symbol:
    name: str
    offset: int
    size: int


def get_data_from_elf(elf_path: Path) -> tuple[bytearray, list[Symbol], int]:
    uncompressed_data = bytearray()
    symbol_list: list[Symbol] = []
    data_offset = -1

    with elf_path.open("rb") as elfFile:
        elf = ELFFile(elfFile)
        for section in elf.iter_sections():
            if section.name == ".data":
                assert len(uncompressed_data) == 0
                uncompressed_data.extend(section.data())
                assert len(uncompressed_data) == section["sh_size"]
                data_offset = section["sh_offset"]
            elif section.name == ".symtab":
                assert isinstance(section, SymbolTableSection)
                for sym in section.iter_symbols():
                    if sym["st_shndx"] == "SHN_UNDEF":
                        continue
                    if sym["st_info"]["type"] != "STT_OBJECT":
                        continue
                    symbol_list.append(
                        Symbol(sym.name, sym["st_value"], sym["st_size"])
                    )
    return uncompressed_data, symbol_list, data_offset


def align_16(val: int) -> int:
    return (val + 0xF) & ~0xF


def create_archive(
    uncompressed_data: bytearray, symbol_list: list[Symbol]
) -> bytearray:
    archive = bytearray()

    first_entry_offset = (len(symbol_list) + 1) * 4

    # Fill with zeroes until the compressed data start
    archive.extend([0] * first_entry_offset)

    write_word_as_bytes(archive, 0, first_entry_offset)

    offset = first_entry_offset

    i = 0
    for sym in symbol_list:
        uncompressed_size = sym.size
        uncompressed_size_aligned = align_16(uncompressed_size)

        input_buf = uncompressed_data[sym.offset : sym.offset + uncompressed_size]
        # Make sure to pad each entry to a 0x10 boundary
        if uncompressed_size_aligned > uncompressed_size:
            input_buf.extend([0x00] * (uncompressed_size_aligned - uncompressed_size))

        compressed = bytearray(crunch64.yaz0.compress(input_buf))
        compressed_size = len(compressed)

        # Pad to 0x10
        compressed_size_aligned = align_16(compressed_size)
        if compressed_size_aligned > compressed_size:
            compressed.extend([0xFF] * (compressed_size_aligned - compressed_size))

        archive.extend(compressed)

        if i > 0:
            write_word_as_bytes(archive, i * 4, offset - first_entry_offset)

        i += 1
        offset += len(compressed)

    write_word_as_bytes(archive, i * 4, offset - first_entry_offset)

    archive_len = len(archive)
    archive_len_aligned = align_16(archive_len)
    if archive_len_aligned > archive_len:
        archive.extend([0x00] * (archive_len_aligned - archive_len))

    return archive


def main():
    parser = argparse.ArgumentParser(
        description="Program to generate compressed yar (Yaz0 ARchive) files from a built C file. Said file must only contain data symbols that do not reference other symbols (i.e. textures)."
    )
    parser.add_argument("in_file", help="Path to built .o file")
    parser.add_argument(
        "out_bin", help="Output path for the generated compressed yar binary"
    )
    parser.add_argument("out_sym", help="Output path for the generated syms elf file")

    args = parser.parse_args()

    in_path = Path(args.in_file)
    out_bin_path = Path(args.out_bin)
    out_sym_path = Path(args.out_sym)

    # Delete output files if they already exist
    out_bin_path.unlink(missing_ok=True)
    out_sym_path.unlink(missing_ok=True)

    elf_bytes = bytearray(in_path.read_bytes())

    uncompressed_data, symbol_list, data_offset = get_data_from_elf(in_path)
    assert len(uncompressed_data) > 0
    assert len(symbol_list) > 0
    assert data_offset > 0

    archive = create_archive(uncompressed_data, symbol_list)

    # Write the compressed archive file as a raw binary
    out_bin_path.write_bytes(archive)

    # Zero out data
    for i in range(data_offset, data_offset + len(uncompressed_data)):
        elf_bytes[i] = 0

    out_sym_path.write_bytes(elf_bytes)


if __name__ == "__main__":
    main()