microecs.command_buffer

  1# pylint: disable=protected-access
  2"""command_buffer.py - data structure that eagerly manages the staging area (add/rm entity/cmpt) before world.upate()"""
  3from __future__ import annotations
  4from typing import Any, TYPE_CHECKING
  5from dataclasses import dataclass, field
  6from enum import StrEnum
  7
  8from .utils import EntityId
  9from .component import ComponentType
 10
 11if TYPE_CHECKING:
 12    from microecs import World
 13
 14class CommandType(StrEnum):
 15    """The types of commands in the command pattern below"""
 16    ADD_ENTITY       = "add_entity"
 17    REMOVE_ENTITY    = "remove_entity"
 18    ADD_COMPONENT    = "add_component"
 19    REMOVE_COMPONENT = "remove_component"
 20
 21@dataclass
 22class Command:
 23    """A command that can happen between two world.updates(), e.g. add/rm entity or components"""
 24    command_type: CommandType                            # The type of command in the command buffer
 25    entity_id: EntityId                                  # The entity this command applies to
 26    args: dict[str, Any] = field(default_factory=dict)   # The arguments of this command. Each command has its own shape
 27
 28class CommandBuffer:
 29    """A data structure that holds all the uncommited commands between two world updates. Support eager exceptions
 30       on things like adding the same component twice on the same entity"""
 31    def __init__(self, world: "World"): # noqa
 32        self.data: list[Command] = [] # should we make this dict[EntityId, list[Command]] ?
 33        self.world = world
 34        self.removed_this_tick: set[EntityId] = set()
 35
 36    def clear(self):
 37        """Clears the buffer"""
 38        self.data.clear()
 39        self.removed_this_tick.clear()
 40
 41    def _get_projected_entity_components(self, entity_id: EntityId) -> set[ComponentType]:
 42        """Gets the projected components of an entity given this command buffer's state (how it'd look after update)"""
 43        if entity_id in self.world._eid_to_pool_ix:
 44            pool, _ = self.world._eid_to_pool_ix[entity_id]
 45            projected = set(self.world.pool_to_components[pool])
 46        else:
 47            for cmd in self.data:
 48                if cmd.entity_id == entity_id and cmd.command_type == CommandType.ADD_ENTITY:
 49                    projected = set(cmd.args["components"])
 50                    break
 51            else:
 52                assert False, "Invariant: live entity either must be in a pool (world) or staged (in self.data)"
 53
 54        for command in self.data:
 55            if command.entity_id == entity_id:
 56                if command.command_type == CommandType.ADD_COMPONENT:
 57                    projected.add(command.args["component"])
 58                elif command.command_type == CommandType.REMOVE_COMPONENT:
 59                    projected.discard(command.args["component"])
 60        return projected
 61
 62    def _get_component_in_buffer_state(self, entity_id: EntityId, component: ComponentType) -> bool | None:
 63        """Look for the latest state of this entity w.r.t this component given the unstaged command buffer.
 64        If the component is not in the buffer at all, return None. Return True if the last command is an ADD_COMPONENT
 65        Return False if the last command is an REMOVE_COMPONENT. Tri-state used to determine what we do with a new
 66        component called twice in the same tick before world.update()"""
 67        for old_command in reversed(self.data):
 68            if old_command.entity_id == entity_id and old_command.args.get("component", None) == component:
 69                return old_command.command_type == CommandType.ADD_COMPONENT # true if ADD, false if REMOVE
 70        return None
 71
 72    def append(self, command: Command):
 73        """Appends a command to the buffer"""
 74        world = self.world
 75        if (eid := command.entity_id) not in world.live_entities:
 76            raise ValueError(f"Entity: {eid} not in live entities ({command})")
 77
 78        # Below, we know for sure that the entity is either live or in staging (e.g. from ADD_ENTITY)
 79
 80        if command.command_type == CommandType.ADD_ENTITY:
 81            # nothing to do here: world.add_entity already ensures validated args (& defaults) come here.
 82            pass
 83
 84        elif command.command_type == CommandType.REMOVE_ENTITY:
 85            # needed so we can fast check in world.remove_enitity if this is a no-op (same tick) or error (stale eid).
 86            self.removed_this_tick.add(eid)
 87
 88        elif command.command_type == CommandType.ADD_COMPONENT:
 89            component = command.args["component"]
 90            fk = {k: v for k, v in command.args.items() if k != "component"}
 91            # same convert + validate + snapshot as add_entity (lists/tuples/scalars, None semantics)
 92            fk = world._validate_component(component, check_extra=True, **fk)
 93
 94            component_state = self._get_component_in_buffer_state(eid, component)
 95            if component_state is True:
 96                raise ValueError(f"Component: {component} already was already buffered this tick (id: {eid}). "
 97                                 f"Run world.update() or you may have a bug where 2 systems write without a lock in "
 98                                  "the same frame which may cause data loss on entities.")
 99
100            if component_state is None: # not staged, compare against projected components
101                projected = self._get_projected_entity_components(eid)
102                if component in projected:
103                    raise ValueError(f"Component: {component} already exists already, use e.set_data(). (id: {eid})")
104
105            # If we reach here either the component is not in the entity or we flip-flop in this frame (REMOVE->ADD)
106            command.args = {"component": component, **fk}
107
108        elif command.command_type == CommandType.REMOVE_COMPONENT:
109            component = command.args["component"]
110            if component not in world.component_types:
111                raise ValueError(f"Unknown component: {component} not in world components {world.component_types}")
112
113            component_state = self._get_component_in_buffer_state(eid, component)
114            if component_state is False: # repeated remove -> noop
115                return
116
117            projected = self._get_projected_entity_components(eid)
118            if component not in projected:
119                raise ValueError(f"{component=} is not part of the entity. Use e.add_component(). (id: {eid})")
120
121            if len(projected) == 1:
122                raise ValueError(f"Cannot remove the last component; use world.remove_entity({eid})")
123
124        self.data.append(command)
125
126    def __iter__(self):
127        return iter(self.data)
128
129    def __len__(self):
130        return len(self.data)
131
132    def __eq__(self, other: CommandBuffer | list[Command]):
133        if isinstance(other, list):
134            return self.data == other
135        elif isinstance(other, CommandBuffer):
136            return self.data == other.data
137        else:
138            return NotImplemented