diff --git a/src/retro_data_structures/construct_extensions/dict.py b/src/retro_data_structures/construct_extensions/dict.py index 8a57158..22fcd4d 100644 --- a/src/retro_data_structures/construct_extensions/dict.py +++ b/src/retro_data_structures/construct_extensions/dict.py @@ -1,38 +1,72 @@ from __future__ import annotations -from construct import Adapter, PrefixedArray, Struct, VarInt +import construct from retro_data_structures.common_types import String -class DictAdapter(Adapter): - def __init__(self, subcon, objisdict=True): - if not objisdict: - subcon = Struct("*Key" / VarInt, "Value" / subcon) - super().__init__(PrefixedArray(VarInt, subcon)) - self.objisdict = objisdict - - def _decode(self, obj, context, path): - D = {} - for v in obj: - if self.objisdict: - D[v["*Key"]] = v - del v["*Key"] - else: - D[v["*Key"]] = v["Value"] - return D - - def _encode(self, obj, context, path): - L = [] - for k, v in obj.items(): - new_item = v - if self.objisdict: - v["*Key"] = k - else: - new_item = {"*Key": k, "Value": v} - L.append(new_item) - return L - - -def DictStruct(*fields): - return Struct(*fields, "*Key" / String) +class DictConstruct(construct.Construct): + def __init__( + self, + key_type: construct.Construct, + value_type: construct.Construct, + count_type: construct.Construct = construct.Int32ub, + ): + super().__init__() + self.key_type = key_type + self.value_type = value_type + self.count_type = count_type + + assert not self.key_type.flagbuildnone + + def _parse(self, stream, context, path) -> construct.Container: + field_count = self.count_type._parsereport(stream, context, path) + + result = construct.Container() + + for i in range(field_count): + field_path = f"{path}.field_{i}" + key = self.key_type._parsereport(stream, context, field_path) + value = self.value_type._parsereport(stream, context, field_path) + result[key] = value + + return result + + def _build(self, obj: construct.Container, stream, context, path): + self.count_type._build(len(obj), stream, context, path) + + for i, (key, value) in enumerate(obj.items()): + field_path = f"{path}.field_{i}" + self.key_type._build(key, stream, context, field_path) + self.value_type._build(value, stream, context, field_path) + + def _emitparse(self, code): + return "Container(({}, {}) for i in range({}))".format( + self.key_type._compileparse(code), + self.value_type._compileparse(code), + self.count_type._compileparse(code), + ) + + def _emitbuild(self, code): + fname = f"build_dict_{code.allocateId()}" + block = f""" + def {fname}(key, value, io, this): + obj = key + {self.key_type._compilebuild(code)} + + obj = value + {self.value_type._compilebuild(code)} + """ + code.append(block) + return ( + f"(reuse(len(obj), " + f"lambda obj: {self.count_type._compilebuild(code)}), " + f"list({fname}(key, value, io, this) for key, value in obj.items()), obj)[2]" + ) + + +def make_dict(value: construct.Construct, key=String): + return DictConstruct( + key_type=key, + value_type=value, + ) diff --git a/src/retro_data_structures/formats/pak.py b/src/retro_data_structures/formats/pak.py index f771501..dc05329 100644 --- a/src/retro_data_structures/formats/pak.py +++ b/src/retro_data_structures/formats/pak.py @@ -3,7 +3,7 @@ import typing from retro_data_structures.base_resource import AssetId, RawResource -from retro_data_structures.formats import pak_gc, pak_wiiu +from retro_data_structures.formats import pak_gc, pak_wii, pak_wiiu from retro_data_structures.formats.pak_gc import PakBody, PakFile from retro_data_structures.game_check import Game @@ -11,6 +11,8 @@ def _pak_for_game(game: Game): if game == Game.PRIME_REMASTER: return pak_wiiu.PAK_WIIU + elif game >= Game.CORRUPTION: + raise ValueError("Unsupported game") else: return pak_gc.PAK_GC @@ -27,6 +29,8 @@ def __init__(self, raw: PakBody, target_game: Game): def header_for_game(game: Game): if game == Game.PRIME_REMASTER: return pak_wiiu.PakWiiUNoData + elif game >= Game.CORRUPTION: + return pak_wii.PAKNoData else: return pak_gc.PAKNoData diff --git a/src/retro_data_structures/formats/pak_wii.py b/src/retro_data_structures/formats/pak_wii.py new file mode 100644 index 0000000..5796efd --- /dev/null +++ b/src/retro_data_structures/formats/pak_wii.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import construct +from construct import Bytes, Const, Int32ub, PrefixedArray, Struct + +from retro_data_structures.common_types import AssetId64, FourCC, String +from retro_data_structures.construct_extensions.dict import make_dict + +PAKHeader = construct.Aligned( + 64, + Struct( + version=Const(2, Int32ub), + header_size=Int32ub, + md5_hash=Bytes(16), + ), +) + +ConstructResourceHeader = Struct( + compressed=Int32ub, + asset_type=FourCC, + asset_id=AssetId64, + size=Int32ub, + offset=Int32ub, +) + + +def _emitparse_header(code: construct.CodeGen) -> str: + code.append("ResourceHeader_Format = struct.Struct('>LLQLL')") + code.append( + """def _create_resource_header(compressed, asset_type, asset_id, size, offset): + return Container(compressed=compressed, asset_type=asset_type.to_bytes(4, "big").decode("ascii"), + asset_id=asset_id, size=size, offset=offset) + """ + ) + return "_create_resource_header(*ResourceHeader_Format.unpack(io.read(24)))" + + +ConstructResourceHeader._emitparse = _emitparse_header + +PAKNoData = Struct( + _start=construct.Tell, + _header=PAKHeader, + table_of_contents=construct.Aligned(64, make_dict(Int32ub, FourCC)), + _named_resources_start=construct.Tell, + named_resources=construct.Aligned( + 64, + PrefixedArray( + Int32ub, + Struct( + name=String, + asset_type=FourCC, + asset_id=AssetId64, + ), + ), + ), + _resources_start=construct.Tell, + _resources_start_assert=construct.Check( + construct.this.table_of_contents.STRG == construct.this._resources_start - construct.this._named_resources_start + ), + resources=construct.Aligned(64, PrefixedArray(Int32ub, ConstructResourceHeader)), + _resources_end=construct.Tell, + _resources_end_assert=construct.Check( + construct.this.table_of_contents.RSHD == construct.this._resources_end - construct.this._resources_start + ), +)