steel_registry/entity_data/

mod.rs

//! Entity data synchronization system.
//!
//! This module provides the infrastructure for syncing entity state (health, pose, flags, etc.)
//! between server and client via the `SetEntityData` packet.
//!
//! # Architecture
//!
//! - [`SyncedValue<T>`] - Wrapper that tracks per-field dirty state
//! - [`EntityData`] - Type-erased enum for network serialization
//! - [`DataValue`] - Network-ready value with index and serializer info
//! - [`EntityPose`] - Entity pose states (standing, sneaking, swimming, etc.)
//! - [`EntityDataSerializerRegistry`] - Registry of serializers with writers
//!
//! # Serialization Flow
//!
//! Serialization happens through the registry, not directly on `EntityData`:
//! ```ignore
//! let writer = REGISTRY.entity_data_serializers.get_writer(serializer_id).unwrap();
//! writer(&entity_data_value, &mut buf)?;
//! ```
//!
//! # Generated Code
//!
//! Per-entity data structs (e.g., `PlayerEntityData`) are generated by the build script
//! from `entities.json` and live in `steel-registry/src/generated/entity_data.rs`.

mod serializer;
mod vanilla_serializers;

pub use serializer::{
    EntityDataSerializerEntry, EntityDataSerializerEntryRef, EntityDataSerializerRegistry,
    EntityDataWriter,
};
pub use vanilla_serializers::register_vanilla_entity_data_serializers;

use std::{io, str::FromStr};

use steel_utils::{
    BlockStateId, Identifier,
    codec::VarInt,
    serial::{ReadFrom, WriteTo},
};
use text_components::TextComponent;
use uuid::Uuid;

use crate::RegistryReference;
use crate::cat_sound_variant::CatSoundVariant;
use crate::cat_variant::CatVariant;
use crate::chicken_sound_variant::ChickenSoundVariant;
use crate::chicken_variant::ChickenVariant;
use crate::cow_sound_variant::CowSoundVariant;
use crate::cow_variant::CowVariant;
use crate::dye_color::DyeColor;
use crate::frog_variant::FrogVariant;
use crate::item_stack::ItemStack;
use crate::painting_variant::PaintingVariant;
pub use crate::particle_type::{ColorParticleOption, ParticleData};
use crate::pig_sound_variant::PigSoundVariant;
use crate::pig_variant::PigVariant;
use crate::wolf_sound_variant::WolfSoundVariant;
use crate::wolf_variant::WolfVariant;
use crate::zombie_nautilus_variant::ZombieNautilusVariant;

// Re-export types used in generated code
pub use crate::blocks::properties::Direction;
pub use steel_utils::BlockPos;

/// Wrapper that tracks modifications per-field.
///
/// Each field in an entity data struct is wrapped in `SyncedValue` to track
/// whether it has been modified and needs to be synced to clients.
#[derive(Debug, Clone)]
pub struct SyncedValue<T> {
    value: T,
    default: T,
    dirty: bool,
}

impl<T: Clone + PartialEq> SyncedValue<T> {
    /// Create a new synced value with the given default.
    pub fn new(default: T) -> Self {
        Self {
            value: default.clone(),
            default,
            dirty: false,
        }
    }

    /// Get a reference to the current value.
    #[inline]
    pub const fn get(&self) -> &T {
        &self.value
    }

    /// Set the value. Only marks as dirty if the value actually changed.
    #[inline]
    pub fn set(&mut self, value: T) {
        if self.value != value {
            self.value = value;
            self.dirty = true;
        }
    }

    /// Returns true if the value has been modified since last sync.
    #[inline]
    pub const fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Clear the dirty flag after syncing.
    #[inline]
    pub const fn clear_dirty(&mut self) {
        self.dirty = false;
    }

    /// Returns true if the current value equals the default.
    #[inline]
    pub fn is_default(&self) -> bool {
        self.value == self.default
    }
}

/// Entity pose states.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[repr(u8)]
pub enum EntityPose {
    #[default]
    Standing = 0,
    FallFlying = 1,
    Sleeping = 2,
    Swimming = 3,
    SpinAttack = 4,
    Sneaking = 5,
    LongJumping = 6,
    Dying = 7,
    Croaking = 8,
    UsingTongue = 9,
    Sitting = 10,
    Roaring = 11,
    Sniffing = 12,
    Emerging = 13,
    Digging = 14,
    Sliding = 15,
    Shooting = 16,
    Inhaling = 17,
}

