Skip to content

More refactoring of MREA for Corruption #218

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

More refactoring of MREA for Corruption #218

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
033254e
Adjusting MREASimple to support Corruption.
henriquegemignani Sep 22, 2025
be0d1bb
Move when compressed block headers are read
henriquegemignani Sep 22, 2025
ea2bd96
Move data section sizes to header, rebuild count
henriquegemignani Sep 22, 2025
987a2b6
Moved the compressed block headers to MREAHeader
henriquegemignani Sep 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 93 additions & 86 deletions src/retro_data_structures/formats/mrea.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,8 @@
from functools import cached_property

import construct
from construct import Adapter, Aligned, If, Int32ub, PrefixedArray, Struct
from construct import Adapter, Aligned, If, Int32ub, PrefixedArray, Rebuild, Struct, this
from construct.core import (
Array,
Computed,
Const,
Enum,
Expand All @@ -29,7 +28,7 @@
from retro_data_structures.base_resource import AssetId, AssetType, BaseResource, Dependency
from retro_data_structures.common_types import AssetId32, FourCC, String, Transform4f
from retro_data_structures.compression import LZOCompressedBlock
from retro_data_structures.construct_extensions.alignment import PrefixedWithPaddingBefore
from retro_data_structures.construct_extensions.alignment import AlignTo, PrefixedWithPaddingBefore
from retro_data_structures.construct_extensions.version import BeforeVersion, WithVersion
from retro_data_structures.data_section import DataSection
from retro_data_structures.exceptions import DependenciesHandledElsewhere, UnknownAssetId
Expand All @@ -44,7 +43,7 @@
from retro_data_structures.game_check import AssetIdCorrect, Game

if typing.TYPE_CHECKING:
from collections.abc import Iterable, Iterator
from collections.abc import Iterable, Iterator, Sequence

from typing_extensions import Self

Expand Down Expand Up @@ -171,56 +170,67 @@ def _encode(self, obj: list[list[str]], context, path):
),
}

MREAHeader = Aligned(
32,
Struct(
"magic" / Const(0xDEADBEEF, Int32ub),
"version" / Enum(Int32ub, MREAVersion),
# Matrix that represents the area's transform from the origin.
# Most area data is pre-transformed, so this matrix is only used occasionally.
"area_transform" / Transform4f,
# Number of world models in this area.
"world_model_count" / Int32ub,
# Number of script layers in this area.
"script_layer_count" / WithVersion(MREAVersion.Echoes, Int32ub),
# Number of data sections in the file.
"data_section_count" / Int32ub,
# Section index for world geometry data. Always 0; starts on materials.
"geometry_section" / Int32ub,
# Section index for script layer data.
"script_layers_section" / Int32ub,
# Section index for generated script object data.
"generated_script_objects_section" / WithVersion(MREAVersion.Echoes, Int32ub),
# Section index for collision data.
"collision_section" / Int32ub,
# Section index for first unknown section.
"unknown_section_1" / Int32ub,
# Section index for light data.
"lights_section" / Int32ub,
# Section index for visibility tree data.
"visibility_tree_section" / Int32ub,
# Section index for path data.
"path_section" / Int32ub,
# Section index for area octree data.
"area_octree_section" / BeforeVersion(MREAVersion.EchoesDemo, Int32ub),
# Section index for second unknown section.
"unknown_section_2" / WithVersion(MREAVersion.Echoes, Int32ub),
# Section index for portal area data.
"portal_area_section" / WithVersion(MREAVersion.Echoes, Int32ub),
# Section index for static geometry map data.
"static_geometry_map_section" / WithVersion(MREAVersion.Echoes, Int32ub),
# Number of compressed data blocks in the file.
"compressed_block_count" / WithVersion(MREAVersion.Echoes, Int32ub),
),
)

CompressedBlockHeader = Struct(
buffer_size=Int32ub,
uncompressed_size=Int32ub,
compressed_size=Int32ub,
data_section_count=Int32ub,
)

