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
  5from dataclasses import dataclass
  6from enum import StrEnum
  7
  8from .utils import EntityId
  9from .component import ComponentType
 10
 11class CommandType(StrEnum):
 12    """The types of commands in the command pattern below"""
 13    ADD_ENTITY       = "add_entity"
 14    REMOVE_ENTITY    = "remove_entity"
 15    ADD_COMPONENT    = "add_component"
 16    REMOVE_COMPONENT = "remove_component"
 17
 18@dataclass
 19class Command:
 20    """A command that can happen between two world.updates(), e.g. add/rm entity or components"""
 21    command_type: CommandType
 22    entity_id: EntityId
 23    args: Any | None = None
 24
 25class CommandBuffer:
 26    """A data structure that holds all the uncommited commands between two world updates. Support eager exceptions
 27       on things like adding the same component twice on the same entity"""
 28    def __init__(self, world: "World"): # noqa
 29        self.data: list[Command] = []
 30        self.world = world
 31        self.removed_this_tick: set[EntityId] = set()
 32
 33    def clear(self):
 34        """Clears the buffer"""
 35        self.data.clear()
 36        self.removed_this_tick.clear()
 37
 38    def _get_entity_components(self, entity_id: EntityId) -> list[ComponentType]: # noqa
 39        # This is the case for uncommited entities
 40        if entity_id not in self.world._eid_to_pool_ix:
 41            # uncommitted spawn: base = the components it was spawned with
 42            for cmd in self.data:
 43                if cmd.entity_id == entity_id and cmd.command_type == CommandType.ADD_ENTITY:
 44                    return cmd.args["components"]
 45            return [] # Entity should exist so this shouldn't be reached technically. We have an assert at call site.
 46        pool, _ = self.world._eid_to_pool_ix[entity_id]
 47        return self.world.pool_to_components[pool]
 48
 49    def _entity_has_buffered_component(self, entity_id: EntityId, component: ComponentType,
 50                                       existing_components: list[ComponentType]) -> bool:
 51        # Look for the latest state of this entity w.r.t this component given the unstaged command buffer.
 52        # We look in the buffer from right to left and return True if the component was added, False otherwise.
 53        # If the component is not in the buffer at all, we check if it is already in the entity. Returns a bool.
 54        for old_command in reversed(self.data):
 55            if old_command.entity_id != entity_id:
 56                continue
 57            if old_command.command_type == CommandType.ADD_COMPONENT:
 58                old_component = old_command.args["component"]
 59                if component == old_component:
 60                    return True
 61            if old_command.command_type == CommandType.REMOVE_COMPONENT:
 62                old_component = old_command.args
 63                if component == old_component:
 64                    return False
 65        return component in existing_components
 66
 67    def append(self, command: Command):
 68        """Appends a command to the buffer"""
 69        world = self.world
 70        entity_id = command.entity_id
 71        if entity_id not in world.live_entities:
 72            raise ValueError(f"Entity: {entity_id} not in live entities ({command})")
 73
 74        if command.command_type == CommandType.ADD_ENTITY:
 75            # nothing to do here: world.add_entity already ensures validated args (& defaults) come here.
 76            pass
 77
 78        elif command.command_type == CommandType.REMOVE_ENTITY:
 79            # needed so we can fast check in world.remove_enitity if this is a no-op (same tick) or error (stale eid).
 80            self.removed_this_tick.add(command.entity_id)
 81
 82        elif command.command_type == CommandType.ADD_COMPONENT:
 83            component = command.args["component"]
 84            fk = {k: v for k, v in command.args.items() if k != "component"}
 85            world._validate_component(component, strict=True, check_extra=True, **fk)
 86
 87            components = self._get_entity_components(entity_id)
 88            assert len(components) > 0, f"guaranteed to be >0 {entity_id} {components}"
 89            has_component = self._entity_has_buffered_component(entity_id, component, existing_components=components)
 90            if has_component:
 91                raise ValueError(f"Component: {component} either added twice or exists already (id: {entity_id})")
 92
 93        elif command.command_type == CommandType.REMOVE_COMPONENT:
 94            component = command.args # TODO: use command.args["component"] for consistency
 95            if component not in world.component_types:
 96                raise ValueError(f"Unknown component: {component} not in world components {world.component_types}")
 97
 98            components = self._get_entity_components(entity_id)
 99            assert len(components) > 0, f"guaranteed to be >0 {entity_id} {components}"
100            has_component = self._entity_has_buffered_component(entity_id, component, existing_components=components)
101            if not has_component:
102                raise ValueError(f"Component: {component} either removed twice or doesn't exist (id: {entity_id})")
103
104        self.data.append(command)
105
106    def __iter__(self):
107        return iter(self.data)
108
109    def __len__(self):
110        return len(self.data)
111
112    def __eq__(self, other: CommandBuffer | list[Command]):
113        if isinstance(other, list):
114            return self.data == other
115        elif isinstance(other, CommandBuffer):
116            return self.data == other.data
117        else:
118            return NotImplemented