1
0
Fork 0
mirror of https://github.com/zeldaret/oot.git synced 2025-05-10 19:13:42 +00:00
oot/tools/assets/extract/extract_xml_z64.py

383 lines
13 KiB
Python

import dataclasses
from pathlib import Path
from pprint import pprint
from ..descriptor.base import (
BaseromFileBackingMemory,
get_resources_desc,
ResourceDesc,
ResourcesDescCollection,
ResourcesDescCollectionsPool,
SegmentStartAddress,
VRAMStartAddress,
)
from .extase import File, Resource
from .extase.memorymap import MemoryContext
from . import z64_resource_handlers
#
# main
#
VERBOSE1 = False
VERBOSE2 = False
# "options"
RM_SOURCE = True
WRITE_SOURCE = True
RM_EXTRACT = True
WRITE_EXTRACT = True
from ..conf import WRITE_HINTS, I_D_OMEGALUL
@dataclasses.dataclass
class ExtractionContext:
oot_version: str
version_memctx_base: MemoryContext
baserom_path: Path
build_path: Path
extracted_path: Path
def get_baserom_file_data(self, baserom_file_name: str):
return memoryview((self.baserom_path / baserom_file_name).read_bytes())
def create_file_resources(rescoll: ResourcesDescCollection, file: File):
"""Collect resources described by the collection into the file object.
Does not do any processing of the data itself: no parsing, no discovering more resources.
"""
file_resources_by_desc: dict[ResourceDesc, Resource] = dict()
list_ResourceNeedsPostProcessWithPoolResourcesException: list[
z64_resource_handlers.ResourceNeedsPostProcessWithPoolResourcesException
] = []
for resource_desc in rescoll.resources:
try:
resource = z64_resource_handlers.get_resource_from_desc(file, resource_desc)
except (
z64_resource_handlers.ResourceNeedsPostProcessWithPoolResourcesException
) as e:
resource = e.resource
list_ResourceNeedsPostProcessWithPoolResourcesException.append(e)
# TODO nice hack right here.
# probably instead rework the "c declaration" system into a more opaque object
# not that this is really a required long term feature as it's only relevant
# for writing the source files (main .c/.h), not extracting
if file.name.startswith("ovl_") and file.name != "ovl_file_choose":
resource.HACK_IS_STATIC_ON = ...
file.add_resource(resource)
file_resources_by_desc[resource_desc] = resource
if VERBOSE1:
print(file)
print(file.name, file._resources)
# Check if described resources overlap
file.sort_resources()
file.check_overlapping_resources()
return (
file_resources_by_desc,
list_ResourceNeedsPostProcessWithPoolResourcesException,
)
def process_pool(
extraction_ctx: ExtractionContext, pool_desc: ResourcesDescCollectionsPool
):
if VERBOSE2:
print("> process_pool")
if len(pool_desc.collections) == 1:
print(", ".join(map(str, (_c.out_path for _c in pool_desc.collections))))
file_by_rescoll: dict[ResourcesDescCollection, File] = dict()
# 1) Create File and Resource objects
pool_resources_by_desc: dict[ResourceDesc, Resource] = dict()
list_all_ResourceNeedsPostProcessWithPoolResourcesException: list[
z64_resource_handlers.ResourceNeedsPostProcessWithPoolResourcesException
] = []
for rescoll in pool_desc.collections:
if not isinstance(rescoll.backing_memory, BaseromFileBackingMemory):
raise NotImplementedError(rescoll.backing_memory)
data = extraction_ctx.get_baserom_file_data(rescoll.backing_memory.name)
if rescoll.backing_memory.range is not None:
range_start, range_end = rescoll.backing_memory.range
data = data[range_start:range_end]
# TODO File.name
file = File(rescoll.backing_memory.name, data=data)
file_by_rescoll[rescoll] = file
(
file_resources_by_desc,
list_ResourceNeedsPostProcessWithPoolResourcesException,
) = create_file_resources(rescoll, file)
pool_resources_by_desc |= file_resources_by_desc
list_all_ResourceNeedsPostProcessWithPoolResourcesException.extend(
list_ResourceNeedsPostProcessWithPoolResourcesException
)
for e in list_all_ResourceNeedsPostProcessWithPoolResourcesException:
e.callback(pool_resources_by_desc)
# 2) Build a MemoryContext for each File
memctx_base = extraction_ctx.version_memctx_base.copy()
files_by_segment: dict[int, list[File]] = dict()
for rescoll, file in file_by_rescoll.items():
if rescoll.start_address is None:
pass
elif isinstance(rescoll.start_address, SegmentStartAddress):
files_by_segment.setdefault(rescoll.start_address.segment, []).append(file)
elif isinstance(rescoll.start_address, VRAMStartAddress):
memctx_base.set_direct_file(rescoll.start_address.vram, file)
else:
raise NotImplementedError(rescoll.start_address)
disputed_segments = []
for segment, files in files_by_segment.items():
if len(files) == 1:
memctx_base.set_segment_file(segment, files[0])
else:
disputed_segments.append(segment)
if VERBOSE2:
print(f"{disputed_segments=}")
memctx_by_file: dict[File, MemoryContext] = dict()
for rescoll, file in file_by_rescoll.items():
file_memctx = memctx_base.copy()
if (
isinstance(rescoll.start_address, SegmentStartAddress)
and rescoll.start_address.segment in disputed_segments
):
if VERBOSE2:
print(
"disputed segment",
rescoll.start_address.segment,
"set to own file for file's memctx",
file.name,
)
file_memctx.set_segment_file(rescoll.start_address.segment, file)
memctx_by_file[file] = file_memctx
# 3) parse: iteratively discover and parse data
# (discover = add resources, parse = make friendlier than binary)
def parse_all_files():
while True:
any_progress = False
for file, file_memctx in memctx_by_file.items():
if VERBOSE2:
print(file.name, ".try_parse_resources_data()")
if file.try_parse_resources_data(file_memctx):
any_progress = True
if not any_progress:
break
for file in memctx_by_file.keys():
file.check_non_parsed_resources()
if VERBOSE2:
print("parse_all_files() ...")
parse_all_files()
for file in memctx_by_file.keys():
file.commit_resource_buffers()
if VERBOSE2:
print("parse new resources (resource buffers) ...")
parse_all_files() # parse new resources (resource buffers)
for file in memctx_by_file.keys():
file.sort_resources()
file.check_overlapping_resources()
# 4) add dummy (binary) resources for the unaccounted gaps
if VERBOSE2:
print("unaccounted...")
for file in memctx_by_file.keys():
file.add_unaccounted_resources(I_D_OMEGALUL=I_D_OMEGALUL)
parse_all_files() # FIXME this is to set is_data_parsed=True on binary blob unaccounteds, handle better
for file in memctx_by_file.keys():
file.sort_resources()
assert not file.get_overlapping_resources()
# At this point all files are completely mapped with resources
# 5)
for rescoll, file in file_by_rescoll.items():
file.set_source_path(
extraction_ctx.extracted_path / "assets" / rescoll.out_path
)
file.set_resources_paths(
extraction_ctx.extracted_path,
extraction_ctx.build_path,
Path("assets") / rescoll.out_path,
)
for file, file_memctx in memctx_by_file.items():
# write to extracted/
if WRITE_EXTRACT:
file.write_resources_extracted(file_memctx)
# "source" refers to the main .c and .h `#include`ing all the extracted resources
if WRITE_SOURCE:
# TODO fill referenced_files properly or something
file.referenced_files = set(memctx_by_file.keys()) - {file}
file.write_source()
def process_pool_wrapped(extraction_ctx, pd):
try:
process_pool(extraction_ctx, pd)
except Exception as e:
import traceback
import sys
# Some exceptions can't be pickled for passing back to the main process
# so print them now as well as reraising
traceback.print_exc(file=sys.stdout)
raise Exception(
"ERROR with pool_desc collections:",
[str(rescoll.out_path) for rescoll in pd.collections],
) from e
def main():
import argparse
import re
from tools import version_config
from tools import dmadata
parser = argparse.ArgumentParser()
parser.add_argument(
"baserom_segments_dir",
type=Path,
help="Directory of uncompressed ROM segments",
)
parser.add_argument(
"output_dir",
type=Path,
help="Output directory to place files in",
)
parser.add_argument("-v", dest="oot_version", default="gc-eu-mq-dbg")
parser.add_argument("-j", dest="use_multiprocessing", action="store_true")
parser.add_argument("-s", dest="single", default=None)
parser.add_argument("-r", dest="single_is_regex", default=None, action="store_true")
args = parser.parse_args()
vc = version_config.load_version_config(args.oot_version)
dma_entries = dmadata.read_dmadata(
memoryview((args.baserom_segments_dir / "dmadata").read_bytes()), 0
)
dmadata_table_rom_file_name_by_vrom = dict()
for dma_entry, name in zip(dma_entries, vc.dmadata_segments.keys()):
dmadata_table_rom_file_name_by_vrom[
(dma_entry.vrom_start, dma_entry.vrom_end)
] = name
pools_desc = get_resources_desc(vc)
version_memctx_base = MemoryContext(dmadata_table_rom_file_name_by_vrom)
from .extase_oot64.dlist_resources import MtxResource, TextureResource
from ..n64 import G_IM_FMT, G_IM_SIZ
file_gMtxClear = File("sys_matrix__gMtxClear", size=0x40)
file_gMtxClear.add_resource(MtxResource(file_gMtxClear, 0, "gMtxClear"))
version_memctx_base.set_direct_file(vc.variables["gMtxClear"], file_gMtxClear)
file_sShadowTex = File("z_en_jsjutan__sShadowTex", size=0x800)
file_sShadowTex.add_resource(
TextureResource(
file_sShadowTex, 0, "sShadowTex", G_IM_FMT.I, G_IM_SIZ._8b, 32, 64
)
)
version_memctx_base.set_direct_file(vc.variables["sShadowTex"], file_sShadowTex)
extraction_ctx = ExtractionContext(
args.oot_version,
version_memctx_base,
args.baserom_segments_dir,
Path("build") / args.oot_version,
args.output_dir,
)
z64_resource_handlers.register_resource_handlers()
# TODO extract only when a pool xml was modified since last extract
try:
if args.single is not None:
any_do_process_pool = False
for pool_desc in pools_desc:
do_process_pool = False
for coll in pool_desc.collections:
if isinstance(coll.backing_memory, BaseromFileBackingMemory):
if args.single_is_regex:
if re.fullmatch(args.single, coll.backing_memory.name):
do_process_pool = True
else:
if coll.backing_memory.name == args.single:
do_process_pool = True
if do_process_pool:
process_pool(extraction_ctx, pool_desc)
any_do_process_pool = True
if any_do_process_pool:
print("Done!")
else:
print("Not found:", args.single)
elif not args.use_multiprocessing: # everything on one process
for pool_desc in pools_desc:
process_pool(extraction_ctx, pool_desc)
print("All done!")
else: # multiprocessing
import multiprocessing
with multiprocessing.Pool() as pool:
pool.starmap(
process_pool_wrapped,
zip([extraction_ctx] * len(pools_desc), pools_desc),
)
print("All done!")
except Exception as e:
import traceback
import sys
traceback.print_exc(file=sys.stdout)
try:
import rich.pretty
except:
print("Install rich for prettier output (`pip install rich`)")
else:
# TODO implement more __rich_repr__
if e.__class__ in (Exception, AssertionError, NotImplementedError):
print("rich.pretty.pprint(e.args):")
rich.pretty.pprint(e.args, indent_guides=False)
else:
print("rich.pretty.pprint(e):")
rich.pretty.pprint(e, indent_guides=False)