MREAHeader = Struct(
"magic" / Const(0xDEADBEEF, Int32ub),
"version" / Enum(Int32ub, MREAVersion),
# Matrix that represents the area's transform from the origin.
# Most area data is pre-transformed, so this matrix is only used occasionally.
"area_transform" / Transform4f,
# Number of world models in this area.
"world_model_count" / Int32ub,
# Number of script layers in this area.
"script_layer_count" / WithVersion(MREAVersion.Echoes, Int32ub),
# Number of data sections in the file.
"_data_section_count" / Rebuild(Int32ub, construct.len_(this.data_section_sizes)),
# Index table for sections.
"section_index"
/ BeforeVersion(
MREAVersion.Corruption,
Struct(
# Section index for world geometry data. Always 0; starts on materials.
"geometry_section" / Int32ub,
# Section index for script layer data.
"script_layers_section" / Int32ub,
# Section index for generated script object data.
"generated_script_objects_section" / WithVersion(MREAVersion.Echoes, Int32ub),
# Section index for collision data.
"collision_section" / Int32ub,
# Section index for first unknown section.
"unknown_section_1" / Int32ub,
# Section index for light data.
"lights_section" / Int32ub,
# Section index for visibility tree data.
"visibility_tree_section" / Int32ub,
# Section index for path data.
"path_section" / Int32ub,
# Section index for area octree data.
"area_octree_section" / BeforeVersion(MREAVersion.EchoesDemo, Int32ub),
# Section index for second unknown section.
"unknown_section_2" / WithVersion(MREAVersion.Echoes, Int32ub),
# Section index for portal area data.
"portal_area_section" / WithVersion(MREAVersion.Echoes, Int32ub),
# Section index for static geometry map data.
"static_geometry_map_section" / WithVersion(MREAVersion.Echoes, Int32ub),
),
),
# Number of compressed data blocks in the file.
"compressed_block_count"
/ WithVersion(MREAVersion.Echoes, Rebuild(Int32ub, construct.len_(this._compressed_block_headers))),
# Number of section numbers at the end of the header.
"section_number_count" / WithVersion(MREAVersion.Corruption, Int32ub),
AlignTo(32),
"data_section_sizes" / Aligned(32, Int32ub[this._data_section_count]),
"_compressed_block_headers"
/ WithVersion(MREAVersion.Echoes, Aligned(32, CompressedBlockHeader[this.compressed_block_count])),
)


