Updates to collision.py

src/mercury_engine_data_structures/formats/collision.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import copy
+
 import construct
 from construct import Array, Container, Flag, Hex, Int8ul, Int16ul, Rebuild, Struct, Switch
 
@@ -10,36 +12,56 @@
     Float,
     StrId,
     UInt,
+    Vec2,
+    Vec3,
+    Vec4,
     make_vector,
 )
 from mercury_engine_data_structures.construct_extensions.misc import ErrorWithMessage, OptionalValue
-from mercury_engine_data_structures.game_check import Game
+
+
+def calculate_poly_boundings(ctx: Container) -> Container:
+    x1 = min(point.position.x for point in ctx.points)
+    y1 = min(point.position.y for point in ctx.points)
+    x2 = max(point.position.x for point in ctx.points)
+    y2 = max(point.position.y for point in ctx.points)
+    return Container({"min": Vec2(x1, y1), "max": Vec2(x2, y2)})
+
+
+def calculate_collection_boundings(ctx: Container) -> Container:
+    x1 = min(poly.boundings.min.x for poly in ctx.polys)
+    y1 = min(poly.boundings.min.y for poly in ctx.polys)
+    x2 = max(poly.boundings.max.x for poly in ctx.polys)
+    y2 = max(poly.boundings.max.y for poly in ctx.polys)
+    return Container({"min": Vec2(x1, y1), "max": Vec2(x2, y2)})
+
+
+BoundingBox2D = Struct("min" / CVector2D, "max" / CVector2D)
 
 CollisionPoint = Struct(
-    "x" / Float,
-    "y" / Float,
+    "position" / CVector2D,
     "material_attribute" / UInt,
 )
 CollisionPolySR = Struct(
     "num_points" / Rebuild(UInt, construct.len_(construct.this.points)),
     "unk4" / Hex(construct.Byte),
     "unk5" / Hex(UInt),
     "points" / Array(construct.this.num_points, CollisionPoint),
-    "boundings" / Array(4, Float),
+    "boundings" / Rebuild(BoundingBox2D, calculate_poly_boundings),
 )
 CollisionPolyDread = Struct(
     "num_points" / Rebuild(UInt, construct.len_(construct.this.points)),
     "unk" / Hex(UInt),
     "points" / Array(construct.this.num_points, CollisionPoint),
     "loop" / Flag,
-    "boundings" / Array(4, Float),
+    "boundings" / Rebuild(BoundingBox2D, calculate_poly_boundings),
 )
-CollisionPoly = game_check.is_at_most(Game.SAMUS_RETURNS, CollisionPolySR, CollisionPolyDread)
+CollisionPoly = game_check.is_sr_or_else(CollisionPolySR, CollisionPolyDread)
 
 BinarySearchTree = Struct(
     "binary_search_index1" / Int16ul,
     "binary_search_index2" / Int16ul,
-    "boundings" / Array(4, Float),
+    "boundings" / BoundingBox2D,
 )
 
 collision_formats = {
@@ -66,7 +88,7 @@
     "POLYCOLLECTION2D": Struct(
         "position" / CVector3D,
         "polys" / make_vector(CollisionPoly),
-        "total_boundings" / Array(4, Float),
+        "total_boundings" / Rebuild(BoundingBox2D, calculate_collection_boundings),
         "binary_search_trees" / OptionalValue(make_vector(BinarySearchTree)),
     ),
 }
@@ -91,26 +113,111 @@
 )
 
 
