From 680d368e6c40d75d8385f4c363d44390a1894435 Mon Sep 17 00:00:00 2001 From: Axy Date: Mon, 21 Sep 2026 17:02:24 +0200 Subject: [PATCH] Big work on archetype based querying --- src/pacman/ecs/world.py | 470 ++++++++++++++++----------- src/pacman/utils/type_resolve.py | 36 +- src/pacman/utils/variadics_please.py | 219 +++++++++++-- 3 files changed, 504 insertions(+), 221 deletions(-) diff --git a/src/pacman/ecs/world.py b/src/pacman/ecs/world.py index b92df32..974ebd1 100644 --- a/src/pacman/ecs/world.py +++ b/src/pacman/ecs/world.py @@ -1,15 +1,16 @@ -import contextlib +import inspect from abc import abstractmethod -from collections.abc import Callable, Generator, Iterable, Iterator +from collections.abc import Callable, Container, Iterable, Iterator from dataclasses import dataclass from typing import ( Any, Literal, Protocol, Self, + Union, cast, get_args, - overload, + get_origin, override, ) @@ -19,11 +20,10 @@ from pacman.ecs.schedule import ( ScheduleLabel, SystemSet, ) -from pacman.utils.variadics_please import VarargCallable, VarargTuple - - -class Entity: - pass +from pacman.utils.variadics_please import ( + RecursiveTuple, + fast_tuple_constructor, +) @dataclass @@ -31,13 +31,21 @@ class Resource2[T]: storage: type[T] -class SkipSystemError(Exception): - pass +class Archetype: + def __init__(self, types: Iterable[type["Component"]]) -> None: + 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._remove_cache: dict[ + int, tuple[Any, Archetype, tuple[int, ...]] + ] = {} -type RawSystem = Callable[["World"], None] +type RawSystem[T] = Callable[[], T] -type Plugin = Callable[["World"], None] | tuple[Plugin, ...] +type Plugin = RecursiveTuple[Callable[["World"], None]] type ComponentHook[T] = Callable[["World", "Entity", T], None] @@ -45,233 +53,307 @@ type ComponentHook[T] = Callable[["World", "Entity", T], None] class FromWorld(Protocol): @classmethod @abstractmethod - def extract_from_world(cls, world: "World") -> Self | None: ... + def extract_from_world( + cls, world: "World" + ) -> Callable[[], Self | None]: ... -type SystemParam = VarargTuple[SystemParam] | World | FromWorld +type SystemParam = RecursiveTuple[FromWorld] -type System = VarargCallable[SystemParam, None] +type System[T] = Callable[..., T] -class QueryItem: - pass +class QueryFilter: + @classmethod + @abstractmethod + def filter( + cls, archetype: "Archetype" + ) -> bool | Callable[[int], bool]: ... -type QueryItems = VarargTuple[QueryItems] | QueryItem +type QueryFilters = RecursiveTuple[QueryFilter] -class QueryFilter: - pass +class QueryItem(QueryFilter): + @override + @classmethod + def filter(cls, archetype: "Archetype") -> bool | Callable[[int], bool]: + filter_fetch = cls.filter_fetch(archetype) + if filter_fetch is None: + return lambda _: False + return lambda i: filter_fetch(i) is not None + @classmethod + @abstractmethod + def filter_fetch( + cls, archetype: "Archetype" + ) -> None | Callable[[int], Self]: ... -type QueryFilters = VarargTuple[QueryFilters] | QueryFilter +type QueryItems = RecursiveTuple[QueryItem] -class Query[Items: QueryItems, Filters: QueryFilters = tuple[()]]( - FromWorld, Iterable[Items] -): - def __init__( - self, world: "World", items: type[Items], filters: type[Filters] - ) -> None: - pass +class Component(QueryItem): @override @classmethod - def extract_from_world(cls, world: "World") -> Self | None: - return cls(world, *get_args(cls)) + def filter(cls, archetype: "Archetype") -> bool: + return cls in archetype._components @override - def __iter__(self) -> Iterator[Items]: - if False: - yield + @classmethod + def filter_fetch( + cls, archetype: "Archetype" + ) -> None | Callable[[int], Self]: + idx = archetype._components.get(cls) + if idx is None: + return None + return cast(Callable[[int], Self], lambda i: idx[i]) -class Resource(FromWorld): +class Entity(QueryItem): @override @classmethod - def extract_from_world(cls, world: "World") -> Self | None: + def filter(cls, archetype: "Archetype") -> bool: + return True - pass + @override + @classmethod + def filter_fetch( + cls, archetype: "Archetype" + ) -> None | Callable[[int], "Entity"]: + return lambda i: archetype._entities[i] -@dataclass -class EntityThunk: - world: "World" - entity: Entity - - def __getitem__[T](self, ty: type[T]) -> T: - return cast(T, self.world._components[ty][self.entity]) - - def __setitem__[T](self, ty: type[T], val: T) -> None: - if self.entity not in self.world._entities: - self.world._entities[self.entity] = set() - if ty not in self.world._entities[self.entity]: - self.world._entities[self.entity].add(ty) - if ty not in self.world._components: - self.world._components[ty] = {} - if self.entity in self.world._components[ty]: - del self[ty] - self.world._components[ty][self.entity] = val - if ty in self.world._insert_hooks: - self.world._insert_hooks[ty](self.world, self.entity, val) - - def __delitem__[T](self, ty: type[T]) -> None: - if ty not in self.world._entities[self.entity]: +class Query[Items: QueryItems, Filters: QueryFilters = tuple[()]]( + FromWorld, Iterable[Items], Container[Entity] +): + def __init__( + self, world: "World", fetchers: type[Items], filters: type[Filters] + ) -> None: + self._world: World = world + self._fetchers_raw = cast( + Callable[[Archetype], Callable[[int], Items] | None], + Query.preprocess_fetchers(fetchers), + ) + self._filters_raw = Query.preprocess_filters(filters) + self._seen_archetypes: int = 0 + self._archetypes_idx: dict[Archetype, int] = {} + self._archetypes: list[Archetype] = [] + self._filters: list[Callable[[int], bool]] = [] + self._fetchers: list[Callable[[int], Items]] = [] + self.updated_from_world() + + @staticmethod + def preprocess_fetchers( + fetcher: type[QueryItems], + ) -> Callable[[Archetype], Callable[[int], QueryItems] | None]: + args = list(map(Query.preprocess_fetchers, get_args(fetcher))) + + def tuple_constructor( + archetype: Archetype, + ) -> Callable[[int], QueryItems] | None: + lst = [] + for arg in args: + curr = arg(archetype) + if curr is None: + return None + lst.append(curr) + return fast_tuple_constructor(tuple(lst)) + + def union_constructor( + archetype: Archetype, + ) -> Callable[[int], QueryItems] | None: + lst: list[Callable[[int], QueryItems]] = [] + for arg in args: + curr = arg(archetype) + if curr is None: + continue + lst.append(curr) + if len(lst) == 0: + return None + return lst.pop() + + if get_origin(fetcher) is tuple: + return tuple_constructor + if get_origin(fetcher) is Union: + return union_constructor + + return cast(QueryItem, fetcher).filter_fetch + + @staticmethod + def preprocess_filters( + filter: type[QueryFilters], + ) -> Callable[[Archetype], Callable[[int], bool] | bool]: + args = list(map(Query.preprocess_filters, get_args(filter))) + + def composed_by( + a: Callable[[int], bool] | None, + b: Callable[[int], bool], + op: Callable[[bool, bool], bool], + ) -> Callable[[int], bool]: + if a is None: + return b + return lambda i: op(a(i), b(i)) + + def tuple_constructor( + archetype: Archetype, + ) -> Callable[[int], bool] | bool: + cond: Callable[[int], bool] | None = None + for arg in args: + curr = arg(archetype) + if isinstance(curr, bool): + if curr: + continue + return False + cond = composed_by(cond, curr, bool.__and__) + if cond is None: + return lambda _: True + return cond + + def union_constructor( + archetype: Archetype, + ) -> Callable[[int], bool] | bool: + cond: Callable[[int], bool] | None = None + for arg in args: + curr = arg(archetype) + if isinstance(curr, bool): + if curr: + return True + continue + cond = composed_by(cond, curr, bool.__or__) + if cond is None: + return lambda _: False + return cond + + if get_origin(filter) is tuple: + return tuple_constructor + if get_origin(filter) is Union: + return union_constructor + + return cast(QueryFilter, filter).filter + + def add_archetype(self, archetype: Archetype) -> None: + cond: Callable[[int], bool] + match self._filters_raw(archetype): + case True: + cond = lambda _: True + case False: + return + case e: + cond = e + + fetcher = self._fetchers_raw(archetype) + if fetcher is None: return - if ty in self.world._remove_hooks: - self.world._remove_hooks[ty]( - self.world, - self.entity, - self.world._components[ty][self.entity], + self._archetypes_idx[archetype] = len(self._archetypes) + self._archetypes.append(archetype) + self._filters.append(cond) + self._fetchers.append(fetcher) + + def updated_from_world(self) -> Self: + for arch in self._world._archetypes[self._seen_archetypes :]: + self.add_archetype(arch) + self._seen_archetypes = len(self._world._archetypes) + return self + + @override + @classmethod + def extract_from_world(cls, world: "World") -> Callable[[], Self]: + slf: Self = cls(world, *get_args(cls)) + return slf.updated_from_world + + @override + def __iter__(self) -> Iterator[Items]: + return ( + fetch(i) + for filt, fetch, arch in zip( + self._filters, self._fetchers, self._archetypes, strict=True ) - self.world._entities[self.entity].remove(ty) - del self.world._components[ty][self.entity] + for i in range(len(arch._entities)) + if filt(i) + ) - def __contains__[T](self, ty: type[T]) -> bool: - return ty in self.world._entities[self.entity] + @override + def __contains__(self, entity: Entity) -> bool: + arch: Archetype = self._world._entities[entity] + if (arch_idx := self._archetypes_idx.get(arch)) is None: + return False + components_idx = arch._positions[entity] + return self._filters[arch_idx](components_idx) + + def single(self) -> Items | None: + itr: Iterator[Items] = self.__iter__() + res = next(itr, None) + if next(itr, None) is not None: + return None + return res + + +class Resource(FromWorld, Component): + @override + @classmethod + def extract_from_world(cls, world: "World") -> Callable[[], Self | None]: + _query = Query[cls] # type: ignore + query: Callable[[], Query[Self]] = cast( + type[Query[Self]], _query + ).extract_from_world(world) + return lambda: query().single() -class World: +class World(FromWorld): + _BASE_ARCHETYPE: int = 0 + def __init__(self) -> None: - self._entities: dict[Entity, set[type]] = {} - self._components: dict[type, dict[Entity, Any]] = {} - self._resources: dict[type, Any] = {} + 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._schedules: dict[ScheduleLabel, Schedule[RawSystem]] = {} self._insert_hooks: dict[type, ComponentHook[Any]] = {} self._remove_hooks: dict[type, ComponentHook[Any]] = {} - def entity(self, entity: Entity) -> EntityThunk: - return EntityThunk(self, entity) + def _run_cond(self, _: RawSystem | SystemSet) -> bool: + return True - def query[*T, *U]( - self, ty: type[tuple[*T]], without: type[tuple[*U]] | None = None - ) -> Generator[tuple[*T]]: - args = get_args(ty) - neg_args = get_args(without) - if len(args) == 0: - return - sets: list[dict[Entity, Any]] = sorted( - ( - self._entities - if arg is Entity - else self._components.get(arg, {}) - for arg in args - if arg is Entity or arg in self._components - ), - key=len, - reverse=True, - ) - neg_sets: list[dict[Entity, Any]] = sorted( - ( - self._entities - if arg is Entity - else self._components.get(arg, {}) - for arg in neg_args - if arg is Entity or arg in self._components - ), - key=len, - reverse=True, + def add_system_raw[T]( + self, system: System[T], default: Callable[[], T] + ) -> RawSystem[T]: + sig = inspect.signature(system) + extractors = tuple( + cast(type[FromWorld], param.annotation).extract_from_world(self) + for param in sig.parameters.values() ) - if len(sets) == 0: - return - for entity in sets.pop(): - if any(entity not in e for e in sets) or any( - entity in e for e in neg_sets - ): - continue - yield cast( - tuple[*T], - tuple( - entity if e is Entity else self._components[e][entity] - for e in args - ), - ) - - @overload - def __getitem__[T](self, arg: Resource2[T]) -> T: ... - @overload - def __getitem__[*T]( - self, arg: type[tuple[*T]] - ) -> Generator[tuple[*T]]: ... + def inner() -> T: + args = [] + for extract in extractors: + curr = extract() + if curr is None: + return default() + args.append(curr) + return system(*args) - @overload - def __getitem__(self, arg: Entity) -> EntityThunk: ... + return inner - def __getitem__( - self, - arg: Any, - ) -> Any: - if isinstance(arg, Entity): - return self.entity(arg) - if isinstance(arg, Resource2): - return self._resources[arg.storage] - return self.query(arg) - - @overload - def __setitem__[T](self, key: Resource2[T], val: T) -> None: ... - - @overload - def __setitem__[*T](self, key: Entity, val: tuple[*T]) -> None: ... - - def __setitem__(self, key: Any, val: Any) -> None: - if isinstance(key, Entity) and isinstance(val, tuple): - if key in self: - del self[key] - self._entities[key] = set() - for component in val: - self[key][type(component)] = component - elif isinstance(key, Resource2) and isinstance(val, key.storage): - self._resources[key.storage] = val - else: - raise TypeError() - - @overload - def __delitem__[T](self, key: Resource2[T]) -> None: ... - - @overload - def __delitem__(self, key: Entity) -> None: ... - - def __delitem__[T](self, key: Entity | Resource2[T]) -> None: - if isinstance(key, Entity): - components = self._entities[key] - del self._entities[key] - for component in components: - del self._components[component][key] - else: - del self._resources[key.storage] - - @overload - def __contains__[T](self, item: Resource2[T]) -> bool: ... - @overload - def __contains__(self, item: Entity) -> bool: ... - - def __contains__[T](self, item: Entity | Resource2[T]) -> bool: - if isinstance(item, Resource2): - return item.storage in self._resources - return item in self._entities - - def _run_cond(self, _: RawSystem | SystemSet) -> bool: - return True + @classmethod + @override + def extract_from_world(cls, world: "World") -> Callable[[], "World"]: + return lambda: world def tick(self, schedule: ScheduleLabel) -> None: if schedule not in self._schedules: return for system in self._schedules[schedule].traverse(self._run_cond): - with contextlib.suppress(SkipSystemError): - system(self) - - def res_s[T](self, ty: type[T]) -> T: - if Resource2(ty) not in self: - raise SkipSystemError() - return self[Resource2(ty)] + system() def run( self, schedule: ScheduleLabel, - stop_cond: Callable[["World"], bool] = lambda _: False, + stop_cond: RawSystem[bool] = lambda: False, ) -> None: - while not stop_cond(self): + while not stop_cond(): self.tick(schedule) def with_systems( @@ -320,9 +402,7 @@ class World: return self def run_main(self) -> None: - self.run( - MainSchedule, lambda world: Resource2(WorldShouldExit) in world - ) + self.run(MainSchedule, self.add_system_raw(should_exit, lambda: False)) _order_groups: int = 0 @@ -375,5 +455,9 @@ class WorldIsInit: pass -class WorldShouldExit: +class WorldShouldExit(Resource): pass + + +def should_exit(_: WorldShouldExit) -> bool: + return True diff --git a/src/pacman/utils/type_resolve.py b/src/pacman/utils/type_resolve.py index d555c13..210cc94 100644 --- a/src/pacman/utils/type_resolve.py +++ b/src/pacman/utils/type_resolve.py @@ -1,12 +1,20 @@ # ruff: noqa: ANN401 +import types +import typing from collections.abc import Iterable from typing import ( Any, + cast, get_args, get_origin, ) +from pacman.utils.variadics_please import ( + RecursiveTuple, + recursive_tuple_create, +) + def map_type_generics(ty: Any, mapping: dict[Any, type]) -> Any: """Map generics for this type, maybe creating a new type.""" @@ -17,6 +25,8 @@ def map_type_generics(ty: Any, mapping: dict[Any, type]) -> Any: args = tuple(map_type_generics(e, mapping) for e in get_args(ty)) if hasattr(ty, "copy_with") and callable(ty.copy_with): return ty.copy_with(args) + if orig is types.UnionType: + orig = typing.Union return orig[args] @@ -36,17 +46,25 @@ def resolve_type_aliases(ty: Any) -> Any: ) +def resolve_ty_to_tup[T]( + tup: type[RecursiveTuple[T]], +) -> RecursiveTuple[type[T]]: + def inner( + ty: type[RecursiveTuple[T]], + ) -> type[T] | tuple[type[RecursiveTuple[T]], ...]: + base = get_origin(ty) + if base is tuple: + return cast(tuple[type[RecursiveTuple[T]], ...], get_args(ty)) + return cast(type[T], ty) + + return recursive_tuple_create(tup, inner) + + type B[T] = list[list[T]] type A[T] = list[tuple[T, B[int]]] if __name__ == "__main__": - alp = A[str] - print(alp) - al = resolve_type_aliases(alp) - print(al) - al2 = list[tuple[str, list[list[int]]]] - print(al2) - al3 = resolve_type_aliases(al2) - print(al3) - print(al == al2) + ty = tuple[int, tuple[float, str], tuple[int, int] | float] + print(ty) + print(resolve_ty_to_tup(resolve_type_aliases(ty))) diff --git a/src/pacman/utils/variadics_please.py b/src/pacman/utils/variadics_please.py index bd773fa..93ecbf1 100644 --- a/src/pacman/utils/variadics_please.py +++ b/src/pacman/utils/variadics_please.py @@ -1,4 +1,5 @@ -from collections.abc import Callable +from collections.abc import Callable, Generator +from typing import cast type VarargCallable[T, U] = ( Callable[[T], U] @@ -17,21 +18,201 @@ type VarargCallable[T, U] = ( | Callable[[T, T, T, T, T, T, T, T, T, T, T, T, T, T], U] | Callable[[T, T, T, T, T, T, T, T, T, T, T, T, T, T, T], U] ) -type VarargTuple[T] = ( - tuple[()] - | tuple[T] - | tuple[T, T] - | tuple[T, T, T] - | tuple[T, T, T, T] - | tuple[T, T, T, T, T] - | tuple[T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T, T, T, T, T, T] - | tuple[T, T, T, T, T, T, T, T, T, T, T, T, T, T, T] -) + +type RecursiveTuple[T] = T | tuple[RecursiveTuple[T], ...] + + +def recursive_tuple_map[T, U]( + tup: RecursiveTuple[T], map: Callable[[T], U] +) -> RecursiveTuple[U]: + if not isinstance(tup, tuple): + return map(tup) + return tuple(recursive_tuple_map(e, map) for e in tup) + + +def recursive_tuple_create[T, U]( + tup: T, map: Callable[[T], U | tuple[T, ...]] +) -> RecursiveTuple[U]: + if not isinstance((res := map(tup)), tuple): + return res + return tuple( + cast(RecursiveTuple[U], recursive_tuple_create(tup, map)) + for tup in res + ) + + +def recursive_tuple_iter[T](tup: RecursiveTuple[T]) -> Generator[T]: + if not isinstance(tup, tuple): + yield tup + else: + for sub in cast(tuple[RecursiveTuple[T], ...], tup): + yield from recursive_tuple_iter(sub) + + +def fast_tuple_constructor[T, U]( + cbs: tuple[Callable[[U], T], ...], +) -> Callable[[U], tuple[T, ...]]: + match cbs: + case (e1,): + return lambda arg: (e1(arg),) + case (e1, e2): + return lambda arg: (e1(arg), e2(arg)) + case (e1, e2, e3): + return lambda arg: (e1(arg), e2(arg), e3(arg)) + case (e1, e2, e3, e4): + return lambda arg: (e1(arg), e2(arg), e3(arg), e4(arg)) + case (e1, e2, e3, e4, e5): + return lambda arg: (e1(arg), e2(arg), e3(arg), e4(arg), e5(arg)) + case (e1, e2, e3, e4, e5, e6): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8, e9): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8, e9, e10): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + e10(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8, e9, e10, e11): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + e10(arg), + e11(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8, e9, e10, e11, e12): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + e10(arg), + e11(arg), + e12(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8, e9, e10, e11, e12, e13): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + e10(arg), + e11(arg), + e12(arg), + e13(arg), + ) + case (e1, e2, e3, e4, e5, e6, e7, e8, e9, e10, e11, e12, e13, e14): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + e10(arg), + e11(arg), + e12(arg), + e13(arg), + e14(arg), + ) + case ( + e1, + e2, + e3, + e4, + e5, + e6, + e7, + e8, + e9, + e10, + e11, + e12, + e13, + e14, + e15, + ): + return lambda arg: ( + e1(arg), + e2(arg), + e3(arg), + e4(arg), + e5(arg), + e6(arg), + e7(arg), + e8(arg), + e9(arg), + e10(arg), + e11(arg), + e12(arg), + e13(arg), + e14(arg), + e15(arg), + ) + case e: + return lambda arg: tuple(cb(arg) for cb in e) -- 2.53.0