Merge pull request #232 from steven11sjf/version-checking-2

Version checking

src/mercury_engine_data_structures/_dread_data_construct.py

@@ -49,4 +49,46 @@ def _build(self, obj: dict[str, int], stream, context, path):
         return self._build_construct._build(list(obj.items()), stream, context, path)
 
 
+class VersionedHashesDict(construct.Construct):
+    def __init__(self):
+        super().__init__()
+        self._build_construct = construct.PrefixedArray(
+            construct.Int32un,
+            construct.Sequence(
+                construct.PascalString(construct.Int16un, "ascii"),  # key
+                construct.Int64un,  # hash
+                construct.Int16un,  # versions
+            ),
+        )
+
+    def _parse(self, stream, context, path) -> dict[str, int]:
+        key_struct = struct.Struct("=H")
+        value_struct = struct.Struct("=QH")
+
+        count = construct.Int32un._parse(stream, None, "")
+
+        result = {}
+        for _ in range(count):
+            key = stream.read(key_struct.unpack(stream.read(2))[0]).decode()
+            value, versions = value_struct.unpack(stream.read(10))
+            result[key] = {"crc": value, "versions": versions}
+
+        return result
+
+    def _build(self, obj: dict[str, dict], stream, context, path):
+        ver_to_val = context.versions
+        all_vers = sum(ver_to_val.values())
+        for a in obj.values():
+            vers = a.get("versions")
+            if vers is not None:
+                a["versions"] = sum([ver_to_val[v] for v in vers])
+            else:
+                a["versions"] = all_vers
+
+        return self._build_construct._build(
+            [(k, v["crc"], v["versions"]) for k, v in obj.items()], stream, context, path
+        )
+
+
 KnownHashes = CompressedZSTD(HashesDict(), 15)
+VersionedHashes = CompressedZSTD(VersionedHashesDict(), 15)

src/mercury_engine_data_structures/dread_data.py

@@ -5,9 +5,12 @@
 import typing
 from pathlib import Path
 
-from mercury_engine_data_structures._dread_data_construct import KnownHashes
+from mercury_engine_data_structures._dread_data_construct import KnownHashes, VersionedHashes
+from mercury_engine_data_structures.game_check import Game, GameVersion
 
 _root = Path(__file__).parent
+DREAD_VERSIONS = GameVersion.versions_for_game(Game.DREAD)
+ALL_VERSIONS_BITMASK = -1
 
 
 @functools.lru_cache
@@ -18,23 +21,38 @@ def get_raw_types() -> dict[str, typing.Any]:
 
 
 @functools.lru_cache
-def all_name_to_asset_id() -> dict[str, int]:
+def all_name_to_asset_id_and_version() -> dict[str, dict[str, int]]:
     bin_path = _root.joinpath("dread_resource_names.bin")
     if bin_path.exists():
-        return dict(KnownHashes.parse_file(bin_path))
+        return dict(VersionedHashes.parse_file(bin_path))
 
     path = _root.joinpath("dread_resource_names.json")
     with path.open() as names_file:
-        return json.load(names_file)
+        data: dict[str, dict] = json.load(names_file)
+
+    for a in data.values():
+        vers = a.get("versions")
+        if vers is not None:
+            a["versions"] = sum([DREAD_VERSIONS[v].bitmask for v in vers])
+        else:
+            a["versions"] = ALL_VERSIONS_BITMASK
+
+    return data
+
+
+@functools.lru_cache
+def all_name_to_asset_id(ver: GameVersion | None = None) -> dict[str, int]:
+    bitmask = ver.bitmask if ver else ALL_VERSIONS_BITMASK
+    return {k: v["crc"] for k, v in all_name_to_asset_id_and_version().items() if v["versions"] & bitmask != 0}
 
 
 @functools.lru_cache
-def all_asset_id_to_name() -> dict[int, str]:
-    return {asset_id: name for name, asset_id in all_name_to_asset_id().items()}
+def all_asset_id_to_name(ver: GameVersion | None = None) -> dict[int, str]:
+    return {asset_id: name for name, asset_id in all_name_to_asset_id(ver).items()}
 
 
-def name_for_asset_id(asset_id: int) -> str | None:
-    return all_asset_id_to_name().get(asset_id)
+def name_for_asset_id(asset_id: int, ver: GameVersion | None = None) -> str | None:
+    return all_asset_id_to_name(ver).get(asset_id)
 
 
 @functools.lru_cache
@@ -55,11 +73,11 @@ def all_property_id_to_name() -> dict[int, str]:
     return {asset_id: name for name, asset_id in names.items()}
 
 
-def all_files_ending_with(ext: str, exclusions: list[str] | None = None) -> list[str]:
+def all_files_ending_with(ext: str, exclusions: list[str] = None, ver: GameVersion | None = None) -> list[str]:
     if not ext.startswith("."):
         ext = "." + ext
 
     if exclusions is None:
         exclusions = []
 
-    return [name for name in all_name_to_asset_id().keys() if name.endswith(ext) and name not in exclusions]
+    return [name for name in all_name_to_asset_id(ver).keys() if name.endswith(ext) and name not in exclusions]