impl FromStr for EntityPose {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "STANDING" => Ok(Self::Standing),
            "FALL_FLYING" => Ok(Self::FallFlying),
            "SLEEPING" => Ok(Self::Sleeping),
            "SWIMMING" => Ok(Self::Swimming),
            "SPIN_ATTACK" => Ok(Self::SpinAttack),
            "SNEAKING" | "CROUCHING" => Ok(Self::Sneaking),
            "LONG_JUMPING" => Ok(Self::LongJumping),
            "DYING" => Ok(Self::Dying),
            "CROAKING" => Ok(Self::Croaking),
            "USING_TONGUE" => Ok(Self::UsingTongue),
            "SITTING" => Ok(Self::Sitting),
            "ROARING" => Ok(Self::Roaring),
            "SNIFFING" => Ok(Self::Sniffing),
            "EMERGING" => Ok(Self::Emerging),
            "DIGGING" => Ok(Self::Digging),
            "SLIDING" => Ok(Self::Sliding),
            "SHOOTING" => Ok(Self::Shooting),
            "INHALING" => Ok(Self::Inhaling),
            _ => Err(()),
        }
    }
}

/// Rotations for armor stands and similar entities (pitch, yaw, roll in degrees).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Rotations {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

impl Rotations {
    pub const ZERO: Self = Self {
        x: 0.0,
        y: 0.0,
        z: 0.0,
    };

    #[must_use]
    pub const fn new(x: f32, y: f32, z: f32) -> Self {
        Self { x, y, z }
    }
}

/// Main hand preference (left or right).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[repr(u8)]
pub enum HumanoidArm {
    Left = 0,
    #[default]
    Right = 1,
}

impl ReadFrom for HumanoidArm {
    fn read(data: &mut std::io::Cursor<&[u8]>) -> io::Result<Self> {
        match VarInt::read(data)?.0 {
            0 => Ok(Self::Left),
            1 => Ok(Self::Right),
            value => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("invalid humanoid arm id {value}"),
            )),
        }
    }
}

/// Villager profession, type, and level data.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct VillagerData {
    /// Villager type (biome variant) - registry ID.
    pub villager_type: i32,
    /// Profession - registry ID.
    pub profession: i32,
    /// Trading level (1-5).
    pub level: i32,
}

impl VillagerData {
    #[must_use]
    pub const fn new(villager_type: i32, profession: i32, level: i32) -> Self {
        Self {
            villager_type,
            profession,
            level,
        }
    }
}

/// Selects a villager profession by numeric registry order, preserving the fallback when empty.
pub(crate) fn random_villager_profession_id(
    random: &mut impl steel_utils::random::Random,
    profession_count: usize,
    fallback: i32,
) -> i32 {
    if profession_count == 0 {
        return fallback;
    }

    let Ok(profession_count) = i32::try_from(profession_count) else {
        panic!("villager profession registry exceeds Vanilla's signed integer range");
    };
    random.next_i32_bounded(profession_count)
}

/// A global position (dimension + block position).
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct GlobalPos {
    pub dimension: Identifier,
    pub pos: BlockPos,
}

impl GlobalPos {
    #[must_use]
    pub const fn new(dimension: Identifier, pos: BlockPos) -> Self {
        Self { dimension, pos }
    }
}

/// A 3D vector (for display entities).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Vector3f {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

impl Vector3f {
    pub const ZERO: Self = Self {
        x: 0.0,
        y: 0.0,
        z: 0.0,
    };

    #[must_use]
    pub const fn new(x: f32, y: f32, z: f32) -> Self {
        Self { x, y, z }
    }
}

/// A quaternion rotation (for display entities).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Quaternionf {
    pub x: f32,
    pub y: f32,
    pub z: f32,
    pub w: f32,
}

impl Quaternionf {
    pub const IDENTITY: Self = Self {
        x: 0.0,
        y: 0.0,
        z: 0.0,
        w: 1.0,
    };

    #[must_use]
    pub const fn new(x: f32, y: f32, z: f32, w: f32) -> Self {
        Self { x, y, z, w }
    }
}

/// A list of particle effects.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ParticleList {
    pub particles: Vec<ParticleData>,
}

pub use crate::resolvable_profile::ResolvableProfile;

