microecs.world

  1"""world.py - The world container for ECS. It manages all the pools (one per archetype). Entities are id-based."""
  2from typing import get_type_hints
  3from dataclasses import fields
  4import numpy as np
  5
  6from .utils import Shape, EntityId, PoolKey, logger
  7from .component import ComponentType
  8from .query_result import QueryResult, QUERY_RESULT_RESERVED_NAMES
  9from .entity import Entity, ENTITY_RESERVED_NAMES
 10from .pool import Pool, POOL_RESERVED_NAMES
 11from .command_buffer import CommandBuffer, Command, CommandType
 12
 13class World:
 14    """
 15    Generic container for pools of components. Newly added components are assigned a unique id and go in a pool
 16    Parameters
 17    - components The list of components that the world accepts
 18    - extra_metadata The list of required extra metadata for each field besides shape and dtype.
 19    """
 20    def __init__(self, components: list[ComponentType], extra_metadata: list[str] | None = None):
 21        self._default_metadata = {"shape", "dtype", "default"}
 22        self.extra_metadata = extra_metadata or []
 23        if not isinstance(self.extra_metadata, list):
 24            raise TypeError(type(self.extra_metadata))
 25        self._check_components(components)
 26
 27        # Pools management
 28        self.pools: dict[PoolKey, Pool] = {}
 29        self.pool_to_components: dict[Pool, list[ComponentType]] = {}
 30
 31        # Components management
 32        self.field_to_component: dict[str, ComponentType] = {}
 33        self.component_names = [x.__name__ for x in components]
 34        self.component_types = set(components)
 35        self.component_name_to_type = {x.__name__: x for x in components}
 36        self.component_to_bit: dict[ComponentType, int] = {t: 2**i for i, t in enumerate(components)} # bit for querying
 37        self.component_to_field_names: dict[ComponentType, list[str]] = {c: [] for c in components}
 38        self.component_to_shapes: dict[ComponentType, list[Shape]] = {c: [] for c in components}
 39        self.component_to_dtypes: dict[ComponentType, list[str]] = {c: [] for c in components}
 40        self.component_to_defaults: dict[ComponentType, list[np.ndarray]] = {c: [] for c in components}
 41        # setup the obligatory metadata at each fields
 42        for c in components:
 43            for f in fields(c):
 44                self.component_to_field_names[c].append(f.name)
 45                self.component_to_shapes[c].append(field_shape := f.metadata["shape"])
 46                self.component_to_dtypes[c].append(field_dtype := f.metadata["dtype"])
 47                self.component_to_defaults[c].append(field_default := f.metadata["default"])
 48
 49                if f.name in self.field_to_component:
 50                    raise ValueError(f"Duplicate field '{c.__name__}/{f.name}': {self.field_to_component[f.name]}")
 51                self.field_to_component[f.name] = c
 52
 53                if field_default is not None:
 54                    if (dt := field_default.dtype) != field_dtype:
 55                        raise TypeError(f"'{c.__name__}/{f.name}'. Expected dtype: {field_dtype}. Got: {dt}")
 56                    if (sh := field_default.shape) != field_shape:
 57                        raise ValueError(f"'{c.__name__}/{f.name}'. Expected shape: {field_shape}. Got: {sh}")
 58
 59        # Entities management
 60        self._eid_to_pool_ix: dict[EntityId, tuple[Pool, int]] = {}
 61        self._pool_ids: dict[Pool, list[EntityId]] = {}
 62        self._last_id: EntityId = -1
 63        # A dictionary of all live entities in 'eager' mode (before update()). The actual entity is created at request
 64        # in get_entity, so we don't pay for the Entity object unless it's explicitly requested by the user.
 65        self.live_entities: dict[EntityId, Entity | None] = {}
 66
 67        # Command buffer management. {add/remove}_{entity/component} are lazy. Taken into account after update().
 68        self._command_buffer = CommandBuffer(self)
 69
 70        # QueryResult cache (key:tuple[include, exclude] - see query()). Useful so we re-use qrs between world.updates.
 71        self._qr_cache: dict[tuple[PoolKey, PoolKey], QueryResult] = {}
 72        logger.debug(f"Created scene with components: {self.component_names}")
 73
 74    # public api
 75
 76    def add_entity(self, components: list[ComponentType], **kwargs) -> EntityId:
 77        """Adds an entity to the world based on components (data->kwargs). Returns an entity id. Lazy; call update()"""
 78        self._validate_components(components, **kwargs)
 79        default_kwargs = self._defaults_for(components, **kwargs)
 80        self._last_id += 1
 81        self.live_entities[self._last_id] = None # add the id the live_entities, but the object is created in get_entity
 82        self._command_buffer.append(Command(CommandType.ADD_ENTITY, self._last_id,
 83                                            args={"components": components, **kwargs, **default_kwargs}))
 84        return self._last_id
 85
 86    def remove_entity(self, entity_id: EntityId):
 87        """Removes an entity based on its unique entity id. The same entity can be safely removed multiple times in the
 88        same tick (e.g. by different systems). Raises if the eid is not (anymore) in the world. Lazy; call update()"""
 89        if entity_id not in self.live_entities:
 90            # NOTE: the only reason this may happen is if we called remove_entity >=2 times before world.update()
 91            if entity_id not in self._command_buffer.removed_this_tick:
 92                raise ValueError(f"Entity: {entity_id} is not in the world (stale). Either wrong id or removed earlier")
 93            return
 94        self._command_buffer.append(Command(CommandType.REMOVE_ENTITY, entity_id))
 95        del self.live_entities[entity_id]
 96
 97    def get_entity(self, entity_id: EntityId) -> Entity:
 98        """Gets the entity reference given an entity id. Used for 'object-like' ops. Structural updates
 99        (add/remove_component) are lazy, so you need to call world.update(). Data updates (set_data) are immediate."""
