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>

Makefile

@@ -65,6 +65,7 @@ endif
 PROJECT_DIR := $(dir $(realpath $(firstword $(MAKEFILE_LIST))))
 BUILD_DIR := build/$(VERSION)
 EXPECTED_DIR := expected/$(BUILD_DIR)
+BASEROM_DIR := baseroms/$(VERSION)
 VENV := .venv
 
 MAKE = make
@@ -210,8 +211,7 @@ ASSET_FILES_OUT := $(foreach f,$(ASSET_FILES_XML:.xml=.c),$f) \
 
 UNDECOMPILED_DATA_DIRS := $(shell find data -type d)
 
-# TODO: for now, ROM segments are still taken from the Debug ROM even when building other versions
+BASEROM_SEGMENTS_DIR := $(BASEROM_DIR)/segments
-BASEROM_SEGMENTS_DIR := baseroms/gc-eu-mq-dbg/segments
 BASEROM_BIN_FILES := $(wildcard $(BASEROM_SEGMENTS_DIR)/*)
 
 # source files
@@ -223,7 +223,7 @@ O_FILES       := $(foreach f,$(S_FILES:.s=.o),$(BUILD_DIR)/$f) \
 
 OVL_RELOC_FILES := $(shell $(CPP) $(CPPFLAGS) $(SPEC) | $(SPEC_REPLACE_VARS) | grep -o '[^"]*_reloc.o' )
 
-DISASM_BASEROM := baseroms/$(VERSION)/baserom-decompressed.z64
+DISASM_BASEROM := $(BASEROM_DIR)/baserom-decompressed.z64
 DISASM_DATA_FILES := $(wildcard $(DISASM_DATA_DIR)/*.csv) $(wildcard $(DISASM_DATA_DIR)/*.txt)
 DISASM_S_FILES := $(shell test -e $(PYTHON) && $(PYTHON) tools/disasm/list_generated_files.py -o $(EXPECTED_DIR) --config-dir $(DISASM_DATA_DIR))
 DISASM_O_FILES := $(DISASM_S_FILES:.s=.o)
@@ -334,13 +334,13 @@ all: rom compress
 rom: $(ROM)
 ifneq ($(COMPARE),0)
 	@md5sum $(ROM)
-	@md5sum -c baseroms/$(VERSION)/checksum.md5
+	@md5sum -c $(BASEROM_DIR)/checksum.md5
 endif
 
 compress: $(ROMC)
 ifneq ($(COMPARE),0)
 	@md5sum $(ROMC)
-	@md5sum -c baseroms/$(VERSION)/checksum-compressed.md5
+	@md5sum -c $(BASEROM_DIR)/checksum-compressed.md5
 endif
 
 clean:
@@ -366,9 +366,9 @@ venv:
 setup: venv
 	$(MAKE) -C tools
 	$(PYTHON) tools/decompress_baserom.py $(VERSION)
-# TODO: for now, we only extract ROM segments and assets from the Debug ROM
+	$(PYTHON) tools/extract_baserom.py $(BASEROM_DIR)/baserom-decompressed.z64 -o $(BASEROM_SEGMENTS_DIR) --dmadata-start `cat $(BASEROM_DIR)/dmadata_start.txt` --dmadata-names $(BASEROM_DIR)/dmadata_names.txt
+# TODO: for now, we only extract assets from the Debug ROM
 ifeq ($(VERSION),gc-eu-mq-dbg)
-	$(PYTHON) extract_baserom.py
 	$(PYTHON) extract_assets.py -j$(N_THREADS)
 endif
 
@@ -390,7 +390,7 @@ $(ROM): $(ELF)
 	$(ELF2ROM) -cic 6105 $< $@
 
 $(ROMC): $(ROM) $(ELF) $(BUILD_DIR)/compress_ranges.txt
-	$(PYTHON) tools/compress.py --in $(ROM) --out $@ --dma-start `./tools/dmadata_start.sh $(NM) $(ELF)` --compress `cat $(BUILD_DIR)/compress_ranges.txt` --threads $(N_THREADS)
+	$(PYTHON) tools/compress.py --in $(ROM) --out $@ --dmadata-start `./tools/dmadata_start.sh $(NM) $(ELF)` --compress `cat $(BUILD_DIR)/compress_ranges.txt` --threads $(N_THREADS)
 	$(PYTHON) -m ipl3checksum sum --cic 6105 --update $@
 
 $(ELF): $(TEXTURE_FILES_OUT) $(ASSET_FILES_OUT) $(O_FILES) $(OVL_RELOC_FILES) $(LDSCRIPT) $(BUILD_DIR)/undefined_syms.txt

baseroms/gc-eu-mq-dbg/dmadata_start.txt

@@ -0,0 +1 @@
+0x12f70

baseroms/gc-eu-mq/dmadata_start.txt

@@ -0,0 +1 @@
+0x07170

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",
+        "--dmadata-start",
-        dest="dma_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]
+    uncompressed_path = baserom_dir / "baserom-decompressed.z64"
-    correct_str_hash = VERSIONS_MD5S[version]
+
+    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()