sonolus.py 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

sonolus/script/archetype.py

@@ -10,8 +10,9 @@ from typing import Annotated, Any, ClassVar, Self, get_origin
 from sonolus.backend.ir import IRConst, IRInstr
 from sonolus.backend.mode import Mode
 from sonolus.backend.ops import Op
+from sonolus.backend.place import BlockPlace
 from sonolus.script.bucket import Bucket, Judgment
-from sonolus.script.callbacks import PLAY_CALLBACKS, PREVIEW_CALLBACKS, WATCH_ARCHETYPE_CALLBACKS, CallbackInfo
+from sonolus.script.internal.callbacks import PLAY_CALLBACKS, PREVIEW_CALLBACKS, WATCH_ARCHETYPE_CALLBACKS, CallbackInfo
 from sonolus.script.internal.context import ctx
 from sonolus.script.internal.descriptor import SonolusDescriptor
 from sonolus.script.internal.generic import validate_concrete_type
@@ -20,16 +21,16 @@ from sonolus.script.internal.introspection import get_field_specifiers
 from sonolus.script.internal.native import native_call
 from sonolus.script.internal.value import Value
 from sonolus.script.num import Num
-from sonolus.script.pointer import deref
+from sonolus.script.pointer import _deref
 from sonolus.script.record import Record
 from sonolus.script.values import zeros
 
-ENTITY_MEMORY_SIZE = 64
-ENTITY_DATA_SIZE = 32
-ENTITY_SHARED_MEMORY_SIZE = 32
+_ENTITY_MEMORY_SIZE = 64
+_ENTITY_DATA_SIZE = 32
+_ENTITY_SHARED_MEMORY_SIZE = 32
 
 
-class StorageType(Enum):
+class _StorageType(Enum):
     IMPORTED = "imported"
     EXPORTED = "exported"
     MEMORY = "memory"
@@ -37,53 +38,55 @@ class StorageType(Enum):
 
 
 @dataclass
-class ArchetypeFieldInfo:
+class _ArchetypeFieldInfo:
     name: str | None
-    storage: StorageType
+    storage: _StorageType
 
 
-class ArchetypeField(SonolusDescriptor):
-    def __init__(self, name: str, data_name: str, storage: StorageType, offset: int, type_: type[Value]):
+class _ArchetypeField(SonolusDescriptor):
+    def __init__(self, name: str, data_name: str, storage: _StorageType, offset: int, type_: type[Value]):
         self.name = name
         self.data_name = data_name  # name used in level data
         self.storage = storage
         self.offset = offset
         self.type = type_
 
-    def __get__(self, instance: BaseArchetype, owner):
+    def __get__(self, instance: _BaseArchetype, owner):
         if instance is None:
             return self
         result = None
         match self.storage:
-            case StorageType.IMPORTED:
+            case _StorageType.IMPORTED:
                 match instance._data_:
-                    case ArchetypeSelfData():
-                        result = deref(ctx().blocks.EntityData, self.offset, self.type)
-                    case ArchetypeReferenceData(index=index):
-                        result = deref(ctx().blocks.EntityDataArray, self.offset + index * ENTITY_DATA_SIZE, self.type)
-                    case ArchetypeLevelData(values=values):
+                    case _ArchetypeSelfData():
+                        result = _deref(ctx().blocks.EntityData, self.offset, self.type)
+                    case _ArchetypeReferenceData(index=index):
+                        result = _deref(
+                            ctx().blocks.EntityDataArray, self.offset + index * _ENTITY_DATA_SIZE, self.type
+                        )
+                    case _ArchetypeLevelData(values=values):
                         result = values[self.name]
-            case StorageType.EXPORTED:
+            case _StorageType.EXPORTED:
                 raise RuntimeError("Exported fields are write-only")
-            case StorageType.MEMORY:
+            case _StorageType.MEMORY:
                 match instance._data_:
-                    case ArchetypeSelfData():
-                        result = deref(ctx().blocks.EntityMemory, self.offset, self.type)
-                    case ArchetypeReferenceData():
+                    case _ArchetypeSelfData():
+                        result = _deref(ctx().blocks.EntityMemory, self.offset, self.type)
+                    case _ArchetypeReferenceData():
                         raise RuntimeError("Entity memory of other entities is not accessible")
-                    case ArchetypeLevelData():
+                    case _ArchetypeLevelData():
                         raise RuntimeError("Entity memory is not available in level data")
-            case StorageType.SHARED:
+            case _StorageType.SHARED:
                 match instance._data_:
