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