+class Bounds2D:
+    def __init__(self, raw: Container):
+        self._raw = raw
+
+    @property
+    def min(self) -> Vec2:
+        return self._raw.min
+
+    @min.setter
+    def min(self, value: Vec2) -> None:
+        self._raw.min = value
+
+    @property
+    def max(self) -> Vec2:
+        return self._raw.max
+
+    @max.setter
+    def max(self, value: Vec2) -> None:
+        self._raw.max = value
+
+
+class PolyData:
+    def __init__(self, raw: Container):
+        self._raw = raw
+
+    @property
+    def num_points(self) -> int:
+        return len(self._raw.points)
+
+    @property
+    def boundings(self) -> Vec4:
+        return self._raw.boundings
+
+    def get_boundings(self) -> Bounds2D:
+        return Bounds2D(calculate_poly_boundings(self._raw))
+
+    def get_point(self, point_idx: int) -> PointData:
+        return PointData(self._raw.points[point_idx])
+
+    def add_point(self, position: Vec2, idx: int = 0) -> None:
+        """
+        Adds a new point by copying an existing point and inserting it at a specified index
+        param position: the x,y position of the new point
+        param idx: the index the new point will be placed in the poly
+        """
+        new_point = copy.deepcopy(self.get_point(0))
+        new_point.position = position
+        self._raw.points.insert(idx, new_point)
+
+    def remove_point(self, idx: int) -> None:
+        """
+        Removes a point from a poly by index
+        param idx: the index of the point to remove
+        """
+        self._raw.points.pop(idx)
+
+
+class PointData:
+    def __init__(self, raw: Container):
+        self._raw = raw
+
+    @property
+    def position(self) -> Vec2:
+        return self._raw.position
+
+    @position.setter
+    def position(self, value: Vec2) -> None:
+        self._raw.position = value
+
+    @property
+    def material_attribute(self) -> int:
+        return self._raw.material_attribute
+
+    @material_attribute.setter
+    def material_attribute(self, value: int = 1) -> None:
+        self._raw.material_attribute = value
+
+
 class CollisionEntry:
     def __init__(self, raw: Container):
         self._raw = raw
 
-    def get_data(self) -> Container:
-        """Returns all data of collision/collision_camera/logic shape"""
+    @property
+    def data(self) -> dict[Vec3, list[dict], Vec4, list[dict]]:
         return self._raw.data
 
-    def get_poly(self, poly_idx: int):
-        """Returns all data associated with a poly (points, boundings)"""
-        return self.get_data().polys[poly_idx]
+    @data.setter
+    def data(self, value: dict[Vec3, list[dict], Vec4, list[dict]]) -> None:
+        self._raw.data = value
+
+    @property
+    def position(self) -> Vec3:
+        return self.data.position
+
+    @position.setter
+    def position(self, value: Vec3) -> None:
+        self.data.position = value
 
-    def get_point(self, poly_idx: int, point_idx: int) -> Container:
-        """Returns a specific point in a poly"""
-        return self.get_poly(poly_idx).points[point_idx]
+    @property
+    def total_boundings(self) -> Vec4:
+        return self.data.total_boundings
 
-    def get_total_boundings(self) -> Container:
-        """Returns the total boundary of collision/collision_camera/logic shape"""
-        return self.get_data().total_boundings
+    def get_total_boundings(self) -> Bounds2D:
+        return Bounds2D(calculate_collection_boundings(self._raw))
 
-    def get_poly_boundings(self, poly_idx: int) -> Container:
-        """Returns the boundary of a poly"""
-        return self.get_poly(poly_idx).boundings
+    def get_poly(self, poly_idx: int) -> PolyData:
+        """Returns all data associated with a poly"""
+        return PolyData(self._raw.data.polys[poly_idx])

tests/formats/test_collision.py

@@ -4,6 +4,7 @@
 from tests.test_lib import parse_build_compare_editor_parsed
 
 from mercury_engine_data_structures import dread_data, samus_returns_data