-                    case ArchetypeSelfData():
-                        result = deref(ctx().blocks.EntitySharedMemory, self.offset, self.type)
-                    case ArchetypeReferenceData(index=index):
-                        result = deref(
+                    case _ArchetypeSelfData():
+                        result = _deref(ctx().blocks.EntitySharedMemory, self.offset, self.type)
+                    case _ArchetypeReferenceData(index=index):
+                        result = _deref(
                             ctx().blocks.EntitySharedMemoryArray,
-                            self.offset + index * ENTITY_SHARED_MEMORY_SIZE,
+                            Num._accept_(self.offset) + index * _ENTITY_SHARED_MEMORY_SIZE,
                             self.type,
                         )
-                    case ArchetypeLevelData():
+                    case _ArchetypeLevelData():
                         raise RuntimeError("Entity shared memory is not available in level data")
         if result is None:
             raise RuntimeError("Invalid storage type")
@@ -92,53 +95,55 @@ class ArchetypeField(SonolusDescriptor):
         else:
             return result._as_py_()
 
-    def __set__(self, instance: BaseArchetype, value):
+    def __set__(self, instance: _BaseArchetype, value):
         if instance is None:
             raise RuntimeError("Cannot set field on class")
         if not self.type._accepts_(value):
             raise TypeError(f"Expected {self.type}, got {type(value)}")
         target = None
         match self.storage:
-            case StorageType.IMPORTED:
+            case _StorageType.IMPORTED:
                 match instance._data_:
-                    case ArchetypeSelfData():
-                        target = deref(ctx().blocks.EntityData, self.offset, self.type)
-                    case ArchetypeReferenceData(index=index):
-                        target = deref(ctx().blocks.EntityDataArray, self.offset + index * ENTITY_DATA_SIZE, self.type)
-                    case ArchetypeLevelData(values=values):
+                    case _ArchetypeSelfData():
+                        target = _deref(ctx().blocks.EntityData, self.offset, self.type)
+                    case _ArchetypeReferenceData(index=index):
+                        target = _deref(
+                            ctx().blocks.EntityDataArray, self.offset + index * _ENTITY_DATA_SIZE, self.type
+                        )
+                    case _ArchetypeLevelData(values=values):
                         target = values[self.name]
-            case StorageType.EXPORTED:
+            case _StorageType.EXPORTED:
                 match instance._data_:
-                    case ArchetypeSelfData():
+                    case _ArchetypeSelfData():
                         if not isinstance(value, self.type):
                             raise TypeError(f"Expected {self.type}, got {type(value)}")
                         for k, v in value._to_flat_dict_(self.data_name).items():
                             index = instance._exported_keys_[k]
                             ctx().add_statements(IRInstr(Op.ExportValue, [IRConst(index), Num._accept_(v).ir()]))
                         return
-                    case ArchetypeReferenceData():
+                    case _ArchetypeReferenceData():
                         raise RuntimeError("Exported fields of other entities are not accessible")
-                    case ArchetypeLevelData():
+                    case _ArchetypeLevelData():
                         raise RuntimeError("Exported fields are not available in level data")
-            case StorageType.MEMORY:
+            case _StorageType.MEMORY:
                 match instance._data_:
-                    case ArchetypeSelfData():
-                        target = deref(ctx().blocks.EntityMemory, self.offset, self.type)
-                    case ArchetypeReferenceData():
+                    case _ArchetypeSelfData():
+                        target = _deref(ctx().blocks.EntityMemory, self.offset, self.type)
+                    case _ArchetypeReferenceData():
                         raise RuntimeError("Entity memory of other entities is not accessible")
-                    case ArchetypeLevelData():
+                    case _ArchetypeLevelData():
                         raise RuntimeError("Entity memory is not available in level data")
-            case StorageType.SHARED:
+            case _StorageType.SHARED:
                 match instance._data_:
-                    case ArchetypeSelfData():
-                        target = deref(ctx().blocks.EntitySharedMemory, self.offset, self.type)
-                    case ArchetypeReferenceData(index=index):
-                        target = deref(
+                    case _ArchetypeSelfData():
+                        target = _deref(ctx().blocks.EntitySharedMemory, self.offset, self.type)
+                    case _ArchetypeReferenceData(index=index):
+                        target = _deref(
                             ctx().blocks.EntitySharedMemoryArray,
-                            self.offset + index * ENTITY_SHARED_MEMORY_SIZE,
+                            Num._accept_(self.offset) + index * _ENTITY_SHARED_MEMORY_SIZE,
                             self.type,
                         )
-                    case ArchetypeLevelData():
+                    case _ArchetypeLevelData():
                         raise RuntimeError("Entity shared memory is not available in level data")
         if target is None:
             raise RuntimeError("Invalid storage type")
@@ -150,22 +155,75 @@ class ArchetypeField(SonolusDescriptor):
 
 
 def imported(*, name: str | None = None) -> Any:
-    return ArchetypeFieldInfo(name, StorageType.IMPORTED)
+    """Declare a field as imported.
+
+    Imported fields may be loaded from the level data.
+
+    In watch mode, data may also be loaded from a corresponding exported field in play mode.
+
+    Imported fields may only be updated in the `preprocess` callback, and are read-only in other callbacks.
+
+    Usage:
+        ```
+        class MyArchetype(PlayArchetype):
+            field: int = imported()
+            field_with_explicit_name: int = imported(name="field_name")
+        ```
+    """
+    return _ArchetypeFieldInfo(name, _StorageType.IMPORTED)
 
 
 def exported(*, name: str | None = None) -> Any:
-    return ArchetypeFieldInfo(name, StorageType.EXPORTED)
+    """Declare a field as exported.
+
+    This is only usable in play mode to export data to be loaded in watch mode.
+
+    Exported fields are write-only.
+
+    Usage:
+        ```
+        class MyArchetype(PlayArchetype):
+            field: int = exported()
+            field_with_explicit_name: int = exported(name="#FIELD")
+        ```
+    """
+    return _ArchetypeFieldInfo(name, _StorageType.EXPORTED)
 
 
 def entity_memory() -> Any:
-    return ArchetypeFieldInfo(None, StorageType.MEMORY)
+    """Declare a field as entity memory.
+
+    Entity memory is private to the entity and is not accessible from other entities.
+
+    Entity memory fields may also be set when an entity is spawned using the `spawn()` method.
+
+    Usage:
+        ```
+        class MyArchetype(PlayArchetype):
+            field: int = entity_memory()
+
+        ```
+    """
+    return _ArchetypeFieldInfo(None, _StorageType.MEMORY)
 
 
 def shared_memory() -> Any:
-    return ArchetypeFieldInfo(None, StorageType.SHARED)
+    """Declare a field as shared memory.
+
+    Shared memory is accessible from other entities.
+
+    Shared memory may only be updated by sequential callbacks such as `preprocess`, `update_sequential`, and `touch`.
+
+    Usage:
+        ```
+        class MyArchetype(PlayArchetype):
+            field: int = shared_memory()
+        ```
+    """
+    return _ArchetypeFieldInfo(None, _StorageType.SHARED)
 
 
-_annotation_defaults: dict[Callable, ArchetypeFieldInfo] = {
+_annotation_defaults: dict[Callable, _ArchetypeFieldInfo] = {
     imported: imported(),
     exported: exported(),
     entity_memory: entity_memory(),
@@ -174,14 +232,53 @@ _annotation_defaults: dict[Callable, ArchetypeFieldInfo] = {
 
 
 class StandardImport:
+    """Standard import annotations for Archetype fields.
+
+    Usage:
+        ```
+        class MyArchetype(WatchArchetype):
+            judgment: StandardImport.JUDGMENT
+        ```
+    """
+
     BEAT = Annotated[float, imported(name="#BEAT")]
+    """The beat of the entity."""
+
     BPM = Annotated[float, imported(name="#BPM")]
+    """The bpm, for bpm change markers."""
+
     TIMESCALE = Annotated[float, imported(name="#TIMESCALE")]
+    """The timescale, for timescale change markers."""
+
     JUDGMENT = Annotated[int, imported(name="#JUDGMENT")]
+    """The judgment of the entity.
+
+    Automatically supported in watch mode for archetypes with a corresponding scored play mode archetype.
+    """
     ACCURACY = Annotated[float, imported(name="#ACCURACY")]
+    """The accuracy of the entity.
+
+    Automatically supported in watch mode for archetypes with a corresponding scored play mode archetype.
+    """
+
+
+def callback[T: Callable](*, order: int = 0) -> Callable[[T], T]:
+    """Annotate a callback with its order.
+
+    Callbacks are execute from lowest to highest order. By default, callbacks have an order of 0.
+
+    Usage:
+        ```
+        class MyArchetype(PlayArchetype):
+            @callback(order=1)
+            def update_sequential(self):
+                pass
+        ```
 
+    Args:
+        order: The order of the callback. Lower values are executed first.
+    """
 
-def callback[T: Callable](order: int) -> Callable[[T], T]:
     def decorator(func: T) -> T:
         func._callback_order_ = order
         return func
@@ -189,37 +286,37 @@ def callback[T: Callable](order: int) -> Callable[[T], T]:
     return decorator
 
 
-class ArchetypeSelfData:
+class _ArchetypeSelfData:
     pass
 
 
-class ArchetypeReferenceData:
+class _ArchetypeReferenceData:
     index: Num
 
     def __init__(self, index: Num):
         self.index = index
 
 
-class ArchetypeLevelData:
+class _ArchetypeLevelData:
     values: dict[str, Value]
 
     def __init__(self, values: dict[str, Value]):
         self.values = values
 
 
-type ArchetypeData = ArchetypeSelfData | ArchetypeReferenceData | ArchetypeLevelData
+type _ArchetypeData = _ArchetypeSelfData | _ArchetypeReferenceData | _ArchetypeLevelData
 
 
-class BaseArchetype:
+class _BaseArchetype:
     _is_comptime_value_ = True
 
     _supported_callbacks_: ClassVar[dict[str, CallbackInfo]]
     _default_callbacks_: ClassVar[set[Callable]]
 
-    _imported_fields_: ClassVar[dict[str, ArchetypeField]]
-    _exported_fields_: ClassVar[dict[str, ArchetypeField]]
-    _memory_fields_: ClassVar[dict[str, ArchetypeField]]
-    _shared_memory_fields_: ClassVar[dict[str, ArchetypeField]]
+    _imported_fields_: ClassVar[dict[str, _ArchetypeField]]
+    _exported_fields_: ClassVar[dict[str, _ArchetypeField]]
+    _memory_fields_: ClassVar[dict[str, _ArchetypeField]]
+    _shared_memory_fields_: ClassVar[dict[str, _ArchetypeField]]
 
     _imported_keys_: ClassVar[dict[str, int]]
     _exported_keys_: ClassVar[dict[str, int]]
@@ -227,9 +324,16 @@ class BaseArchetype:
     _data_constructor_signature_: ClassVar[inspect.Signature]
     _spawn_signature_: ClassVar[inspect.Signature]
 
-    _data_: ArchetypeData
+    _data_: _ArchetypeData
 
     name: ClassVar[str | None] = None
+    """The name of the archetype.
+
+    If not set, the name will be the class name.
+
+    The name is used in level data.
+    """
+
     is_scored: ClassVar[bool] = False
 
     def __init__(self, *args, **kwargs):
@@ -241,7 +345,7 @@ class BaseArchetype:
             field.name: field.type._accept_(bound.arguments.get(field.name) or zeros(field.type))._get_()
             for field in self._imported_fields_.values()
         }
-        self._data_ = ArchetypeLevelData(values=values)
+        self._data_ = _ArchetypeLevelData(values=values)
 
     @classmethod
     def _new(cls):
@@ -250,14 +354,14 @@ class BaseArchetype:
     @classmethod
     def _for_compilation(cls):
         result = cls._new()
-        result._data_ = ArchetypeSelfData()
+        result._data_ = _ArchetypeSelfData()
         return result
 
     @classmethod
     @meta_fn
     def at(cls, index: Num) -> Self:
         result = cls._new()
-        result._data_ = ArchetypeReferenceData(index=index)
+        result._data_ = _ArchetypeReferenceData(index=Num._accept_(index))
         return result
 
     @classmethod
@@ -272,7 +376,21 @@ class BaseArchetype:
 
     @classmethod
     @meta_fn
-    def spawn(cls, **kwargs):
+    def spawn(cls, **kwargs: Any) -> None:
+        """Spawn an entity of this archetype, injecting the given values into entity memory.
+
+        Usage:
+            ```
+            class MyArchetype(PlayArchetype):
+                field: int = entity_memory()
+
+            def f():
+                MyArchetype.spawn(field=123)
+            ```
+
+        Args:
+            **kwargs: Entity memory values to inject by field name as defined in the Archetype.
+        """
         if not ctx():
             raise RuntimeError("Spawn is only allowed within a callback")
         archetype_id = cls.id()
@@ -283,18 +401,18 @@ class BaseArchetype:
             data.extend(field.type._accept_(bound.arguments[field.name] or zeros(field.type))._to_list_())
         native_call(Op.Spawn, archetype_id, *(Num(x) for x in data))
 
-    def _level_data_entries(self):
-        if not isinstance(self._data_, ArchetypeLevelData):
+    def _level_data_entries(self, level_refs: dict[Any, int] | None = None):
+        if not isinstance(self._data_, _ArchetypeLevelData):
             raise RuntimeError("Entity is not level data")
         entries = []
         for name, value in self._data_.values.items():
             field_info = self._imported_fields_.get(name)
-            for k, v in value._to_flat_dict_(field_info.data_name).items():
+            for k, v in value._to_flat_dict_(field_info.data_name, level_refs).items():
                 entries.append({"name": k, "value": v})
         return entries
 
     def __init_subclass__(cls, **kwargs):
-        if cls.__module__ == BaseArchetype.__module__:
+        if cls.__module__ == _BaseArchetype.__module__:
             if cls._supported_callbacks_ is None:
                 raise TypeError("Cannot directly subclass Archetype, use the Archetype subclass for your mode")
             cls._default_callbacks_ = {getattr(cls, cb_info.py_name) for cb_info in cls._supported_callbacks_.values()}
@@ -303,6 +421,7 @@ class BaseArchetype:
             raise TypeError("Cannot subclass Archetypes")
         if cls.name is None:
             cls.name = cls.__name__
+        field_specifiers = get_field_specifiers(cls, skip={"name", "is_scored"}).items()
         cls._imported_fields_ = {}
         cls._exported_fields_ = {}
         cls._memory_fields_ = {}
@@ -311,7 +430,7 @@ class BaseArchetype:
         exported_offset = 0
         memory_offset = 0
         shared_memory_offset = 0
-        for name, value in get_field_specifiers(cls).items():
+        for name, value in field_specifiers:
             if value is ClassVar or get_origin(value) is ClassVar:
                 continue
             if get_origin(value) is not Annotated:
@@ -322,7 +441,7 @@ class BaseArchetype:
             for metadata in value.__metadata__:
                 if isinstance(metadata, FunctionType):
                     metadata = _annotation_defaults.get(metadata, metadata)
-                if isinstance(metadata, ArchetypeFieldInfo):
+                if isinstance(metadata, _ArchetypeFieldInfo):
                     if field_info is not None:
                         raise TypeError(
                             f"Unexpected multiple field annotations for '{name}', "
@@ -336,26 +455,26 @@ class BaseArchetype:
                 )
             field_type = validate_concrete_type(value.__args__[0])
             match field_info.storage:
-                case StorageType.IMPORTED:
-                    cls._imported_fields_[name] = ArchetypeField(
+                case _StorageType.IMPORTED:
+                    cls._imported_fields_[name] = _ArchetypeField(
                         name, field_info.name or name, field_info.storage, imported_offset, field_type
                     )
                     imported_offset += field_type._size_()
                     setattr(cls, name, cls._imported_fields_[name])
-                case StorageType.EXPORTED:
-                    cls._exported_fields_[name] = ArchetypeField(
+                case _StorageType.EXPORTED:
+                    cls._exported_fields_[name] = _ArchetypeField(
                         name, field_info.name or name, field_info.storage, exported_offset, field_type
                     )
                     exported_offset += field_type._size_()
                     setattr(cls, name, cls._exported_fields_[name])
-                case StorageType.MEMORY:
-                    cls._memory_fields_[name] = ArchetypeField(
+                case _StorageType.MEMORY:
+                    cls._memory_fields_[name] = _ArchetypeField(
                         name, field_info.name or name, field_info.storage, memory_offset, field_type
                     )
                     memory_offset += field_type._size_()
                     setattr(cls, name, cls._memory_fields_[name])
-                case StorageType.SHARED:
-                    cls._shared_memory_fields_[name] = ArchetypeField(
+                case _StorageType.SHARED:
+                    cls._shared_memory_fields_[name] = _ArchetypeField(
                         name, field_info.name or name, field_info.storage, shared_memory_offset, field_type
                     )
                     shared_memory_offset += field_type._size_()
@@ -386,41 +505,97 @@ class BaseArchetype:
             cls._callbacks_.append(cb)
 
 
-class PlayArchetype(BaseArchetype):
+class PlayArchetype(_BaseArchetype):
+    """Base class for play mode archetypes.
+
+    Usage:
+        ```
+        class MyArchetype(PlayArchetype):
+            # Set to True if the entity is a note and contributes to combo and score
+            # Default is False
+            is_scored: bool = True
+
+            imported_field: int = imported()
+            exported_field: int = exported()
+            entity_memory_field: int = entity_memory()
+            shared_memory_field: int = shared_memory()
+
+            @callback(order=1)
+            def preprocess(self):
+                ...
+        ```
+    """
+
     _supported_callbacks_ = PLAY_CALLBACKS
 
+    is_scored: ClassVar[bool] = False
+    """Whether the entity contributes to combo and score."""
+
     def preprocess(self):
-        pass
+        """Perform upfront processing.
+
+        Runs first when the level is loaded.
+        """
 
     def spawn_order(self) -> float:
-        pass
+        """Return the spawn order of the entity.
+
+        Runs when the level is loaded after `preprocess`.
+        """
 
     def should_spawn(self) -> bool:
-        pass
+        """Return whether the entity should be spawned.
+
+        Runs when this entity is first in the spawn queue.
+        """
 
     def initialize(self):
-        pass
+        """Initialize this entity.
+
+        Runs when this entity is spawned.
+        """
 
     def update_sequential(self):
-        pass
+        """Perform non-parallel actions for this frame.
+
+        Runs first each frame.
+
+        This is where logic affecting shared memory should be placed.
+        Other logic should be placed in `update_parallel` for better performance.
+        """
 
     def update_parallel(self):
-        pass
+        """Perform parallel actions for this frame.
+
+        Runs after `touch` each frame.
+
+        This is where most gameplay logic should be placed.
+        """
 
     def touch(self):
-        pass
+        """Handle user input.
+
+        Runs after `update_sequential` each frame.
+        """
 
     def terminate(self):
-        pass
+        """Finalize before despawning.
+
+        Runs when the entity is despawned.
+        """
 
     @property
     @meta_fn
     def despawn(self):
+        """Whether the entity should be despawned after this frame.
+
+        Setting this to True will despawn the entity.
+        """
         if not ctx():
             raise RuntimeError("Calling despawn is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                return deref(ctx().blocks.EntityDespawn, 0, Num)
+            case _ArchetypeSelfData():
+                return _deref(ctx().blocks.EntityDespawn, 0, Num)
             case _:
                 raise RuntimeError("Despawn is only accessible from the entity itself")
 
@@ -430,8 +605,8 @@ class PlayArchetype(BaseArchetype):
         if not ctx():
             raise RuntimeError("Calling despawn is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                deref(ctx().blocks.EntityDespawn, 0, Num)._set_(value)
+            case _ArchetypeSelfData():
+                _deref(ctx().blocks.EntityDespawn, 0, Num)._set_(value)
             case _:
                 raise RuntimeError("Despawn is only accessible from the entity itself")
 
@@ -441,73 +616,128 @@ class PlayArchetype(BaseArchetype):
         if not ctx():
             raise RuntimeError("Calling info is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                return deref(ctx().blocks.EntityInfo, 0, PlayEntityInfo)
-            case ArchetypeReferenceData(index=index):
-                return deref(ctx().blocks.EntityInfoArray, index * PlayEntityInfo._size_(), PlayEntityInfo)
+            case _ArchetypeSelfData():
+                return _deref(ctx().blocks.EntityInfo, 0, PlayEntityInfo)
+            case _ArchetypeReferenceData(index=index):
+                return _deref(ctx().blocks.EntityInfoArray, index * PlayEntityInfo._size_(), PlayEntityInfo)
             case _:
                 raise RuntimeError("Info is only accessible from the entity itself")
 
     @property
     def index(self) -> int:
+        """The index of this entity."""
         return self._info.index
 
     @property
     def is_waiting(self) -> bool:
+        """Whether this entity is waiting to be spawned."""
         return self._info.state == 0
 
     @property
     def is_active(self) -> bool:
+        """Whether this entity is active."""
         return self._info.state == 1
 
     @property
     def is_despawned(self) -> bool:
+        """Whether this entity is despawned."""
         return self._info.state == 2
 
     @property
     def life(self) -> ArchetypeLife:
+        """How this entity contributes to life."""
         if not ctx():
             raise RuntimeError("Calling life is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData() | ArchetypeReferenceData():
-                return deref(ctx().blocks.ArchetypeLife, self.id() * ArchetypeLife._size_(), ArchetypeLife)
+            case _ArchetypeSelfData() | _ArchetypeReferenceData():
+                return _deref(ctx().blocks.ArchetypeLife, self.id() * ArchetypeLife._size_(), ArchetypeLife)
             case _:
                 raise RuntimeError("Life is not available in level data")
 
     @property
     def result(self) -> PlayEntityInput:
+        """The result of this entity.
+
+        Only meaningful for scored entities.
+        """
         if not ctx():
             raise RuntimeError("Calling result is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                return deref(ctx().blocks.EntityInput, 0, PlayEntityInput)
+            case _ArchetypeSelfData():
+                return _deref(ctx().blocks.EntityInput, 0, PlayEntityInput)
             case _:
                 raise RuntimeError("Result is only accessible from the entity itself")
 
+    def ref(self):
+        """Get a reference to this entity for creating level data.
+
+        Not valid elsewhere.
+        """
+        if not isinstance(self._data_, _ArchetypeLevelData):
+            raise RuntimeError("Entity is not level data")
+        result = EntityRef[type(self)](index=-1)
+        result._ref_ = self
+        return result
+
+
+class WatchArchetype(_BaseArchetype):
+    """Base class for watch mode archetypes.
+
+    Usage:
+        ```
+        class MyArchetype(WatchArchetype):
+            imported_field: int = imported()
+            entity_memory_field: int = entity_memory()
+            shared_memory_field: int = shared_memory()
+
+            @callback(order=1)
+            def update_sequential(self):
+                ...
+        ```
+    """
 
-class WatchArchetype(BaseArchetype):
     _supported_callbacks_ = WATCH_ARCHETYPE_CALLBACKS
 
     def preprocess(self):
-        pass
+        """Perform upfront processing.
+
+        Runs first when the level is loaded.
+        """
 
     def spawn_time(self) -> float:
-        pass
+        """Return the spawn time of the entity."""
 
     def despawn_time(self) -> float:
-        pass
+        """Return the despawn time of the entity."""
 
     def initialize(self):
-        pass
+        """Initialize this entity.
+
+        Runs when this entity is spawned.
+        """
 
     def update_sequential(self):
-        pass
+        """Perform non-parallel actions for this frame.
+
+        Runs first each frame.
+
+        This is where logic affecting shared memory should be placed.
+        Other logic should be placed in `update_parallel` for better performance.
+        """
 
     def update_parallel(self):
-        pass
+        """Parallel update callback.
+
+        Runs after `touch` each frame.
+
+        This is where most gameplay logic should be placed.
+        """
 
     def terminate(self):
-        pass
+        """Finalize before despawning.
+
+        Runs when the entity is despawned.
+        """
 
     @property
     @meta_fn
@@ -515,43 +745,56 @@ class WatchArchetype(BaseArchetype):
         if not ctx():
             raise RuntimeError("Calling info is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                return deref(ctx().blocks.EntityInfo, 0, WatchEntityInfo)
-            case ArchetypeReferenceData(index=index):
-                return deref(ctx().blocks.EntityInfoArray, index * WatchEntityInfo._size_(), PlayEntityInfo)
+            case _ArchetypeSelfData():
+                return _deref(ctx().blocks.EntityInfo, 0, WatchEntityInfo)
+            case _ArchetypeReferenceData(index=index):
+                return _deref(ctx().blocks.EntityInfoArray, index * WatchEntityInfo._size_(), PlayEntityInfo)
             case _:
                 raise RuntimeError("Info is only accessible from the entity itself")
 
     @property
     def index(self) -> int:
+        """The index of this entity."""
         return self._info.index
 
     @property
     def is_active(self) -> bool:
+        """Whether this entity is active."""
         return self._info.state == 1
 
     @property
     def life(self) -> ArchetypeLife:
+        """How this entity contributes to life."""
         if not ctx():
             raise RuntimeError("Calling life is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData() | ArchetypeReferenceData():
-                return deref(ctx().blocks.ArchetypeLife, self.id() * ArchetypeLife._size_(), ArchetypeLife)
+            case _ArchetypeSelfData() | _ArchetypeReferenceData():
+                return _deref(ctx().blocks.ArchetypeLife, self.id() * ArchetypeLife._size_(), ArchetypeLife)
             case _:
                 raise RuntimeError("Life is not available in level data")
 
     @property
     def result(self) -> WatchEntityInput:
+        """The result of this entity.
+
+        Only meaningful for scored entities.
+        """
         if not ctx():
             raise RuntimeError("Calling result is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                return deref(ctx().blocks.EntityInput, 0, WatchEntityInput)
+            case _ArchetypeSelfData():
+                return _deref(ctx().blocks.EntityInput, 0, WatchEntityInput)
             case _:
                 raise RuntimeError("Result is only accessible from the entity itself")
 
     @property
     def target_time(self) -> float:
+        """The target time of this entity.
+
+        Only meaningful for scored entities. Determines when combo and score are updated.
+
+        Alias of `result.target_time`.
+        """
         return self.result.target_time
 
     @target_time.setter
@@ -559,56 +802,86 @@ class WatchArchetype(BaseArchetype):
         self.result.target_time = value
 
 
-class PreviewArchetype(BaseArchetype):
+class PreviewArchetype(_BaseArchetype):
+    """Base class for preview mode archetypes.
+
+    Usage:
+        ```
+        class MyArchetype(PreviewArchetype):
+            imported_field: int = imported()
+            entity_memory_field: int = entity_memory()
+            shared_memory_field: int = shared_memory()
+
+            @callback(order=1)
+            def preprocess(self):
+                ...
+        ```
+    """
+
     _supported_callbacks_ = PREVIEW_CALLBACKS
 
     def preprocess(self):
-        pass
+        """Perform upfront processing.
+
+        Runs first when the level is loaded.
+        """
 
     def render(self):
-        pass
+        """Render the entity.
+
+        Runs after `preprocess`.
+        """
 
     @property
     def _info(self) -> PreviewEntityInfo:
         if not ctx():
             raise RuntimeError("Calling info is only allowed within a callback")
         match self._data_:
-            case ArchetypeSelfData():
-                return deref(ctx().blocks.EntityInfo, 0, PreviewEntityInfo)
-            case ArchetypeReferenceData(index=index):
-                return deref(ctx().blocks.EntityInfoArray, index * PreviewEntityInfo._size_(), PreviewEntityInfo)
+            case _ArchetypeSelfData():
+                return _deref(ctx().blocks.EntityInfo, 0, PreviewEntityInfo)
+            case _ArchetypeReferenceData(index=index):
+                return _deref(ctx().blocks.EntityInfoArray, index * PreviewEntityInfo._size_(), PreviewEntityInfo)
             case _:
                 raise RuntimeError("Info is only accessible from the entity itself")
 
     @property
     def index(self) -> int:
+        """The index of this entity."""
         return self._info.index
 
 
 @meta_fn
 def entity_info_at(index: Num) -> PlayEntityInfo | WatchEntityInfo | PreviewEntityInfo:
+    """Retrieve entity info of the entity at the given index.
+
+    Available in play, watch, and preview mode.
+    """
     if not ctx():
         raise RuntimeError("Calling entity_info_at is only allowed within a callback")
     match ctx().global_state.mode:
         case Mode.PLAY:
-            return deref(ctx().blocks.EntityInfoArray, index * PlayEntityInfo._size_(), PlayEntityInfo)
+            return _deref(ctx().blocks.EntityInfoArray, index * PlayEntityInfo._size_(), PlayEntityInfo)
         case Mode.WATCH:
-            return deref(ctx().blocks.EntityInfoArray, index * WatchEntityInfo._size_(), WatchEntityInfo)
+            return _deref(ctx().blocks.EntityInfoArray, index * WatchEntityInfo._size_(), WatchEntityInfo)
         case Mode.PREVIEW:
-            return deref(ctx().blocks.EntityInfoArray, index * PreviewEntityInfo._size_(), PreviewEntityInfo)
+            return _deref(ctx().blocks.EntityInfoArray, index * PreviewEntityInfo._size_(), PreviewEntityInfo)
         case _:
             raise RuntimeError(f"Entity info is not available in mode '{ctx().global_state.mode}'")
 
 
 @meta_fn
-def archetype_life_of(archetype: type[BaseArchetype] | BaseArchetype) -> ArchetypeLife:
+def archetype_life_of(archetype: type[_BaseArchetype] | _BaseArchetype) -> ArchetypeLife:
+    """Retrieve the archetype life of the given archetype.
+
+    Available in play and watch mode.
+    """
     archetype = validate_value(archetype)
     archetype = archetype._as_py_()
     if not ctx():
         raise RuntimeError("Calling archetype_life_of is only allowed within a callback")
     match ctx().global_state.mode:
         case Mode.PLAY | Mode.WATCH:
-            return deref(ctx().blocks.ArchetypeLife, archetype.id() * ArchetypeLife._size_(), ArchetypeLife)
+            return _deref(ctx().blocks.ArchetypeLife, archetype.id() * ArchetypeLife._size_(), ArchetypeLife)
         case _:
             raise RuntimeError(f"Archetype life is not available in mode '{ctx().global_state.mode}'")
 
@@ -631,10 +904,19 @@ class PreviewEntityInfo(Record):
 
 
 class ArchetypeLife(Record):
+    """How an entity contributes to life."""
+
     perfect_increment: Num
+    """Life increment for a perfect judgment."""
+
     great_increment: Num
+    """Life increment for a great judgment."""
+
     good_increment: Num
+    """Life increment for a good judgment."""
+
     miss_increment: Num
+    """Life increment for a miss judgment."""
 
     def update(
         self,
@@ -643,6 +925,7 @@ class ArchetypeLife(Record):
         good_increment: Num | None = None,
         miss_increment: Num | None = None,
     ):
+        """Update the life increments."""
         if perfect_increment is not None:
             self.perfect_increment = perfect_increment
         if great_increment is not None:
@@ -666,17 +949,33 @@ class WatchEntityInput(Record):
     bucket_value: float
 
 
-class EntityRef[A: BaseArchetype](Record):
+class EntityRef[A: _BaseArchetype](Record):
+    """Reference to another entity."""
+
     index: int
 
     @classmethod
     def archetype(cls) -> type[A]:
-        return cls._get_type_arg_(A)
+        return cls.type_var_value(A)
 
     def get(self) -> A:
-        return self.archetype().at(Num(self.index))
+        return self.archetype().at(self.index)
+
+    def _to_list_(self, level_refs: dict[Any, int] | None = None) -> list[float | BlockPlace]:
+        ref = getattr(self, "_ref_", None)
+        if ref is None:
+            return [self.index]
+        else:
+            if ref not in level_refs:
+                raise KeyError("Reference to entity not in level data")
+            return [level_refs[ref]]
 
 
 class StandardArchetypeName(StrEnum):
+    """Standard archetype names."""
+
     BPM_CHANGE = "#BPM_CHANGE"
+    """Bpm change marker"""
+
     TIMESCALE_CHANGE = "#TIMESCALE_CHANGE"
+    """Timescale change marker"""