class MREACompressedBlock(construct.Construct):
"""
Expand Down Expand Up @@ -254,17 +264,20 @@ def _parse(self, stream, context, path) -> bytes:
return result


SectionIndex = Struct(
"type" / FourCC,
"index" / Int32ub,
)

# This construct decodes the minimum necessary to read the whole file and decompress everything.
MREAPrime2Simple = construct.Struct(
"header" / construct.Aligned(32, MREAHeader),
"data_section_sizes" / construct.Aligned(32, construct.Int32ub[construct.this.header.data_section_count]),
"_compressed_block_headers"
/ construct.Aligned(32, CompressedBlockHeader[construct.this.header.compressed_block_count]),
# Not compatible with Prime 1
MREASimple = Struct(
"header" / Aligned(32, MREAHeader),
"version" / Computed(this.header.version),
"section_index" / WithVersion(MREAVersion.Corruption, Aligned(32, SectionIndex[this.header.section_number_count])),
"compressed_blocks"
/ construct.Aligned(
32,
MREACompressedBlock(construct.this._compressed_block_headers),
)[construct.this.header.compressed_block_count],
/ Aligned(32, MREACompressedBlock(this.header._compressed_block_headers))[this.header.compressed_block_count],
WithVersion(MREAVersion.Corruption, AlignTo(64, b"\xff")),
construct.Terminated,
)

Expand Down Expand Up @@ -308,11 +321,13 @@ class MREAConstruct(construct.Construct):
def _aligned_parse(self, conn: construct.Construct, stream, context, path):
return Aligned(32, conn)._parsereport(stream, context, path)

def _decode_compressed_blocks(self, mrea_header, data_section_sizes, stream, context, path) -> list[bytes]:
compressed_block_headers = self._aligned_parse(
Array(mrea_header.compressed_block_count, CompressedBlockHeader), stream, context, path
)

def _decode_compressed_blocks(
self, compressed_block_headers: Sequence[Container], data_section_sizes: Sequence[int], stream, context, path
) -> list[bytes]:
"""
Reads and decodes a list of compressed blocks, described by the given headers,
and splits it into a list of data sections.
"""
context.compressed_block_headers = compressed_block_headers
compressed_block_construct = Aligned(32, MREACompressedBlock(construct.this.compressed_block_headers))

Expand All @@ -333,19 +348,21 @@ def _decode_compressed_blocks(self, mrea_header, data_section_sizes, stream, con

def _parse(self, stream, context, path):
mrea_header = MREAHeader._parsereport(stream, context, path)
data_section_sizes = self._aligned_parse(Array(mrea_header.data_section_count, Int32ub), stream, context, path)

if mrea_header.compressed_block_count is not None:
data_sections = self._decode_compressed_blocks(mrea_header, data_section_sizes, stream, context, path)
if mrea_header._compressed_block_headers is not None:
data_sections = self._decode_compressed_blocks(
mrea_header._compressed_block_headers, mrea_header.data_section_sizes, stream, context, path
)
else:
data_sections = Array(
mrea_header.data_section_count,
Aligned(32, FixedSized(lambda ctx: data_section_sizes[ctx._index], GreedyBytes)),
)._parsereport(stream, context, path)
data_sections = []
for section_size in mrea_header.data_section_sizes:
data_sections.append(self._aligned_parse(FixedSized(section_size, GreedyBytes), stream, context, path))

# Split data sections into the named sections
categories = [
{"label": label, "value": mrea_header[label]} for label in _all_categories if mrea_header[label] is not None
{"label": label, "value": mrea_header.section_index[label]}
for label in _all_categories
if mrea_header.section_index[label] is not None
]
categories.sort(key=lambda c: c["value"])

Expand Down Expand Up @@ -462,18 +479,18 @@ def _build(self, obj: Container, stream, context, path):

# Combine all sections into the data sections array
data_sections = ListContainer()
mrea_header.section_index = Container()

for category in _all_categories:
if category in raw_sections:
mrea_header[category] = len(data_sections)
mrea_header.section_index[category] = len(data_sections)
data_sections.extend(raw_sections[category])
else:
mrea_header[category] = None
mrea_header.section_index[category] = None

# Compress the data sections
if int(obj.version) >= MREAVersion.Echoes.value:
compressed_blocks = self._encode_compressed_blocks(data_sections, mrea_header, context, path)
mrea_header.compressed_block_count = len(compressed_blocks)
compressed_blocks = self._encode_compressed_blocks(data_sections, mrea_header.section_index, context, path)
else:
compressed_blocks = None
raise NotImplementedError
Expand All @@ -482,22 +499,12 @@ def _build(self, obj: Container, stream, context, path):
mrea_header.area_transform = obj.area_transform
mrea_header.world_model_count = obj.world_model_count
mrea_header.script_layer_count = len(raw_sections.script_layers_section)
mrea_header.data_section_count = len(data_sections)
mrea_header.section_number_count = len(mrea_header.section_index)
mrea_header.data_section_sizes = [len(section) for section in data_sections]
mrea_header._compressed_block_headers = [block.header for block in compressed_blocks]

MREAHeader._build(mrea_header, stream, context, path)
Aligned(32, Array(mrea_header.data_section_count, Int32ub))._build(
[len(section) for section in data_sections],
stream,
context,
path,
)
if compressed_blocks is not None:
Aligned(32, Array(mrea_header.compressed_block_count, CompressedBlockHeader))._build(
[block.header for block in compressed_blocks],
stream,
context,
path,
)
for compressed_block in compressed_blocks:
block_header = compressed_block.header
if block_header.compressed_size:
Expand Down
9 changes: 8 additions & 1 deletion tests/formats/test_mrea.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ def test_compare_p1(prime1_asset_manager):


def test_compare_p2(prime2_asset_manager, mrea_asset_id: AssetId):
compare_all_instances(prime2_asset_manager, mrea_asset_id, _all_instances_p1_p2, mrea.MREAPrime2Simple)
compare_all_instances(prime2_asset_manager, mrea_asset_id, _all_instances_p1_p2, mrea.MREASimple)


# @pytest.mark.skip(reason="Corruption MREA not implemented correctly")
Expand Down Expand Up @@ -149,3 +149,10 @@ def test_compare_p2_add_layer_hashes(prime2_asset_manager, mrea_asset_id: AssetI
mrea_encoded = area.mrea.build()

_compare_mrea_hashes("mrea_hashes_echoes_add_layer.json", mlvl_encoded + mrea_encoded, mrea_asset_id)


def test_p3_header(prime3_asset_manager, mrea_asset_id: AssetId):
raw = prime3_asset_manager.get_raw_asset(mrea_asset_id)
header = mrea.MREASimple.parse(raw.data)
assert header is not None
# decoded = Mrea.parse(raw.data, prime3_asset_manager.target_game, prime3_asset_manager)