MetaCrate/crates/libremetaverse/src/animesh.rs

//! Deterministic animation playback and object-animation asset coordination.

#![allow(clippy::needless_pass_by_value, clippy::unnecessary_wraps)] // Mapped constructors are fixed.

use crate::animation::{BinBVHAnimationReader, BinBVHJointKey};
use crate::client_core::ClientWeakHandle;
use crate::packet_catalog::PacketType;
use crate::packets::ObjectAnimationPacket;
use crate::{Error, GridClient};
use libremetaverse_types::compat::{CancellationTokenSource, Subscription};
use libremetaverse_types::{AssetType, Quaternion, UUID, Vector3};
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, RwLock};

const MAX_TRACKS_PER_PLAYER: usize = 64;
const MAX_PLAYERS: usize = 4096;
const MAX_ANIMATIONS_PER_UPDATE: usize = 64;
const MAX_CONCURRENT_FETCHES: usize = 16;

fn mutex<T>(lock: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
    lock.lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

fn read<T>(lock: &RwLock<T>) -> std::sync::RwLockReadGuard<'_, T> {
    lock.read()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

fn write<T>(lock: &RwLock<T>) -> std::sync::RwLockWriteGuard<'_, T> {
    lock.write()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

#[derive(Clone, Copy, Debug)]
pub struct JointPose {
    pub ease_weight: f32,
    pub has_position: bool,
    pub has_rotation: bool,
    pub position: Vector3,
    pub priority: i32,
    pub rotation: Quaternion,
}

#[derive(Default)]
struct TrackState {
    data: Option<BinBVHAnimationReader>,
    current_time: f32,
    is_finished: bool,
}

pub struct AnimationTrack {
    animation_id: UUID,
    state: Mutex<TrackState>,
}

impl AnimationTrack {
    pub(crate) fn native_new(animation_id: UUID) -> Result<Self, Error> {
        Ok(Self {
            animation_id,
            state: Mutex::new(TrackState::default()),
        })
    }

    pub(crate) fn native_animation_id(&self) -> UUID {
        self.animation_id
    }

    pub(crate) fn native_data(&self) -> Option<BinBVHAnimationReader> {
        mutex(&self.state).data.clone()
    }

    pub(crate) fn native_set_data(&self, data: Option<BinBVHAnimationReader>) {
        let mut state = mutex(&self.state);
        state.current_time = 0.0;
        state.is_finished = false;
        state.data = data;
    }

    pub(crate) fn native_current_time(&self) -> f32 {
        mutex(&self.state).current_time
    }

    pub(crate) fn native_set_current_time(&self, value: f32) {
        if value.is_finite() {
            mutex(&self.state).current_time = value;
        }
    }

    pub(crate) fn native_is_finished(&self) -> bool {
        mutex(&self.state).is_finished
    }

    pub(crate) fn native_set_is_finished(&self, value: bool) {
        mutex(&self.state).is_finished = value;
    }

    pub(crate) fn native_advance(&self, dt: f32) -> Result<(), Error> {
        if !dt.is_finite() || dt < 0.0 {
            return Err(Error::Argument);
        }
        let mut state = mutex(&self.state);
        if state.is_finished || state.data.is_none() {
            return Ok(());
        }
        let data = state.data.as_ref().expect("checked animation data");
        let (looping, in_point, out_point) = (data.loop_, data.in_point, data.out_point);
        state.current_time += dt;
        if looping {
            let span = out_point - in_point;
            if span > 0.0 && state.current_time > out_point {
                state.current_time = in_point + (state.current_time - in_point) % span;
            }
        } else if state.current_time >= out_point {
            state.current_time = out_point;
            state.is_finished = true;
        }
        Ok(())
    }

    pub(crate) fn native_ease_weight(&self) -> f32 {
        let state = mutex(&self.state);
        ease_weight(&state)
    }

    pub(crate) fn native_evaluate_pose(
        &self,
        pose: &mut HashMap<String, JointPose>,
    ) -> Result<(), Error> {
        let state = mutex(&self.state);
        let Some(data) = &state.data else {
            return Ok(());
        };
        let ease = ease_weight(&state);
        for joint in data.joints.iter().take(MAX_TRACKS_PER_PLAYER * 8) {
            let rotation = (!joint.rotationkeys.is_empty())
                .then(|| interpolate_rotation(&joint.rotationkeys, state.current_time))
                .transpose()?;
            let position = (!joint.positionkeys.is_empty())
                .then(|| interpolate_position(&joint.positionkeys, state.current_time))
                .transpose()?;
            if rotation.is_none() && position.is_none() {
                continue;
            }
            let (rotation, position) = if let Some(existing) = pose.get(&joint.name).copied() {
                if joint.priority < existing.priority {
                    continue;
                }
                if joint.priority == existing.priority {
                    let total = existing.ease_weight + ease;
                    let weight = if total > 0.0 { ease / total } else { 0.5 };
                    (
                        match (rotation, existing.has_rotation) {
                            (Some(value), true) => {
                                Some(Quaternion::slerp(existing.rotation, value, weight)?)
                            }
                            (Some(value), false) => Some(value),
                            (None, true) => Some(existing.rotation),
                            (None, false) => None,
                        },
                        match (position, existing.has_position) {
                            (Some(value), true) => {
                                Some(Vector3::lerp(existing.position, value, weight)?)
                            }
                            (Some(value), false) => Some(value),
                            (None, true) => Some(existing.position),
                            (None, false) => None,
                        },
                    )
                } else {
                    (rotation, position)
                }
            } else {
                (rotation, position)
            };
            pose.insert(
                joint.name.clone(),
                JointPose {
                    rotation: rotation.unwrap_or_else(Quaternion::identity),
                    has_rotation: rotation.is_some(),
                    position: position.unwrap_or_else(Vector3::zero),
                    has_position: position.is_some(),
                    priority: joint.priority,
                    ease_weight: ease,
                },
            );
        }
        Ok(())
    }

    pub(crate) fn native_decode_rotation(value: Vector3) -> Result<Quaternion, Error> {
        let w_squared = 1.0 - value.x * value.x - value.y * value.y - value.z * value.z;
        Quaternion::normalize(Quaternion {
            x: value.x,
            y: value.y,
            z: value.z,
            w: w_squared.max(0.0).sqrt(),
        })
    }
}

fn ease_weight(state: &TrackState) -> f32 {
    let Some(data) = &state.data else { return 0.0 };
    if data.ease_in_time > 0.0 && state.current_time < data.ease_in_time {
        return (state.current_time / data.ease_in_time).clamp(0.0, 1.0);
    }
    let start = data.out_point - data.ease_out_time;
    if data.ease_out_time > 0.0 && state.current_time > start {
        return ((data.out_point - state.current_time) / data.ease_out_time).clamp(0.0, 1.0);
    }
    1.0
}

fn bracket(keys: &[BinBVHJointKey], time: f32) -> usize {
    keys.partition_point(|key| key.time < time)
}

fn interpolate_rotation(keys: &[BinBVHJointKey], time: f32) -> Result<Quaternion, Error> {
    let upper = bracket(keys, time);
    if upper == 0 {
        return AnimationTrack::native_decode_rotation(keys[0].key_element);
    }
    if upper >= keys.len() {
        return AnimationTrack::native_decode_rotation(keys[keys.len() - 1].key_element);
    }
    let (left, right) = (&keys[upper - 1], &keys[upper]);
    let fraction = if right.time > left.time {
        (time - left.time) / (right.time - left.time)
    } else {
        0.0
    };
    Quaternion::slerp(
        AnimationTrack::native_decode_rotation(left.key_element)?,
        AnimationTrack::native_decode_rotation(right.key_element)?,
        fraction,
    )
}

fn interpolate_position(keys: &[BinBVHJointKey], time: f32) -> Result<Vector3, Error> {
    let upper = bracket(keys, time);
    if upper == 0 {
        return Ok(keys[0].key_element);
    }
    if upper >= keys.len() {
        return Ok(keys[keys.len() - 1].key_element);
    }
    let (left, right) = (&keys[upper - 1], &keys[upper]);
    let fraction = if right.time > left.time {
        (time - left.time) / (right.time - left.time)
    } else {
        0.0
    };
    Vector3::lerp(left.key_element, right.key_element, fraction)
}

struct PlayerInner {
    object_id: UUID,
    tracks: RwLock<HashMap<UUID, Arc<AnimationTrack>>>,
}

#[derive(Clone)]
pub struct AnimeshPlayer(Arc<PlayerInner>);

impl AnimeshPlayer {
    pub(crate) fn native_new(object_id: UUID) -> Result<Self, Error> {
        Ok(Self(Arc::new(PlayerInner {
            object_id,
            tracks: RwLock::new(HashMap::new()),
        })))
    }

    pub(crate) fn native_object_id(&self) -> UUID {
        self.0.object_id
    }

    pub(crate) fn native_get_or_add_track(
        &self,
        animation_id: UUID,
    ) -> Result<Arc<AnimationTrack>, Error> {
        if let Some(track) = read(&self.0.tracks).get(&animation_id).cloned() {
            return Ok(track);
        }
        let mut tracks = write(&self.0.tracks);
        if let Some(track) = tracks.get(&animation_id).cloned() {
            return Ok(track);
        }
        if tracks.len() >= MAX_TRACKS_PER_PLAYER {
            return Err(Error::InvalidOperation);
        }
        let track = Arc::new(AnimationTrack::native_new(animation_id)?);
        tracks.insert(animation_id, Arc::clone(&track));
        Ok(track)
    }

    fn native_track(&self, animation_id: UUID) -> Option<Arc<AnimationTrack>> {
        read(&self.0.tracks).get(&animation_id).cloned()
    }

    pub(crate) fn native_retain_only(&self, active: &HashSet<UUID>) {
        write(&self.0.tracks).retain(|id, _| active.contains(id));
    }

    pub(crate) fn native_track_count(&self) -> i32 {
        i32::try_from(read(&self.0.tracks).len()).unwrap_or(i32::MAX)
    }

    pub(crate) fn native_update(&self, dt: f32) -> Result<(), Error> {
        let mut tracks: Vec<_> = read(&self.0.tracks).values().cloned().collect();
        tracks.sort_by_key(|track| track.native_animation_id().to_string());
        for track in tracks {
            track.native_advance(dt)?;
        }
        Ok(())
    }

    pub(crate) fn native_evaluate_pose(&self) -> Result<HashMap<String, JointPose>, Error> {
        let mut tracks: Vec<_> = read(&self.0.tracks).values().cloned().collect();
        tracks.sort_by_key(|track| track.native_animation_id().to_string());
        let mut pose = HashMap::new();
        for track in tracks {
            track.native_evaluate_pose(&mut pose)?;
        }
        Ok(pose)
    }
}

pub(crate) struct AnimeshManagerInner {
    client: ClientWeakHandle,
    players: RwLock<HashMap<UUID, AnimeshPlayer>>,
    cancellation: CancellationTokenSource,
    fetches_in_flight: Arc<AtomicUsize>,
    subscriptions: Mutex<Vec<Subscription>>,
}

impl Drop for AnimeshManagerInner {
    fn drop(&mut self) {
        self.cancellation.cancel();
    }
}

#[derive(Clone)]
pub struct AnimeshManager(Arc<AnimeshManagerInner>);

impl AnimeshManager {
    pub(crate) fn native_new(client: GridClient) -> Result<Self, Error> {
        let inner = Arc::new(AnimeshManagerInner {
            client: client.native_weak_handle(),
            players: RwLock::new(HashMap::new()),
            cancellation: CancellationTokenSource::new(),
            fetches_in_flight: Arc::new(AtomicUsize::new(0)),
            subscriptions: Mutex::new(Vec::new()),
        });
        let weak = Arc::downgrade(&inner);
        let subscription = client
            .native_network()?
            .subscribe_raw_packet(Arc::new(move |event| {
                if event.packet_type != PacketType::ObjectAnimation {
                    return;
                }
                let Some(inner) = weak.upgrade() else { return };
                handle_object_animation(&inner, event.data);
            }));
        mutex(&inner.subscriptions).push(subscription);
        Ok(Self(inner))
    }

    pub(crate) fn native_inner(&self) -> Arc<AnimeshManagerInner> {
        Arc::clone(&self.0)
    }

    pub(crate) fn native_from_inner(inner: Arc<AnimeshManagerInner>) -> Self {
        Self(inner)
    }

    pub(crate) fn native_get_player(&self, object_id: UUID) -> Option<AnimeshPlayer> {
        read(&self.0.players).get(&object_id).cloned()
    }

    pub(crate) fn native_all_players(&self) -> Vec<AnimeshPlayer> {
        read(&self.0.players).values().cloned().collect()
    }

    pub(crate) fn native_remove_player(&self, object_id: UUID) {
        write(&self.0.players).remove(&object_id);
    }

    pub(crate) fn native_update(&self, dt: f32) -> Result<(), Error> {
        let players = self.native_all_players();
        for player in players {
            player.native_update(dt)?;
        }
        Ok(())
    }
}

fn handle_object_animation(inner: &Arc<AnimeshManagerInner>, data: Vec<u8>) {
    let mut offset = 0;
    let Ok(packet) = ObjectAnimationPacket::new_with_bytes_int32(data, &mut offset) else {
        return;
    };
    if packet.animation_list.len() > MAX_ANIMATIONS_PER_UPDATE {
        return;
    }
    let object_id = packet.sender.id;
    let player = {
        let mut players = write(&inner.players);
        if let Some(player) = players.get(&object_id).cloned() {
            player
        } else {
            if players.len() >= MAX_PLAYERS {
                return;
            }
            let Ok(player) = AnimeshPlayer::native_new(object_id) else {
                return;
            };
            players.insert(object_id, player.clone());
            player
        }
    };
    let active: HashSet<_> = packet
        .animation_list
        .iter()
        .map(|value| value.anim_id)
        .collect();
    player.native_retain_only(&active);
    for animation in packet.animation_list {
        let Ok(track) = player.native_get_or_add_track(animation.anim_id) else {
            continue;
        };
        if track.native_data().is_none() {
            fetch_animation(inner, player.clone(), animation.anim_id);
        }
    }
}

fn fetch_animation(inner: &Arc<AnimeshManagerInner>, player: AnimeshPlayer, animation_id: UUID) {
    if inner
        .fetches_in_flight
        .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
            (count < MAX_CONCURRENT_FETCHES).then_some(count + 1)
        })
        .is_err()
    {
        return;
    }
    let weak = Arc::downgrade(inner);
    let fetches_in_flight = Arc::clone(&inner.fetches_in_flight);
    let spawned = std::thread::Builder::new()
        .name("animesh-asset".to_owned())
        .spawn(move || {
            struct FetchGuard(Arc<AtomicUsize>);
            impl Drop for FetchGuard {
                fn drop(&mut self) {
                    self.0.fetch_sub(1, Ordering::AcqRel);
                }
            }
            let _guard = FetchGuard(fetches_in_flight);
            let Ok(runtime) = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
            else {
                return;
            };
            runtime.block_on(async move {
                let (token, assets) = {
                    let Some(inner) = weak.upgrade() else { return };
                    let token = inner.cancellation.token();
                    let Some(client) = inner.client.upgrade() else {
                        return;
                    };
                    let Ok(assets) = client.native_assets() else {
                        return;
                    };
                    (token, assets)
                };
                let Ok(Some(asset)) = assets
                    .request_asset_with_uuid_asset_type_boolean_cancellation_token(
                        animation_id,
                        AssetType::Animation,
                        false,
                        Some(token),
                    )
                    .await
                else {
                    return;
                };
                let Ok(decoded) = BinBVHAnimationReader::native_new(asset.asset_data) else {
                    return;
                };
                if let Some(track) = player.native_track(animation_id) {
                    track.native_set_data(Some(decoded));
                }
            });
        });
    if spawned.is_err() {
        inner.fetches_in_flight.fetch_sub(1, Ordering::AcqRel);
    }
}