/// Sniffer entity state.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[repr(u8)]
pub enum SnifferState {
    #[default]
    Idling = 0,
    FeelingHappy = 1,
    Scenting = 2,
    Sniffing = 3,
    Searching = 4,
    Digging = 5,
    Rising = 6,
}

/// Armadillo entity state.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[repr(u8)]
pub enum ArmadilloState {
    #[default]
    Idle = 0,
    RollingUp = 1,
    Rolled = 2,
    Unrolling = 3,
    Scared = 4,
}

/// Type-erased entity data for network encoding.
///
/// This enum holds the actual value to be sent over the network.
/// Serialization happens through the [`EntityDataSerializerRegistry`], not directly
/// on this type. Use `REGISTRY.entity_data_serializers.get_writer(id)` to get the
/// writer function for a specific serializer ID.
#[derive(Debug, Clone)]
pub enum EntityData {
    // Primitive types
    Byte(i8),
    Int(i32),
    Long(i64),
    Float(f32),
    String(String),
    Boolean(bool),

    // Text components
    Component(Box<TextComponent>),
    OptionalComponent(Option<Box<TextComponent>>),

    // Item
    ItemStack(ItemStack),

    // Spatial types
    Rotations(Rotations),
    BlockPos(BlockPos),
    OptionalBlockPos(Option<BlockPos>),
    Direction(Direction),
    OptionalLivingEntityRef(Option<Uuid>),
    BlockState(BlockStateId),
    OptionalBlockState(Option<BlockStateId>),
    Vector3(Vector3f),
    Quaternion(Quaternionf),

    // Entity states
    Pose(EntityPose),
    SnifferState(SnifferState),
    ArmadilloState(ArmadilloState),

    // Particles
    Particle(ParticleData),
    Particles(ParticleList),

    // Villager
    VillagerData(VillagerData),

    // Optional numeric
    OptionalUnsignedInt(Option<u32>),

    // Holder/registry reference variants (encoded as VarInt registry IDs)
    CatVariant(RegistryReference<CatVariant>),
    CatSoundVariant(RegistryReference<CatSoundVariant>),
    CowVariant(RegistryReference<CowVariant>),
    CowSoundVariant(RegistryReference<CowSoundVariant>),
    WolfVariant(RegistryReference<WolfVariant>),
    WolfSoundVariant(RegistryReference<WolfSoundVariant>),
    FrogVariant(RegistryReference<FrogVariant>),
    PigVariant(RegistryReference<PigVariant>),
    PigSoundVariant(RegistryReference<PigSoundVariant>),
    ChickenVariant(RegistryReference<ChickenVariant>),
    ChickenSoundVariant(RegistryReference<ChickenSoundVariant>),
    ZombieNautilusVariant(RegistryReference<ZombieNautilusVariant>),
    PaintingVariant(RegistryReference<PaintingVariant>),

    // Global position
    OptionalGlobalPos(Option<GlobalPos>),

    // Copper golem / weathering states (future entities)
    CopperGolemState(i32),
    WeatheringCopperState(i32),

    // Profile
    ResolvableProfile(ResolvableProfile),

    // Player
    HumanoidArm(HumanoidArm),

    // Dye color
    DyeColor(DyeColor),
}

/// Network-ready data value with index and serializer info.
///
/// This is what gets written to the `SetEntityData` packet.
#[derive(Debug, Clone)]
pub struct DataValue {
    /// The index of this data field (0-254, 255 is terminator).
    pub index: u8,
    /// The serializer ID (from registration order in `EntityDataSerializerRegistry`).
    pub serializer_id: i32,
    /// The actual value to write.
    pub value: EntityData,
}

impl DataValue {
    /// Write this data value to the buffer (without terminator).
    ///
    /// Uses the global registry to look up the writer function for this serializer.
    pub fn write_to(&self, buf: &mut Vec<u8>) -> io::Result<()> {
        use crate::REGISTRY;

        self.index.write(buf)?;
        VarInt(self.serializer_id).write(buf)?;

        let writer = REGISTRY
            .entity_data_serializers
            .get_writer(self.serializer_id)
            .ok_or_else(|| {
                io::Error::other(format!(
                    "Unknown entity data serializer ID: {}",
                    self.serializer_id
                ))
            })?;

        writer(&self.value, buf)
    }
}

/// Write a list of data values to a buffer, with the 0xFF terminator.
pub fn write_data_values(values: &[DataValue], buf: &mut Vec<u8>) -> io::Result<()> {
    for value in values {
        value.write_to(buf)?;
    }
    // Write terminator
    0xFFu8.write(buf)
}

