Factored out hierarchy traversal

diff --git a/src/pacman/ecs/__init__.py b/src/pacman/ecs/__init__.py
index 369aa8865a37223eaa8187505ddb1ca8c82d4f08..71428a42671275076264d0063dcfff2acc851d8a 100644 (file)
--- a/src/pacman/ecs/__init__.py
+++ b/src/pacman/ecs/__init__.py
@@ -2,12 +2,17 @@ from collections.abc import Callable, Generator, Iterable
 from dataclasses import dataclass
 from pathlib import Path
 from time import time
-from typing import Any, Literal, cast, get_args, overload
+from typing import Any, Literal, cast, overload
 
-import pygame.sprite
+import pygame
 from pygame import Vector2
 
-from pacman.ecs.hierarchy import Children, Parent, hierachy_plugins
+from pacman.ecs.hierarchy import (
+    Children,
+    Parent,
+    hierachy_plugins,
+    traverse_hierarchy,
+)
 from pacman.ecs.schedule import (
     MainSchedule,
     Schedule,
@@ -66,20 +71,10 @@ def update_coord(world: World) -> None:
 
 
 def update_coord_hierarchy(world: World) -> None:
-    work = []
-    for (children,) in world.query(tuple[Children], without=tuple[Parent]):
-        work.extend(children.children)
-    while work:
-        curr = work.pop()
-        child = world[curr]
-        parent = world[world[curr][Parent].parent]
-        if CoordAbs2D in parent and CoordAbs2D in child:
-            parent_coord = parent[CoordAbs2D]
-            child_coord = child[CoordAbs2D]
-            child_coord.vec += parent_coord.vec
-            child_coord.depth += parent_coord.depth
-        if Children in child:
-            work.extend(child[Children].children)
+    for parent, child in traverse_hierarchy(world, tuple[CoordAbs2D]):
+        if parent is None:
+            continue
+        child[0].vec += parent[0].vec
 
 
 PropagateTransform = SystemSet("propagate-transform")
@@ -156,9 +151,8 @@ def render_sprite2d(world: World) -> None:
 def graphics_plugins(world: World) -> None:
     world.with_systems(StartupSchedule, init_window).with_systems(
         UpdateSchedule,
-        Systems(poll_events, close_on_close, render_sprite2d)
-        .chain()
-        .after(PropagateTransform),
+        Systems(render_sprite2d).in_set(PostUpdate).after(PropagateTransform),
+        Systems(poll_events, close_on_close).chain().in_set(PreUpdate),
     )
 
 

diff --git a/src/pacman/ecs/hierarchy.py b/src/pacman/ecs/hierarchy.py
index 9106a35a8a3c482e9918fc4fbbf6c5cd15aa3ac8..d2f5d9369467bc88c57031bb2f6a51aa00c28b06 100644 (file)
--- a/src/pacman/ecs/hierarchy.py
+++ b/src/pacman/ecs/hierarchy.py
@@ -1,4 +1,5 @@
 from collections.abc import Generator, Iterable
+from typing import cast, get_args
 
 from pacman.ecs.world import Entity, World
 
@@ -92,3 +93,36 @@ def hierachy_plugins(world: World) -> None:
     ).with_hook("insert", Parent, parent_insert_hook).with_hook(
         "remove", Parent, parent_remove_hook
     )
+
+
+def traverse_hierarchy[*T](
+    world: World, query: type[tuple[*T]]
+) -> Generator[tuple[tuple[*T] | None, tuple[*T]]]:
+    params = sorted(
+        get_args(query), key=lambda e: len(world._components.get(e, {}))
+    )
+    work = []
+    for e in world.query(query, without=tuple[Parent, Children]):
+        yield (None, e)
+    for entity, children in world.query(
+        tuple[Entity, Children], without=tuple[Parent]
+    ):
+        work.extend(children.children)
+        ent = world[entity]
+        if all(param in ent for param in params):
+            yield (None, cast(tuple[*T], (*(ent[param] for param in params),)))
+    while work:
+        curr = work.pop()
+        child = world[curr]
+        parent = world[world[curr][Parent].parent]
+        if all(param in child for param in params):
+            parent_res = None
+            if all(param in parent for param in params):
+                parent_res = cast(
+                    tuple[*T], (*(parent[param] for param in params),)
+                )
+            child_res = cast(tuple[*T], (*(child[param] for param in params),))
+            yield (parent_res, child_res)
+
+        if Children in child:
+            work.extend(child[Children].children)