mirror of
https://github.com/zeldaret/oot.git
synced 2024-12-05 01:06:37 +00:00
0ac4448d99
* Archive compression support + small cleanups * UNSET spec and PR comment * UNSET -> SYMS * Syms comment * PR review * remove stderr * Format * format2 * Remove trailing ,s
285 lines
9.1 KiB
Python
Executable file
285 lines
9.1 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
# SPDX-FileCopyrightText: 2024 zeldaret
|
|
# SPDX-License-Identifier: CC0-1.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
import dataclasses
|
|
import time
|
|
import multiprocessing
|
|
import multiprocessing.pool
|
|
|
|
import crunch64
|
|
|
|
import dmadata
|
|
|
|
|
|
def align(v: int):
|
|
v += 0xF
|
|
return v // 0x10 * 0x10
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class RomSegment:
|
|
vrom_start: int
|
|
vrom_end: int
|
|
is_compressed: bool
|
|
is_syms: bool
|
|
data: memoryview | None
|
|
data_async: multiprocessing.pool.AsyncResult | None
|
|
|
|
@property
|
|
def uncompressed_size(self):
|
|
return self.vrom_end - self.vrom_start
|
|
|
|
|
|
# Make interrupting the compression with ^C less jank
|
|
# https://stackoverflow.com/questions/72967793/keyboardinterrupt-with-python-multiprocessing-pool
|
|
def set_sigint_ignored():
|
|
import signal
|
|
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
|
|
|
|
def compress_rom(
|
|
rom_data: memoryview,
|
|
dmadata_start: int,
|
|
compress_entries_indices: set[int],
|
|
n_threads: int = None,
|
|
):
|
|
"""
|
|
rom_data: the uncompressed rom data
|
|
dmadata_start: the offset in the rom where the dmadata starts
|
|
compress_entries_indices: the indices in the dmadata of the segments that should be compressed
|
|
n_threads: how many cores to use for compression
|
|
"""
|
|
|
|
# Segments of the compressed rom (not all are compressed)
|
|
compressed_rom_segments: list[RomSegment] = []
|
|
|
|
dma_entries = dmadata.read_dmadata(rom_data, dmadata_start)
|
|
# We sort the DMA entries by ROM start because `compress_entries_indices`
|
|
# refers to indices in ROM order, but the uncompressed dmadata might not be
|
|
# in ROM order.
|
|
dma_entries.sort(key=lambda dma_entry: dma_entry.vrom_start)
|
|
|
|
with multiprocessing.Pool(n_threads, initializer=set_sigint_ignored) as p:
|
|
# Extract each segment from the input rom
|
|
for entry_index, dma_entry in enumerate(dma_entries):
|
|
segment_rom_start = dma_entry.rom_start
|
|
segment_rom_end = dma_entry.rom_start + (
|
|
dma_entry.vrom_end - dma_entry.vrom_start
|
|
)
|
|
segment_data_uncompressed = rom_data[segment_rom_start:segment_rom_end]
|
|
|
|
is_compressed = entry_index in compress_entries_indices
|
|
|
|
if is_compressed:
|
|
segment_data = None
|
|
segment_data_async = p.apply_async(
|
|
crunch64.yaz0.compress,
|
|
(bytes(segment_data_uncompressed),),
|
|
)
|
|
else:
|
|
segment_data = segment_data_uncompressed
|
|
segment_data_async = None
|
|
|
|
compressed_rom_segments.append(
|
|
RomSegment(
|
|
dma_entry.vrom_start,
|
|
dma_entry.vrom_end,
|
|
is_compressed,
|
|
dma_entry.is_syms(),
|
|
segment_data,
|
|
segment_data_async,
|
|
)
|
|
)
|
|
|
|
# Wait on compression of all compressed segments
|
|
waiting_on_segments = [
|
|
segment for segment in compressed_rom_segments if segment.is_compressed
|
|
]
|
|
total_uncompressed_size_of_data_to_compress = sum(
|
|
segment.uncompressed_size for segment in waiting_on_segments
|
|
)
|
|
uncompressed_size_of_data_compressed_so_far = 0
|
|
while waiting_on_segments:
|
|
# Show progress
|
|
progress = (
|
|
uncompressed_size_of_data_compressed_so_far
|
|
/ total_uncompressed_size_of_data_to_compress
|
|
)
|
|
print(f"Compressing... {progress * 100:.1f}%", end="\r")
|
|
|
|
# The segments for which the compression is not finished yet are
|
|
# added to this list
|
|
still_waiting_on_segments = []
|
|
got_some_results = False
|
|
for segment in waiting_on_segments:
|
|
assert segment.data is None
|
|
assert segment.data_async is not None
|
|
|
|
try:
|
|
compressed_data = segment.data_async.get(0)
|
|
except multiprocessing.TimeoutError:
|
|
# Compression not finished yet
|
|
still_waiting_on_segments.append(segment)
|
|
else:
|
|
# Compression finished!
|
|
assert isinstance(compressed_data, bytes)
|
|
segment.data = memoryview(compressed_data)
|
|
uncompressed_size_of_data_compressed_so_far += (
|
|
segment.uncompressed_size
|
|
)
|
|
got_some_results = True
|
|
segment.data_async = None
|
|
|
|
if not got_some_results and still_waiting_on_segments:
|
|
# Nothing happened this wait iteration, idle a bit
|
|
time.sleep(0.010)
|
|
|
|
waiting_on_segments = still_waiting_on_segments
|
|
|
|
print("Putting together the compressed rom...")
|
|
|
|
# Put together the compressed rom
|
|
compressed_rom_size = sum(
|
|
align(len(segment.data)) for segment in compressed_rom_segments
|
|
)
|
|
pad_to_multiple_of = 8 * 2**20 # 8 MiB
|
|
compressed_rom_size_padded = (
|
|
(compressed_rom_size + pad_to_multiple_of - 1)
|
|
// pad_to_multiple_of
|
|
* pad_to_multiple_of
|
|
)
|
|
compressed_rom_data = memoryview(bytearray(compressed_rom_size_padded))
|
|
compressed_rom_dma_entries: list[dmadata.DmaEntry] = []
|
|
rom_offset = 0
|
|
for segment in compressed_rom_segments:
|
|
assert segment.data is not None
|
|
|
|
segment_rom_start = rom_offset
|
|
segment_rom_end = align(segment_rom_start + len(segment.data))
|
|
|
|
i = segment_rom_start + len(segment.data)
|
|
assert i <= len(compressed_rom_data)
|
|
compressed_rom_data[segment_rom_start:i] = segment.data
|
|
|
|
rom_offset = segment_rom_end
|
|
|
|
if segment.is_syms:
|
|
segment_rom_start = 0xFFFFFFFF
|
|
segment_rom_end = 0xFFFFFFFF
|
|
elif not segment.is_compressed:
|
|
segment_rom_end = 0
|
|
|
|
compressed_rom_dma_entries.append(
|
|
dmadata.DmaEntry(
|
|
segment.vrom_start,
|
|
segment.vrom_end,
|
|
segment_rom_start,
|
|
segment_rom_end,
|
|
)
|
|
)
|
|
|
|
assert rom_offset == compressed_rom_size
|
|
# Pad the compressed rom with the pattern matching the baseroms
|
|
for i in range(compressed_rom_size, compressed_rom_size_padded):
|
|
compressed_rom_data[i] = i % 256
|
|
|
|
# Write the new dmadata
|
|
offset = dmadata_start
|
|
for dma_entry in compressed_rom_dma_entries:
|
|
dma_entry.to_bin(compressed_rom_data[offset:])
|
|
offset += dmadata.DmaEntry.SIZE_BYTES
|
|
|
|
return compressed_rom_data
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--in",
|
|
dest="in_rom",
|
|
required=True,
|
|
help="path to an uncompressed rom to be compressed",
|
|
)
|
|
parser.add_argument(
|
|
"--out",
|
|
dest="out_rom",
|
|
required=True,
|
|
help="path of the compressed rom to write out",
|
|
)
|
|
parser.add_argument(
|
|
"--dmadata-start",
|
|
dest="dmadata_start",
|
|
type=lambda s: int(s, 16),
|
|
required=True,
|
|
help=(
|
|
"The dmadata location in the rom, as a hexadecimal offset (e.g. 0x12f70)."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--compress",
|
|
dest="compress_ranges",
|
|
required=True,
|
|
help=(
|
|
"The indices in the dmadata of the entries to be compressed,"
|
|
" where 0 is the first entry."
|
|
" It is a comma-separated list of individual indices and inclusive ranges."
|
|
" e.g. '0-1,3,5,6-9' is all indices from 0 to 9 (included) except 2 and 4."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--threads",
|
|
dest="n_threads",
|
|
type=int,
|
|
default=1,
|
|
help="how many cores to use for parallel compression",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
in_rom_p = Path(args.in_rom)
|
|
if not in_rom_p.exists():
|
|
parser.error(f"Input rom file {in_rom_p} doesn't exist.")
|
|
|
|
out_rom_p = Path(args.out_rom)
|
|
|
|
dmadata_start = args.dmadata_start
|
|
|
|
compress_ranges_str: str = args.compress_ranges
|
|
compress_entries_indices = set()
|
|
for compress_range_str in compress_ranges_str.split(","):
|
|
compress_range_ends_str = compress_range_str.split("-")
|
|
assert len(compress_range_ends_str) <= 2, (
|
|
compress_range_ends_str,
|
|
compress_range_str,
|
|
compress_ranges_str,
|
|
)
|
|
compress_range_ends = [int(v_str) for v_str in compress_range_ends_str]
|
|
if len(compress_range_ends) == 1:
|
|
compress_entries_indices.add(compress_range_ends[0])
|
|
else:
|
|
assert len(compress_range_ends) == 2
|
|
compress_range_first, compress_range_last = compress_range_ends
|
|
compress_entries_indices.update(
|
|
range(compress_range_first, compress_range_last + 1)
|
|
)
|
|
|
|
n_threads = args.n_threads
|
|
|
|
in_rom_data = in_rom_p.read_bytes()
|
|
out_rom_data = compress_rom(
|
|
memoryview(in_rom_data),
|
|
dmadata_start,
|
|
compress_entries_indices,
|
|
n_threads,
|
|
)
|
|
out_rom_p.write_bytes(out_rom_data)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|