Hooks :D

diff --git a/src/pacman/ecs/__init__.py b/src/pacman/ecs/__init__.py
index 522a3afe90a75562b643c8822bbb7ad60ad252dd..b01cb872263bacda64a15f7f03c15fac5e5b7976 100644 (file)
--- a/src/pacman/ecs/__init__.py
+++ b/src/pacman/ecs/__init__.py
@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -16,9 +17,13 @@ from pacman.ecs.schedule import (
     UpdateSchedule,
 )
 from pacman.ecs.world import (
+    CommandQueue,
     Component,
+    Entity,
     Not,
     Query,
+    RawComponentArg,
+    RawComponentHook,
     Resource,
     Systems,
     World,
@@ -178,11 +183,38 @@ def move_root(root: Query[Coord2D, WiggleRoot]) -> None:
         coord.x = 30.0 + math.cos(tick) * 30.0
 
 
+def sprite_despawn_hook() -> RawComponentHook[WiggleSprite]:
+    def inner(arg: RawComponentArg[WiggleSprite]) -> None:
+        for ent, _ in arg:
+            print("despawn", ent)
+
+    return inner
+
+
+def unspawn_wiggle(
+    queue: CommandQueue, q: Query[Entity, WiggleSprite]
+) -> None:
+    lst = [e for _, e in zip(range(3), q, strict=False)]
+    queue.queue(lambda world: world.despawn_many(lst))
+
+
+def close_on_empty(world: World, q: Query[Entity, WiggleSprite]) -> None:
+    if next(iter(q), None) is None:
+        world.spawn(WorldShouldExit, WorldShouldExit())
+
+
 if __name__ == "__main__":
     world = (
         World()
         .with_plugins(minimal_plugins, graphics_plugins)
-        .with_systems(UpdateSchedule, move_sprites, move_root)
+        .with_systems(
+            UpdateSchedule,
+            move_sprites,
+            move_root,
+            unspawn_wiggle,
+            close_on_empty,
+        )
+        .with_hook("remove", WiggleSprite, sprite_despawn_hook)
     )
     assets = list(Path("./asset").rglob("*.png"))
     root = world.spawn(

diff --git a/src/pacman/ecs/world.py b/src/pacman/ecs/world.py
index 1c787b4f6418795e9455a326091311821a81bc40..c308d59a967f2a8485919772e5e093fde1cd686c 100644 (file)
--- a/src/pacman/ecs/world.py
+++ b/src/pacman/ecs/world.py
@@ -144,13 +144,19 @@ class Archetype:
                 return kv[ty]
             return map(self._components[ty].__getitem__, positions)
 
+        itr = range(len(dst._entities) - len(components), len(dst._entities))
         dst._extend(cb)
-        self._discard_sort(positions)
+        self._discard(positions)
+
+        for k in comps_ty:
+            lst = dst._components[k]
+            hooks = self._world._components_cache[k][1]
+            hooks._insert((dst._entities[i], lst[i]) for i in itr)
 
     def remove_components[T: Components](
         self, ty: type[T], entities: Iterable["Entity"]
     ) -> None:
-        dst, _ = self._compute_remove(ty)
+        dst, comps_ty = self._compute_remove(ty)
         positions = [self._positions[e] for e in entities]
 
         def cb(ty: ComponentID | None) -> Iterable[Any]:
@@ -158,11 +164,31 @@ class Archetype:
                 return entities
             return map(self._components[ty].__getitem__, positions)
 
+        removed_ents = [self._entities[i] for i in positions]
+        removed_comps = [
+            [lst[i] for i in positions]
+            for lst in map(self._components.__getitem__, comps_ty)
+        ]
         dst._extend(cb)
-        self._discard_sort(positions)
+        self._discard(positions)
+        for k, comp in zip(comps_ty, removed_comps, strict=True):
+            hooks = self._world._components_cache[k][1]
+            hooks._remove(zip(removed_ents, comp, strict=True))
 
-    def remove(self, entities: Iterable["Entity"]) -> None:
-        self._discard_sort([self._positions[e] for e in entities])
+    def remove(self, entities: list["Entity"]) -> None:
+        positions = [self._positions[e] for e in entities]
+        comps = [
+            (
+                self._world._components_cache[comp][1],
+                [lst[i] for i in positions],
+            )
+            for comp, lst in self._components.items()
+        ]
+        self._discard(positions)
+        for ent in entities:
+            del self._world._entities[ent]
+        for hook, comp in comps:
+            hook._remove(zip(entities, comp))
 
     def _extend(
         self,
@@ -171,6 +197,7 @@ class Archetype:
         for entity in cb(None):
             self._positions[entity] = len(self._entities)
             self._entities.append(entity)
+            self._world._entities[entity] = self
         for ty, comp in self._components.items():
             itr = cb(ty)
             comp.extend(itr)
@@ -179,20 +206,20 @@ class Archetype:
             for comp in self._components.values()
         ), "Callaback's return lists weren't all of equal size"
 
-    def _discard_sort(self, itr: list[int]) -> None:
-        itr.sort(reverse=True)
+    def _discard(self, itr: Iterable[int]) -> None:
+        sort = sorted(itr, reverse=True)
 
         def remove_idx[T](
             lst: list[T], cb: Callable[[T, int], None] = lambda _1, _2: None
         ) -> None:
-            for i in itr:
+            for i in sort:
                 shifted = lst.pop()
                 if len(lst) == i:
                     continue
                 cb(shifted, i)
                 lst[i] = shifted
 
-        for i in itr:
+        for i in sort:
             self._positions.__delitem__(self._entities[i])
         remove_idx(self._entities, self._positions.__setitem__)
         for comp in self._components.values():
@@ -205,7 +232,8 @@ type System[T] = Callable[..., T]
 
 type Plugin = RecursiveTuple[Callable[["World"], None]]
 
-type RawComponentHook[T] = Callable[[list[tuple["Entity", T]]], None]
+type RawComponentArg[T] = Iterable[tuple["Entity", T]]
+type RawComponentHook[T] = Callable[[RawComponentArg], None]
 type ComponentHook[T] = System[RawComponentHook[T]]
 
 
@@ -530,7 +558,9 @@ class Query[Items, Filters = tuple[()]](
 
     @override
     def __contains__(self, entity: Entity) -> bool:
-        arch: Archetype = self._world._entities[entity]
+        arch: Archetype | None = self._world._entities.get(entity)
+        if arch is None:
+            return False
         if (arch_idx := self._archetypes_idx.get(arch)) is None:
             return False
         components_idx = arch._positions[entity]
@@ -578,6 +608,10 @@ class CommandQueue(FromWorld):
         self._queue.append(sys)
 
 
+def _system_failed_va[T, *P = *tuple[()]](*_: *P) -> T:
+    raise ValueError("Failed to fetch arguments for system")
+
+
 def _system_failed[T]() -> T:
     raise ValueError("Failed to fetch arguments for system")
 
@@ -608,7 +642,6 @@ class World(FromWorld):
 
     def component_id(self, comp: type) -> ComponentID:
         if (res := self._components_lookup.get(comp)) is None:
-            print(f"added comp {comp}")
             res = ComponentID(len(self._components_cache))
             self._components_cache.append((comp, ComponentHooks()))
             self._components_lookup[comp] = res
@@ -628,10 +661,12 @@ class World(FromWorld):
         )
         return res
 
-    def spawn[T](self, ty: type[T], comp: T) -> Entity:
+    def spawn[T: Components](self, ty: type[T], comp: T) -> Entity:
         return self.spawn_many(ty, [comp])[0]
 
-    def spawn_many[T](self, ty: type[T], comps: list[T]) -> list[Entity]:
+    def spawn_many[T: Components](
+        self, ty: type[T], comps: list[T]
+    ) -> list[Entity]:
         ents = self.spawn_many_empty(len(comps))
         self._base_archetype.add_components(
             cast(type[Components], ty),
@@ -639,11 +674,54 @@ class World(FromWorld):
         )
         return ents
 
-    def add_system_raw[T](
-        self,
-        system: Callable[..., T],
-        default: Callable[[], T] = _system_failed,
-    ) -> RawSystem[T]:
+    def despawn(self, ent: Entity) -> None:
+        self.despawn_many([ent])
+
+    def despawn_many(self, ents: Iterable[Entity]) -> None:
+        archs: dict[Archetype, list[Entity]] = {}
+        for ent in ents:
+            arch = self._entities[ent]
+            if arch not in archs:
+                archs[arch] = []
+            archs[arch].append(ent)
+        for arch, lst in archs.items():
+            arch.remove(lst)
+
+    def insert[T: Components](self, ty: type[T], ent: Entity, comp: T) -> None:
+        self.insert_many(ty, [(ent, comp)])
+
+    def insert_many[T: Components](
+        self, ty: type[T], comps: Iterable[tuple[Entity, T]]
+    ) -> None:
+        archs: dict[Archetype, list[tuple[Entity, T]]] = {}
+        for ent, comp in comps:
+            arch = self._entities[ent]
+            if arch not in archs:
+                archs[arch] = []
+            archs[arch].append((ent, comp))
+
+        for arch, lst in archs.items():
+            arch.add_components(ty, lst)
+
+    def remove[T: Components](self, ty: type[T], ent: Entity) -> None:
+        self.remove_many(ty, [ent])
+
+    def remove_many[T: Components](
+        self, ty: type[T], comps: Iterable[Entity]
+    ) -> None:
+        archs: dict[Archetype, list[Entity]] = {}
+        for ent in comps:
+            arch = self._entities[ent]
+            if arch not in archs:
+                archs[arch] = []
+            archs[arch].append(ent)
+
+        for arch, lst in archs.items():
+            arch.remove_components(ty, lst)
+
+    def _sys_extractors[T](
+        self, system: System[T]
+    ) -> tuple[Callable[[], FromWorld | None], ...]:
         sig = inspect.signature(system, eval_str=True)
         if err := next(
             filter(
@@ -655,13 +733,41 @@ class World(FromWorld):
             None,
         ):
             raise TypeError(f"System {system}'s {err[0]} parameter is invalid")
-        extractors = tuple(
+        return tuple(
             cast(type[FromWorld], param.annotation).extract_from_world(
                 param.annotation, self
             )
             for param in sig.parameters.values()
         )
 
+    def add_system[T, *P](
+        self,
+        system: System[Callable[[*P], T]],
+        default: Callable[[*P], T] = _system_failed_va,
+    ) -> Callable[[*P], T]:
+        extractors = self._sys_extractors(system)
+
+        def inner(*cb_args: *P) -> T:
+            args = []
+            for extract in extractors:
+                curr = extract()
+                if curr is None:
+                    return default()
+                args.append(curr)
+            res = system(*args)(*cb_args)
+            for arg in args:
+                arg.cleanup()
+            return res
+
+        return inner
+
+    def add_raw_system[T](
+        self,
+        system: Callable[..., T],
+        default: Callable[[], T] = _system_failed,
+    ) -> RawSystem[T]:
+        extractors = self._sys_extractors(system)
+
         def inner() -> T:
             args = []
             for extract in extractors:
@@ -718,7 +824,7 @@ class World(FromWorld):
     ) -> "World":
         comp_id = self.component_id(resolve_type_aliases(ty))
         hooks = self._components_cache[comp_id][1]
-        sys = self.add_system_raw(hook)
+        sys = self.add_raw_system(hook)
         match cycle:
             case "insert":
                 hooks.add_insert(lambda e: sys()(e))
@@ -735,7 +841,7 @@ class World(FromWorld):
         return self
 
     def run_main(self) -> None:
-        self.run(MainSchedule, self.add_system_raw(should_exit, lambda: False))
+        self.run(MainSchedule, self.add_raw_system(should_exit, lambda: False))
 
 
 _order_groups: int = 0
@@ -775,7 +881,7 @@ class Systems:
             else (
                 e
                 if isinstance(e, SystemSet)
-                else world.add_system_raw(e, default=lambda: None)
+                else world.add_raw_system(e, default=lambda: None)
             )
             for e in self._systems
         )