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