#[cfg(test)]
mod tests {
    use steel_utils::{
        ArgbColor, Identifier,
        codec::VarInt,
        random::{Random as _, legacy_random::LegacyRandom},
        serial::WriteTo,
    };

    use crate::vanilla_entity_data::{EggEntityData, ItemEntityData, ZombieVillagerEntityData};
    use crate::{REGISTRY, RegistryExt, init_vanilla_registry};

    use super::{
        ColorParticleOption, EntityData, EntityDataSerializerRegistry, ParticleData,
        random_villager_profession_id, register_vanilla_entity_data_serializers,
    };

    #[test]
    fn zombie_villager_default_selects_profession_from_registry() {
        init_vanilla_registry();

        let profession_count = REGISTRY.villager_professions.len();
        let Ok(profession_bound) = i32::try_from(profession_count) else {
            panic!("test villager profession registry must fit in i32");
        };
        let Some(nitwit_id) = REGISTRY
            .villager_professions
            .id_from_key(&Identifier::vanilla_static("nitwit"))
        else {
            panic!("nitwit villager profession must be registered");
        };
        let Some(plains_id) = REGISTRY
            .villager_types
            .id_from_key(&Identifier::vanilla_static("plains"))
        else {
            panic!("plains villager type must be registered");
        };

        let mut expected_random = LegacyRandom::from_seed(25);
        let expected_profession = expected_random.next_i32_bounded(profession_bound);
        assert_eq!(usize::try_from(expected_profession), Ok(nitwit_id));

        let mut random = LegacyRandom::from_seed(25);
        let data = ZombieVillagerEntityData::new(&mut random);
        let villager_data = data.villager_data.get();

        assert_eq!(usize::try_from(villager_data.villager_type), Ok(plains_id));
        assert_eq!(villager_data.profession, expected_profession);
        assert_eq!(villager_data.level, 1);
        assert!(data.pack_all().is_empty());
    }

    #[test]
    fn empty_villager_profession_registry_preserves_fallback_without_rng_draw() {
        let mut random = LegacyRandom::from_seed(25);
        let seed_before = random.get_seed();

        assert_eq!(random_villager_profession_id(&mut random, 0, 7), 7);
        assert_eq!(random.get_seed(), seed_before);
    }

    #[test]
    fn projectile_item_stack_defaults_use_extracted_item() {
        let data = EggEntityData::new();
        let stack = data.throwable_item_projectile().item_stack.get();

        assert_eq!(&stack.item().key, &Identifier::vanilla_static("egg"));
        assert_eq!(stack.count(), 1);
        assert!(data.pack_all().is_empty());
    }

    #[test]
    fn empty_item_stack_defaults_remain_empty() {
        let data = ItemEntityData::new();

        assert!(data.item.get().is_empty());
        assert!(data.pack_all().is_empty());
    }

    #[test]
    fn entity_effect_particle_color_options_encode_payload() {
        crate::init_vanilla_registry();
        let registry = &*crate::REGISTRY;
        let entity_effect = Identifier::vanilla_static("entity_effect");
        let Some(particle_type_id) = registry.particle_types.id_from_key(&entity_effect) else {
            panic!("entity_effect particle type must be registered");
        };

        let mut serializers = EntityDataSerializerRegistry::new();
        register_vanilla_entity_data_serializers(&mut serializers);

        let Some(serializer_id) = serializers.id_from_key(&Identifier::vanilla_static("particle"))
        else {
            panic!("particle entity-data serializer must be registered");
        };

        let Some(writer) = serializers.get_writer(serializer_id as i32) else {
            panic!("particle entity-data serializer must have a writer");
        };

        let particle = ParticleData::new(
            &crate::vanilla_particle_types::ENTITY_EFFECT,
            ColorParticleOption::new(ArgbColor::new(-1)),
        );
        let value = EntityData::Particle(particle);
        let mut encoded = Vec::new();
        let result = writer(&value, &mut encoded);
        assert!(result.is_ok(), "{result:?}");

        let mut expected = Vec::new();
        let result = VarInt(particle_type_id as i32).write(&mut expected);
        assert!(result.is_ok(), "{result:?}");
        let result = (-1i32).write(&mut expected);
        assert!(result.is_ok(), "{result:?}");

        assert_eq!(encoded, expected);
    }
}