+from mercury_engine_data_structures.common_types import Vec2, Vec3, Vec4
 from mercury_engine_data_structures.formats.bmscc import Bmscc
 
 bossrush_assets = [
@@ -100,24 +101,74 @@ def surface_bmscc(samus_returns_tree) -> Bmscc:
 
 
 def test_get_data(surface_bmscc: Bmscc):
-    data = surface_bmscc.get_entry().get_data()
-    assert len(data) == 5
+    entry = surface_bmscc.get_entry()
+    assert len(entry.data) == 5
+    assert entry.position == [0.0, 0.0, 0.0]
+
+    assert entry.data != {
+        "position": Vec3(0.0, 0.0, 0.0),
+        "polys": [
+            {
+                "num_points": 4,
+                "unk4": 0x0A,
+                "unk5": 0x01B18000,
+                "points": [
+                    {"position": Vec2(5900.0, -5400.0), "material_attribute": 1},
+                    {"position": Vec2(6400.0, -5400.0), "material_attribute": 1},
+                    {"position": Vec2(6400.0, -5500.0), "material_attribute": 1},
+                    {"position": Vec2(6100.0, -5500.0), "material_attribute": 1},
+                ],
+                "boundings": Vec4(5900.0, -5600.0, 6400.0, -5400.0),
+            }
+        ],
+        "total_boundings": Vec4(-25100.0, -10600.0, 12500.0, 14103.099609375),
+        "binary_search_trees": [
+            {
+                "binary_search_index1": 2149,
+                "binary_search_index2": 0,
+                "boundings": Vec4(-25100.0, -10600.0, 12500.0, 14103.099609375),
+            }
+        ],
+    }
 
 
 def test_modifying_collision(surface_bmscc: Bmscc):
-    point = surface_bmscc.get_entry().get_point(2, 9)
-    assert point["x"] == -800.0
-    assert point["y"] == -7000.0
-
-
-def test_get_boundings(surface_bmscc: Bmscc):
-    total_boundings = surface_bmscc.get_entry().get_total_boundings()
-    polys = surface_bmscc.get_entry().get_data().polys
-    for i, poly in enumerate(polys):
-        poly_boundings = surface_bmscc.get_entry().get_poly_boundings(i)
-        # Boundings for polygons are in the order: x1, y1, x2, y2
-        # Assert that the boundings are confined within the total bounds of the collision_camera
-        assert poly_boundings[0] >= total_boundings[0]
-        assert poly_boundings[1] >= total_boundings[1]
-        assert poly_boundings[2] <= total_boundings[2]
-        assert poly_boundings[3] <= total_boundings[3]
+    poly = surface_bmscc.get_entry().get_poly(2)
+    point = poly.get_point(5)
+    assert point.position == [-900.0, -7400.0]
+
+
+def test_boundings(surface_bmscc: Bmscc):
+    poly_boundings = surface_bmscc.get_entry().get_poly(0).boundings
+    assert poly_boundings == {"min": Vec2(-25100.0, -10600.0), "max": Vec2(12500.0, 14103.099609375)}
+    poly_boundings["min"] = Vec2(30000.0, 30000.0)
+    poly_boundings["max"] = Vec2(-30000.0, -30000.0)
+    assert poly_boundings == {"min": Vec2(30000.0, 30000.0), "max": Vec2(-30000.0, -30000.0)}
+
+    total_boundings = surface_bmscc.get_entry().total_boundings
+    assert total_boundings == {"min": Vec2(-25100.0, -10600.0), "max": Vec2(12500.0, 14103.099609375)}
+    total_boundings["min"] = Vec2(30000.0, 30000.0)
+    total_boundings["max"] = Vec2(-30000.0, -30000.0)
+    assert total_boundings == {"min": Vec2(30000.0, 30000.0), "max": Vec2(-30000.0, -30000.0)}
+
+
+def test_get_poly(surface_bmscc: Bmscc):
+    poly = surface_bmscc.get_entry().get_poly(5)
+    assert poly.num_points == 12
+
+
+def test_add_point(surface_bmscc: Bmscc):
+    point = (6000.0, -5400.0)
+    poly = surface_bmscc.get_entry().get_poly(3)
+    poly.add_point(point)
+    assert poly.get_point(0) is not None
+    assert poly.get_point(0).position == (6000.0, -5400.0)
+    assert poly.get_point(0).material_attribute == 1
+    assert poly.num_points == 9
+
+
+def test_remove_point(surface_bmscc: Bmscc):
+    poly = surface_bmscc.get_entry().get_poly(0)
+    assert poly.get_point(0).position == Vec2(4300.0, -4800.0)
+    poly.remove_point(0)
+    assert poly.get_point(0).position == Vec2(3600.0, -4800.0)