Extract ROM segments from gc-eu-mq (take 2) (#1709)

* Extract ROM segments from gc-eu-mq (take 2)

* Apply suggestions from code review

Co-authored-by: Dragorn421 <Dragorn421@users.noreply.github.com>

* Typo

* dma_{start,names} -> dmadata_{start,names}

* Restore dest

* Don't assume rom location == vrom location

---------

Co-authored-by: Dragorn421 <Dragorn421@users.noreply.github.com>

tools/compress.py

@@ -46,13 +46,13 @@ def set_sigint_ignored():
 
 def compress_rom(
     rom_data: memoryview,
-    dmadata_offset: int,
+    dmadata_start: int,
     compress_entries_indices: set[int],
     n_threads: int = None,
 ):
     """
     rom_data: the uncompressed rom data
-    dmadata_offset: the offset in the rom where the dmadata starts
+    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
     """
@@ -60,7 +60,7 @@ def compress_rom(
     # Segments of the compressed rom (not all are compressed)
     compressed_rom_segments: list[RomSegment] = []
 
-    dma_entries = dmadata.read_dmadata(rom_data, dmadata_offset)
+    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.
@@ -184,7 +184,7 @@ def compress_rom(
         compressed_rom_data[i] = i % 256
 
     # Write the new dmadata
-    offset = dmadata_offset
+    offset = dmadata_start
     for dma_entry in compressed_rom_dma_entries:
         dma_entry.to_bin(compressed_rom_data[offset:])
         offset += dmadata.DmaEntry.SIZE_BYTES
@@ -207,8 +207,8 @@ def main():
         help="path of the compressed rom to write out",
     )
     parser.add_argument(
-        "--dma-start",
-        dest="dma_start",
+        "--dmadata-start",
+        dest="dmadata_start",
         type=lambda s: int(s, 16),
         required=True,
         help=(
@@ -241,7 +241,7 @@ def main():
 
     out_rom_p = Path(args.out_rom)
 
-    dmadata_offset = args.dma_start
+    dmadata_start = args.dmadata_start
 
     compress_ranges_str: str = args.compress_ranges
     compress_entries_indices = set()
@@ -267,7 +267,7 @@ def main():
     in_rom_data = in_rom_p.read_bytes()
     out_rom_data = compress_rom(
         memoryview(in_rom_data),
-        dmadata_offset,
+        dmadata_start,
         compress_entries_indices,
         n_threads,
     )

tools/decompress_baserom.py

@@ -34,17 +34,6 @@ def decompress(data: bytes, is_zlib_compressed: bool) -> bytes:
     return crunch64.yaz0.decompress(data)
 
 
-FILE_TABLE_OFFSET = {
-    "gc-eu-mq": 0x07170,
-    "gc-eu-mq-dbg": 0x12F70,
-}
-
-VERSIONS_MD5S = {
-    "gc-eu-mq": "1a438f4235f8038856971c14a798122a",
-    "gc-eu-mq-dbg": "f0b7f35375f9cc8ca1b2d59d78e35405",
-}
-
-
 def round_up(n, shift):
     mod = 1 << shift
     return (n + mod - 1) >> shift << shift
@@ -64,7 +53,7 @@ def update_crc(decompressed: io.BytesIO) -> io.BytesIO:
 
 def decompress_rom(
     file_content: bytearray,
-    dmadata_offset: int,
+    dmadata_start: int,
     dma_entries: list[dmadata.DmaEntry],
     is_zlib_compressed: bool,
 ) -> bytearray:
@@ -93,7 +82,7 @@ def decompress_rom(
         decompressed.seek(vrom_st)
         decompressed.write(data)
     # write new dmadata
-    decompressed.seek(dmadata_offset)
+    decompressed.seek(dmadata_start)
     for dma_entry in new_dmadata:
         entry_data = bytearray(dmadata.DmaEntry.SIZE_BYTES)
         dma_entry.to_bin(entry_data)
@@ -177,17 +166,21 @@ def main():
     parser.add_argument(
         "version",
         help="Version of the game to decompress.",
-        choices=list(VERSIONS_MD5S.keys()),
     )
 
     args = parser.parse_args()
 
     version = args.version
 
-    uncompressed_path = Path(f"baseroms/{version}/baserom-decompressed.z64")
+    baserom_dir = Path(f"baseroms/{version}")
+    if not baserom_dir.exists():
+        print(f"Error: Unknown version '{version}'.")
+        exit(1)
 
-    dmadata_offset = FILE_TABLE_OFFSET[version]
-    correct_str_hash = VERSIONS_MD5S[version]
+    uncompressed_path = baserom_dir / "baserom-decompressed.z64"
+
+    dmadata_start = int((baserom_dir / "dmadata_start.txt").read_text(), 16)
+    correct_str_hash = (baserom_dir / "checksum.md5").read_text().split()[0]
 
     if check_existing_rom(uncompressed_path, correct_str_hash):
         print("Found valid baserom - exiting early")
@@ -225,13 +218,13 @@ def main():
 
     file_content = per_version_fixes(file_content, version)
 
-    dma_entries = dmadata.read_dmadata(file_content, dmadata_offset)
+    dma_entries = dmadata.read_dmadata(file_content, dmadata_start)
     # Decompress
     if any(dma_entry.is_compressed() for dma_entry in dma_entries):
         print("Decompressing rom...")
         is_zlib_compressed = version in {"ique-cn", "ique-zh"}
         file_content = decompress_rom(
-            file_content, dmadata_offset, dma_entries, is_zlib_compressed
+            file_content, dmadata_start, dma_entries, is_zlib_compressed
         )
 
     file_content = pad_rom(file_content, dma_entries)

tools/extract_baserom.py

@@ -0,0 +1,76 @@
+#!/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 sys
+
+import dmadata
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Extract segments from an uncompressed ROM, based on its dmadata."
+    )
+    parser.add_argument(
+        "rom", metavar="ROM", type=Path, help="Path to uncompressed ROM"
+    )
+    parser.add_argument(
+        "-o",
+        "--output-dir",
+        type=Path,
+        required=True,
+        help="Output directory for segments",
+    )
+    parser.add_argument(
+        "--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(
+        "--dmadata-names",
+        type=Path,
+        required=True,
+        help="Path to file containing segment names",
+    )
+
+    args = parser.parse_args()
+
+    rom_data = memoryview(args.rom.read_bytes())
+
+    dma_names = args.dmadata_names.read_text().splitlines()
+    dma_entries = dmadata.read_dmadata(rom_data, args.dmadata_start)
+    if len(dma_names) != len(dma_entries):
+        print(
+            f"Error: expected {len(dma_names)} DMA entries but found {len(dma_entries)} in ROM",
+            file=sys.stderr,
+        )
+        exit(1)
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    for dma_name, dma_entry in zip(dma_names, dma_entries):
+        if dma_entry.is_compressed():
+            print(f"Error: segment {dma_name} is compressed", file=sys.stderr)
+            exit(1)
+
+        segment_rom_start = dma_entry.rom_start
+        segment_rom_end = dma_entry.rom_start + (
+            dma_entry.vrom_end - dma_entry.vrom_start
+        )
+
+        segment_data = rom_data[segment_rom_start:segment_rom_end]
+        segment_path = args.output_dir / dma_name
+        segment_path.write_bytes(segment_data)
+
+    print(f"Extracted {len(dma_entries)} segments to {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()