100        try:
101            entity = self.live_entities[entity_id]
102        except KeyError:
103            raise ValueError(f"Entity id: {entity_id} not in the world")
104
105        if entity is None: # this can happen if it was just added by add_entity() but not materialized yet
106            self.live_entities[entity_id] = entity = Entity(
107                entity_id, eid_to_pool_ix=self._eid_to_pool_ix, pool_to_components=self.pool_to_components,
108                world_command_buffer=self._command_buffer)
109
110        return entity
111
112    def query(self, *include: ComponentType, exclude: list[ComponentType] | None = None) -> QueryResult:
113        """
114        Queries the world for entities that match the include set of components.
115        Syntax: `world.query(A, B, exclude=[C, D])` is, in logic form, a chain of 'ands': `A & B & ~C & ~D`.
116        Return: A `QueryResult` object with the entities that have all the requested components. Has `EntityIds` too.
117        """
118
119        # Note: we can cache the queries. The only time it can get invalidated (via public API) is at update().
120        include_key = self._make_key(include)
121        exclude_key = self._make_key(exclude or [])
122        if (key := (include_key, exclude_key)) in self._qr_cache:
123            return self._qr_cache[key]
124
125        # archetype_key = (1 0 0 1 1) &
126        #           key = (1 0 0 0 1)
127        #              -> (1 0 0 0 1) OK
128        # but
129        # archetype_key = (1 0 0 1 1) &
130        #           key = (1 0 1 0 0)
131        #              -> (1 0 0 0 0) NOT OK
132        res = []
133        for archetype_key, archetype_pool in self.pools.items():
134            if (archetype_key & include_key) == include_key and (archetype_key & exclude_key) == 0:
135                res.append(archetype_pool)
136
137        field_names = sum([self.component_to_field_names[c] for c in include], [])
138        field_shapes = dict(zip(field_names, sum([self.component_to_shapes[c] for c in include], [])))
139        field_dtypes = dict(zip(field_names, sum([self.component_to_dtypes[c] for c in include], [])))
140        self._qr_cache[key] = QueryResult(res, field_shapes, field_dtypes=field_dtypes, pool_ids=self._pool_ids)
141        return self._qr_cache[key]
142
143    def update(self):
144        """commits the underlying pool changes from the systems between two updates. Should be called in main loop."""
145        for command in self._command_buffer:
146            if command.command_type == CommandType.ADD_ENTITY:
147                components = command.args.pop("components")
148                self._add_to_pool(command.entity_id, components=components, entity_data=command.args)
149            elif command.command_type == CommandType.REMOVE_ENTITY:
150                self._remove_from_pool(command.entity_id)
151            elif command.command_type == CommandType.ADD_COMPONENT:
152                component = command.args.pop("component")
153                self._do_add_component(command.entity_id, component=component, **command.args)
154            elif command.command_type == CommandType.REMOVE_COMPONENT:
155                self._do_remove_component(command.entity_id, component=command.args)
156            else: # CommandType.REMOVE_COMPONENT
157                raise NotImplementedError(command)
158
159        # Check if there's any empty pool since the last movements and remove it, if so.
160        empty_keys = [pool_key for pool_key, pool in self.pools.items() if len(pool) == 0]
161        for pool_key in empty_keys:
162            pool = self.pools.pop(pool_key)
163            del self.pool_to_components[pool]
164            del self._pool_ids[pool]
165
166        if len(self._command_buffer) > 0:
167            self._qr_cache.clear()
168        self._command_buffer.clear() # .clear() here also clears the buffer.removed_this_tick set.
169
170    # private stuff
171
172    # eager mode methods equivalent to add/remove entities and add/remove_components
173
174    def _add_to_pool(self, entity_id: EntityId, components: list[ComponentType], entity_data: dict[str, np.ndarray]):
175        """adds the item to the pool"""
176        pool = self._get_entity_pool(components)
177        pool_index = pool.add_entity(entity_data=entity_data)
178        self._eid_to_pool_ix[entity_id] = (pool, pool_index)
179        self._pool_ids.setdefault(pool, []).append(entity_id)
180        assert len(self._pool_ids[pool]) == len(pool), (pool, len(self._pool_ids[pool]), len(pool))
181
182    def _remove_from_pool(self, entity_id: EntityId):
183        """removes the entity from the pool w/o any internal data copying, like _pop_from_pool"""
184        old_pool, pool_ix = self._eid_to_pool_ix.pop(entity_id)
185        old_pool.remove_entity(pool_ix)
186        # Move the last id from the pool in the place of the recently removed entity
187        id_which_was_last_in_pool = self._pool_ids[old_pool].pop()
188
189        if entity_id != id_which_was_last_in_pool:
190            self._eid_to_pool_ix[id_which_was_last_in_pool] = (old_pool, pool_ix) # we re-use the popped id (swapped)
191            self._pool_ids[old_pool][pool_ix] = id_which_was_last_in_pool
192
193    def _pop_from_pool(self, entity_id: EntityId) -> tuple[dict[str, np.ndarray], list[ComponentType]]:
194        """common function that updates the entities inside a pool (after popswap) and removes them if they get empty"""
195        old_pool, pool_ix = self._eid_to_pool_ix.pop(entity_id)
196        entity_data = old_pool.pop_entity(pool_ix)
197        components = self.pool_to_components[old_pool]
198
199        # Move the last id from the pool in the place of the recently removed entity
200        id_which_was_last_in_pool = self._pool_ids[old_pool].pop()
201        if entity_id != id_which_was_last_in_pool:
202            self._eid_to_pool_ix[id_which_was_last_in_pool] = (old_pool, pool_ix) # we re-use the popped id (swapped)
203            self._pool_ids[old_pool][pool_ix] = id_which_was_last_in_pool
204
205        return entity_data, components
206
207    def _do_add_component(self, entity_id: EntityId, component: ComponentType, **kwargs):
208        curr_data, curr_components = self._pop_from_pool(entity_id)
209        if not curr_data.keys().isdisjoint(kwargs):
210            raise ValueError(f"Duplicate keys: {curr_data.keys()} vs {kwargs.keys()}")
211
212        new_components = [*curr_components, component]
213        default_kwargs = self._defaults_for(new_components, **curr_data, **kwargs)
214        new_data = {**curr_data, **kwargs, **default_kwargs}
215        self._add_to_pool(entity_id, components=new_components, entity_data=new_data)
216
217    def _do_remove_component(self, entity_id: EntityId, component: ComponentType):
218        entity_data, components = self._pop_from_pool(entity_id)
219        for _field in self.component_to_field_names[component]:
220            assert _field in entity_data, f"Field {component}/{_field} not in components: {components} ({entity_id=})"
221            entity_data.pop(_field)
222        new_components = [c for c in components if c != component]
223        self._add_to_pool(entity_id, components=new_components, entity_data=entity_data)
224
225    # other low-level methods
226
227    def _validate_component(self, component: ComponentType, strict: bool, check_extra: bool, **kwargs):
228        """
229        Validates a single component.
230        Parameters:
231        - strict Two modes: True -> kwargs==component fields, False -> kwargs is subset
232        - check_extra If set, extra data in kwargs unrelated to this component also raises (_validate_components)
233        """
234        if check_extra and (extra := set(kwargs) - set(self.component_to_field_names[component])):
235            raise ValueError(f"Extra fields: {extra}; expected {self.component_to_field_names[component]}")
236
237        for name, shape, dtype, default in zip(
238                self.component_to_field_names[component], self.component_to_shapes[component],
239                self.component_to_dtypes[component], self.component_to_defaults[component]):
240            if name not in kwargs:
241                if default is None and strict is True:
242                    raise KeyError(f"'{component.__name__}/{name}' required (default=None) but not supplied")
243                continue                      # omitted but has a default -> fine (or non-strict mode for set_data)
244            if not isinstance(field := kwargs[name], np.ndarray):
245                raise TypeError(f"'{component.__name__}/{name}'. Expected np.ndarray, got {type(field)}")
246            if (dt := field.dtype) != dtype:
247                raise TypeError(f"'{component.__name__}/{name}'. Expected dtype {dtype}, got {dt}")
248            if (sh := field.shape) != shape:
249                raise ValueError(f"'{component.__name__}/{name}'. Expected shape {shape}, got {sh}")
250
251    def _validate_components(self, components: list[ComponentType], **kwargs):
252        """Pure check. Raises on: no components, unknown component, missing-required (default=None),
253            wrong dtype/shape, extra field. No mutation, no return. kwargs == fields data."""
254        if len(cs := set(components)) == 0:
255            raise ValueError(f"Entity has no components: {self.component_names}")
256        if diff := cs - self.component_types:
257            raise ValueError(f"Unknown components: {diff}")
258        if len(cs) != len(components):
259            raise ValueError(f"Duplicate components: {components}")
260
261        expected = set()
262        for component in components:
263            expected.update(self.component_to_field_names[component]) # updated here for error message.
264            self._validate_component(component, strict=True, check_extra=False, **kwargs)
265        if extra := set(kwargs) - expected:
266            raise ValueError(f"Extra fields: {extra}; expected {expected}")
267
268    def _defaults_for(self, components: list[ComponentType], **kwargs) -> dict[str, np.ndarray]:
269        """Defaults for omitted fields. Assumes already validated. No mutation of `kwargs` (data) is done."""
270        res = {}
271        for c in components:
272            for name, default in zip(self.component_to_field_names[c], self.component_to_defaults[c]):
273                if name not in kwargs and default is not None:
274                    res[name] = default.copy()
275        return res
276
277    def _get_entity_pool(self, components: list[ComponentType]) -> Pool:
278        if (key := self._make_key(components)) not in self.pools:
279            _fields = sum([self.component_to_field_names[c] for c in components], []) # merge fields
280            shapes = sum([self.component_to_shapes[c] for c in components], []) # merge shapes
281            dtypes = sum([self.component_to_dtypes[c] for c in components], []) # merge dtypes
282            self.pools[key] = Pool(_fields, shapes, dtypes)
283            self.pool_to_components[self.pools[key]] = components
284        return self.pools[key]
285
286    def _make_key(self, components: list[ComponentType]) -> PoolKey:
287        key = 0
288        for c in components:
289            if c not in self.component_types:
290                raise ValueError(f"c '{c.__name__}' not in {self.component_names}")
291            key |= self.component_to_bit[c]
292        return key
293
294    def _check_components(self, components: list[ComponentType]):
295        reserved_names = ENTITY_RESERVED_NAMES | POOL_RESERVED_NAMES | QUERY_RESULT_RESERVED_NAMES
296        dtypes = {"float32", "int32", "bool", "object"}
297        expected_meta = {*self._default_metadata, *self.extra_metadata}
298
299        for c in components:
300            if not hasattr(c, "__dataclass_fields__"):
301                raise TypeError(f"c '{c.__name__}' is not a dataclass")
302
303            hints = get_type_hints(c) # make it work with from __future__ import annotations
304            for f in fields(c):
305                if hints[f.name] is not np.ndarray:
306                    raise TypeError(f"Field '{c.__name__}/{f.name}' not an array: {f}")
307                if f.name in reserved_names:
308                    raise ValueError(f"Field '{c.__name__}/{f.name}' in {reserved_names}")
309                if f.metadata.keys() != expected_meta:
310                    raise ValueError(f"Field '{c.__name__}/{f.name}'\n{list(f.metadata.keys())}\nvs\n{expected_meta}\n"
311                                      "Perhaps missing World(extra_metadata=[...])?")
312                if not isinstance(f.metadata["shape"], tuple):
313                    raise TypeError(f"Expected tuple, got {type(f.metadata['shape'])}: {f.metadata['shape']}")
314                if not isinstance(fmd := f.metadata["dtype"], str) or fmd not in dtypes:
315                    raise TypeError(f"{fmd} not a string or not in {dtypes}")
316
317    def __len__(self):
318        return len(self.live_entities)
319
320    def __repr__(self):
321        return (f"[World]\n- Entities: {len(self)} (last id: {self._last_id})"
322                f"\n- Components ({len(self.component_names)}): {self.component_names}"
323                f"\n- Pools: {len(self.pools)}\n- Command buffer: {len(self._command_buffer)}")