Archetypes, not finished

diff --git a/src/pacman/ecs/world.py b/src/pacman/ecs/world.py
index 974ebd15011d01cdae7c2c2da3e2957ecd378466..7cef6eb9e71e41d38937300668e42888e11b7fea 100644 (file)
--- a/src/pacman/ecs/world.py
+++ b/src/pacman/ecs/world.py
@@ -12,6 +12,7 @@ from typing import (
     get_args,
     get_origin,
     override,
+    runtime_checkable,
 )
 
 from pacman.ecs.schedule import (
@@ -31,17 +32,142 @@ class Resource2[T]:
     storage: type[T]
 
 
+type Components = RecursiveTuple["Component"]
+
+
+def components_iter(components: Components) -> Iterator["Component"]:
+    pass
+
+
+def components_type_iter(
+    components: type[Components],
+) -> Iterator[type["Component"]]:
+    pass
+
+
 class Archetype:
-    def __init__(self, types: Iterable[type["Component"]]) -> None:
+    def __init__(
+        self, world: "World", types: Iterable[type["Component"]]
+    ) -> None:
+        self._world = world
         self._types = frozenset(types)
         self._positions: dict[Entity, int] = {}
         self._entities: list[Entity] = []
-        self._components: dict[type[Component], list[Any]] = {}
-        self._add_cache: dict[int, tuple[Any, Archetype, tuple[int, ...]]] = {}
+        self._components: dict[type[Component], list[Any]] = {
+            ty: [] for ty in types
+        }
+        self._add_cache: dict[
+            int, tuple[Any, Archetype, list[type[Component]]]
+        ] = {}
         self._remove_cache: dict[
-            int, tuple[Any, Archetype, tuple[int, ...]]
+            int, tuple[Any, Archetype, list[type[Component]]]
         ] = {}
 
+        self._world._archetypes.append(self)
+        self._world._archetypes_lookup[self._types] = self
+
+    def _compute_add(
+        self, ty: type[Components]
+    ) -> tuple["Archetype", list[type["Component"]]]:
+        ptr = id(ty)
+        if res := self._add_cache.get(ptr):
+            return (res[1], res[2])
+        comps = list(components_type_iter(ty))
+        res_arch: Archetype
+        types = frozenset(comps) | self._types
+        if prev := self._world._archetypes_lookup.get(types):
+            res_arch = prev
+        else:
+            res_arch = Archetype(self._world, types)
+        self._add_cache[ptr] = (ty, res_arch, comps)
+        return (res_arch, comps)
+
+    def _compute_remove(
+        self, ty: type[Components]
+    ) -> tuple["Archetype", list[type["Component"]]]:
+        ptr = id(ty)
+        if res := self._remove_cache.get(ptr):
+            return (res[1], res[2])
+        comps = list(components_type_iter(ty))
+        res_arch: Archetype
+        types = self._types - frozenset(comps)
+        if prev := self._world._archetypes_lookup.get(types):
+            res_arch = prev
+        else:
+            res_arch = Archetype(self._world, types)
+        self._remove_cache[ptr] = (ty, res_arch, comps)
+        return (res_arch, comps)
+
+    def add_components[T: Components](
+        self, ty: type[T], components: Iterable[tuple["Entity", T]]
+    ) -> None:
+        dst, comps_ty = self._compute_add(ty)
+        comps = [components_iter(e) for _, e in components]
+        positions = [self._positions[e] for e, _ in components]
+        kv = {}
+        for k in comps_ty:
+            kv[k] = [next(e) for e in comps]
+
+        def cb(ty: type[Entity] | type[Component]) -> Iterable[Any]:
+            if issubclass(ty, Entity):
+                return (e for e, _ in components)
+            if ty in kv:
+                return kv[ty]
+            return map(self._components[ty].__getitem__, positions)
+
+        dst._extend(cb)
+        self._discard_sort(positions)
+
+    def remove_components[T: Components](
+        self, ty: type[T], entities: Iterable["Entity"]
+    ) -> None:
+        dst, _ = self._compute_remove(ty)
+        positions = [self._positions[e] for e in entities]
+
+        def cb(ty: type[Entity] | type[Component]) -> Iterable[Any]:
+            if issubclass(ty, Entity):
+                return entities
+            return map(self._components[ty].__getitem__, positions)
+
+        dst._extend(cb)
+        self._discard_sort(positions)
+
+    def remove(self, entities: Iterable["Entity"]) -> None:
+        self._discard_sort([self._positions[e] for e in entities])
+
+    def _extend(
+        self,
+        cb: Callable[[type["Component"] | type["Entity"]], Iterable[Any]],
+    ) -> None:
+        for entity in cb(Entity):
+            self._positions[entity] = len(self._entities)
+            self._entities.append(entity)
+        for ty, comp in self._components.items():
+            comp.extend(cb(ty))
+        assert all(
+            len(comp) == len(self._entities)
+            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 remove_idx[T](
+            lst: list[T], cb: Callable[[T, int], None] = lambda _1, _2: None
+        ) -> None:
+            for i in itr:
+                shifted = lst.pop()
+                if len(lst) == i:
+                    continue
+                cb(shifted, i)
+                lst[i] = shifted
+
+        for i in itr:
+            self._positions.__delitem__(self._entities[i])
+        remove_idx(self._entities, self._positions.__setitem__)
+        for comp in self._components.values():
+            remove_idx(comp)
+
 
 type RawSystem[T] = Callable[[], T]
 
@@ -50,6 +176,7 @@ type Plugin = RecursiveTuple[Callable[["World"], None]]
 type ComponentHook[T] = Callable[["World", "Entity", T], None]
 
 
+@runtime_checkable
 class FromWorld(Protocol):
     @classmethod
     @abstractmethod
@@ -57,10 +184,8 @@ class FromWorld(Protocol):
         cls, world: "World"
     ) -> Callable[[], Self | None]: ...
 
-
-type SystemParam = RecursiveTuple[FromWorld]
-
-type System[T] = Callable[..., T]
+    def cleanup(self) -> None:
+        pass
 
 
 class QueryFilter:
@@ -300,27 +425,41 @@ class Resource(FromWorld, Component):
         return lambda: query().single()
 
 
-class World(FromWorld):
-    _BASE_ARCHETYPE: int = 0
+def _system_failed[T]() -> T:
+    raise ValueError("Failed to fetch arguments for system")
+
 
+class World(FromWorld):
     def __init__(self) -> None:
         self._entities: dict[Entity, Archetype] = {}
-        self._archetypes: list[Archetype] = [Archetype([])]
         self._archetypes_lookup: dict[
             frozenset[type[Component]], Archetype
-        ] = {frozenset(): self._archetypes[self._BASE_ARCHETYPE]}
-
+        ] = {}
+        self._archetypes: list[Archetype] = []
         self._schedules: dict[ScheduleLabel, Schedule[RawSystem]] = {}
         self._insert_hooks: dict[type, ComponentHook[Any]] = {}
         self._remove_hooks: dict[type, ComponentHook[Any]] = {}
+        self._base_archetype: Archetype = Archetype(self, ())
 
     def _run_cond(self, _: RawSystem | SystemSet) -> bool:
         return True
 
     def add_system_raw[T](
-        self, system: System[T], default: Callable[[], T]
+        self,
+        system: Callable[..., T],
+        default: Callable[[], T] = _system_failed,
     ) -> RawSystem[T]:
-        sig = inspect.signature(system)
+        sig = inspect.signature(system, eval_str=True)
+        if err := next(
+            filter(
+                lambda kv: (
+                    not issubclass(get_origin(kv[1].annotation), FromWorld)
+                ),
+                sig.parameters.items(),
+            ),
+            None,
+        ):
+            raise TypeError(f"System {system}'s {err[0]} parameter is invalid")
         extractors = tuple(
             cast(type[FromWorld], param.annotation).extract_from_world(self)
             for param in sig.parameters.values()
@@ -333,7 +472,10 @@ class World(FromWorld):
                 if curr is None:
                     return default()
                 args.append(curr)
-            return system(*args)
+            res = system(*args)
+            for arg in args:
+                arg.cleanup()
+            return res
 
         return inner