Split baked animation modules

src/fighter/baked_animation/asset_baking.rs

@@ -0,0 +1,456 @@
+use super::{
+    BAKED_ANIMATION_FPS,
+    fixed_transform::FixedMat4,
+    runtime::{BakedAnimationClip, BakedFighterAnimations},
+};
+use crate::fighter::{animation::AnimKind, manifest::FighterManifest};
+use bevy::{
+    asset::{
+        AssetLoader, AssetPath, AsyncWriteExt, LoadContext,
+        io::{Reader, Writer},
+        processor::LoadTransformAndSave,
+        saver::{AssetSaver, SavedAsset},
+        transformer::{AssetTransformer, TransformedAsset},
+    },
+    prelude::*,
+    reflect::TypePath,
+};
+use serde::Deserialize;
+use std::{collections::HashMap, io, path::Path};
+
+#[derive(Deserialize)]
+struct FighterBakeDeclaration {
+    manifest_path: String,
+}
+#[derive(Asset, TypePath)]
+struct FighterBakeSource {
+    model_bytes: Vec<u8>,
+    buffers: Vec<Vec<u8>>,
+    animations: HashMap<AnimKind, String>,
+}
+#[derive(Default, TypePath)]
+struct FighterBakeSourceLoader;
+
+impl AssetLoader for FighterBakeSourceLoader {
+    type Asset = FighterBakeSource;
+    type Settings = ();
+    type Error = io::Error;
+    async fn load(
+        &self,
+        reader: &mut dyn Reader,
+        _: &Self::Settings,
+        load_context: &mut LoadContext<'_>,
+    ) -> Result<Self::Asset, Self::Error> {
+        let mut declaration = Vec::new();
+        reader.read_to_end(&mut declaration).await?;
+        let declaration: FighterBakeDeclaration = ron::de::from_bytes(&declaration)
+            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
+        let manifest_bytes = load_context
+            .read_asset_bytes(&declaration.manifest_path)
+            .await
+            .map_err(|error| io::Error::other(error.to_string()))?;
+        let manifest: FighterManifest = ron::de::from_bytes(&manifest_bytes)
+            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
+        let model_bytes = load_context
+            .read_asset_bytes(&manifest.model_path)
+            .await
+            .map_err(|error| io::Error::other(error.to_string()))?;
+        let gltf = gltf::Gltf::from_slice(&model_bytes)
+            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
+        let mut buffers = Vec::with_capacity(gltf.buffers().len());
+        for buffer in gltf.buffers() {
+            let bytes = match buffer.source() {
+                gltf::buffer::Source::Uri(uri) => load_context
+                    .read_asset_bytes(AssetPath::from_path_buf(
+                        Path::new(&manifest.model_path)
+                            .parent()
+                            .unwrap_or_else(|| Path::new(""))
+                            .join(uri),
+                    ))
+                    .await
+                    .map_err(|error| io::Error::other(error.to_string()))?,
+                gltf::buffer::Source::Bin => gltf.blob.clone().ok_or_else(|| {
+                    io::Error::new(io::ErrorKind::InvalidData, "missing glTF binary blob")
+                })?,
+            };
+            buffers.push(bytes);
+        }
+        Ok(FighterBakeSource {
+            model_bytes,
+            buffers,
+            animations: manifest.animations,
+        })
+    }
+    fn extensions(&self) -> &[&str] {
+        &["fighterbake"]
+    }
+}
+
+#[derive(Default, TypePath)]
+struct BakedFighterAnimationsLoader;
+impl AssetLoader for BakedFighterAnimationsLoader {
+    type Asset = BakedFighterAnimations;
+    type Settings = ();
+    type Error = io::Error;
+    async fn load(
+        &self,
+        reader: &mut dyn Reader,
+        _: &Self::Settings,
+        _: &mut LoadContext<'_>,
+    ) -> Result<Self::Asset, Self::Error> {
+        let mut bytes = Vec::new();
+        reader.read_to_end(&mut bytes).await?;
+        bincode::deserialize(&bytes)
+            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
+    }
+    fn extensions(&self) -> &[&str] {
+        &["fighterbake"]
+    }
+}
+
+#[derive(Default, TypePath)]
+struct FighterBakeTransformer;
+impl AssetTransformer for FighterBakeTransformer {
+    type AssetInput = FighterBakeSource;
+    type AssetOutput = BakedFighterAnimations;
+    type Settings = ();
+    type Error = io::Error;
+    async fn transform<'a>(
+        &'a self,
+        source: TransformedAsset<Self::AssetInput>,
+        _: &'a Self::Settings,
+    ) -> Result<TransformedAsset<Self::AssetOutput>, Self::Error> {
+        let baked = bake_gltf(&source.model_bytes, &source.buffers, &source.animations)?;
+        Ok(source.replace_asset(baked))
+    }
+}
+
+#[derive(Default, TypePath)]
+struct BakedFighterAnimationsSaver;
+impl AssetSaver for BakedFighterAnimationsSaver {
+    type Asset = BakedFighterAnimations;
+    type Settings = ();
+    type OutputLoader = BakedFighterAnimationsLoader;
+    type Error = io::Error;
+    async fn save(
+        &self,
+        writer: &mut Writer,
+        asset: SavedAsset<'_, '_, Self::Asset>,
+        _: &Self::Settings,
+        _: AssetPath<'_>,
+    ) -> Result<(), Self::Error> {
+        let bytes = bincode::serialize(&*asset)
+            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
+        writer.write_all(&bytes).await
+    }
+}
+type FighterBakeProcessor = LoadTransformAndSave<
+    FighterBakeSourceLoader,
+    FighterBakeTransformer,
+    BakedFighterAnimationsSaver,
+>;
+
+pub struct BakedAnimationPlugin;
+impl Plugin for BakedAnimationPlugin {
+    fn build(&self, app: &mut App) {
+        app.init_asset::<BakedFighterAnimations>()
+            .register_asset_loader(FighterBakeSourceLoader)
+            .register_asset_loader(BakedFighterAnimationsLoader)
+            .register_asset_processor(FighterBakeProcessor::new(
+                FighterBakeTransformer,
+                BakedFighterAnimationsSaver,
+            ))
+            .set_default_asset_processor::<FighterBakeProcessor>("fighterbake")
+            .add_systems(
+                bevy_ggrs::prelude::GgrsSchedule,
+                super::runtime::update_fighter_bone_matrices
+                    .in_set(crate::schedule::GameplaySet::Animation),
+            )
+            .add_systems(
+                Update,
+                super::visual_binding::apply_baked_fighter_visual_pose,
+            )
+            .add_observer(super::visual_binding::bind_fighter_visual_bones);
+    }
+}
+
+fn bake_gltf(
+    bytes: &[u8],
+    buffer_bytes: &[Vec<u8>],
+    animations: &HashMap<AnimKind, String>,
+) -> io::Result<BakedFighterAnimations> {
+    let gltf = gltf::Gltf::from_slice(bytes)
+        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
+    let buffers = gltf
+        .buffers()
+        .zip(buffer_bytes)
+        .map(|(buffer, bytes)| {
+            if bytes.len() < buffer.length() {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!(
+                        "glTF buffer {} is too short: expected {}, got {}",
+                        buffer.index(),
+                        buffer.length(),
+                        bytes.len()
+                    ),
+                ));
+            }
+            Ok(gltf::buffer::Data(bytes.clone()))
+        })
+        .collect::<io::Result<Vec<_>>>()?;
+    if buffers.len() != gltf.buffers().len() {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "glTF buffer data is incomplete",
+        ));
+    }
+    bake_gltf_document(gltf.document, buffers, animations)
+}
+
+pub(super) fn bake_gltf_document(
+    document: gltf::Document,
+    buffers: Vec<gltf::buffer::Data>,
+    animations: &HashMap<AnimKind, String>,
+) -> io::Result<BakedFighterAnimations> {
+    let nodes: Vec<_> = document.nodes().collect();
+    if nodes.is_empty() {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "fighter model has no nodes",
+        ));
+    }
+    let mut parents = vec![None; nodes.len()];
+    for node in &nodes {
+        for child in node.children() {
+            parents[child.index()] = Some(node.index());
+        }
+    }
+    let names = nodes
+        .iter()
+        .map(|node| {
+            node.name().map(str::to_owned).ok_or_else(|| {
+                io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    format!("node {} has no name", node.index()),
+                )
+            })
+        })
+        .collect::<io::Result<Vec<_>>>()?;
+    let mut pending = HashMap::<&str, Vec<AnimKind>>::new();
+    for (kind, name) in animations {
+        pending.entry(name).or_default().push(*kind);
+    }
+    let mut clips = Vec::with_capacity(pending.len());
+    let mut clip_indices = HashMap::with_capacity(animations.len());
+    for animation in document.animations() {
+        let name = animation.name().ok_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("animation {} has no name", animation.index()),
+            )
+        })?;
+        let Some(kinds) = pending.remove(name) else {
+            continue;
+        };
+        let channels = animation
+            .channels()
+            .map(|channel| read_channel(channel, &buffers))
+            .collect::<io::Result<Vec<_>>>()?;
+        let duration = channels
+            .iter()
+            .filter_map(Channel::duration)
+            .fold(0.0_f32, f32::max);
+        let frame_count = baked_frame_count(duration, BAKED_ANIMATION_FPS);
+        let mut frames = Vec::with_capacity(frame_count);
+        for frame in 0..frame_count {
+            let time = (frame as f32 / BAKED_ANIMATION_FPS as f32).min(duration);
+            frames.push(
+                nodes
+                    .iter()
+                    .map(|node| local_transform(node.clone(), &channels, time))
+                    .collect::<io::Result<Vec<_>>>()?,
+            );
+        }
+        let index = clips.len();
+        clips.push(BakedAnimationClip { frames });
+        for kind in kinds {
+            clip_indices.insert(kind, index);
+        }
+    }
+    if clips.is_empty() {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "fighter model has no animations",
+        ));
+    }
+    if !pending.is_empty() {
+        let mut missing = pending.keys().copied().collect::<Vec<_>>();
+        missing.sort_unstable();
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!(
+                "fighter model has no animations named {}",
+                missing.join(", ")
+            ),
+        ));
+    }
+    Ok(BakedFighterAnimations {
+        frame_rate: BAKED_ANIMATION_FPS,
+        bone_names: names,
+        parents,
+        clips,
+        clip_indices,
+    })
+}
+
+pub(super) fn baked_frame_count(duration: f32, frame_rate: u32) -> usize {
+    let frame = duration as f64 * frame_rate as f64;
+    let nearest_frame = frame.round();
+    let last_frame = if (frame - nearest_frame).abs() < 0.0001 {
+        nearest_frame
+    } else {
+        frame.ceil()
+    };
+    last_frame as usize + 1
+}
+
+#[derive(Clone)]
+pub(super) struct Channel {
+    pub(super) node: usize,
+    pub(super) property: gltf::animation::Property,
+    pub(super) interpolation: gltf::animation::Interpolation,
+    pub(super) times: Vec<f32>,
+    pub(super) values: ChannelValues,
+}
+impl Channel {
+    fn duration(&self) -> Option<f32> {
+        self.times.last().copied()
+    }
+}
+#[derive(Clone)]
+pub(super) enum ChannelValues {
+    Vec3(Vec<[f32; 3]>),
+    Rotation(Vec<[f32; 4]>),
+}
+
+fn read_channel(
+    channel: gltf::animation::Channel<'_>,
+    buffers: &[gltf::buffer::Data],
+) -> io::Result<Channel> {
+    let reader = channel.reader(|buffer| Some(&buffers[buffer.index()].0));
+    let times = reader
+        .read_inputs()
+        .ok_or_else(|| {
+            io::Error::new(io::ErrorKind::InvalidData, "animation channel has no input")
+        })?
+        .collect();
+    let values = match reader.read_outputs().ok_or_else(|| {
+        io::Error::new(
+            io::ErrorKind::InvalidData,
+            "animation channel has no output",
+        )
+    })? {
+        gltf::animation::util::ReadOutputs::Translations(values)
+        | gltf::animation::util::ReadOutputs::Scales(values) => {
+            ChannelValues::Vec3(values.collect())
+        }
+        gltf::animation::util::ReadOutputs::Rotations(values) => {
+            ChannelValues::Rotation(values.into_f32().collect())
+        }
+        gltf::animation::util::ReadOutputs::MorphTargetWeights(_) => {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "morph target animation is not supported for bone baking",
+            ));
+        }
+    };
+    if channel.sampler().interpolation() == gltf::animation::Interpolation::CubicSpline {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "cubic spline animation is not supported for bone baking",
+        ));
+    }
+    Ok(Channel {
+        node: channel.target().node().index(),
+        property: channel.target().property(),
+        interpolation: channel.sampler().interpolation(),
+        times,
+        values,
+    })
+}
+
+fn local_transform(node: gltf::Node<'_>, channels: &[Channel], time: f32) -> io::Result<FixedMat4> {
+    let (mut translation, mut rotation, mut scale) = match node.transform() {
+        gltf::scene::Transform::Decomposed {
+            translation,
+            rotation,
+            scale,
+        } => (
+            Vec3::from(translation),
+            Quat::from_array(rotation),
+            Vec3::from(scale),
+        ),
+        gltf::scene::Transform::Matrix { matrix } => {
+            return FixedMat4::from_mat4(Mat4::from_cols_array_2d(&matrix));
+        }
+    };
+    for channel in channels
+        .iter()
+        .filter(|channel| channel.node == node.index())
+    {
+        match (&channel.values, channel.property) {
+            (ChannelValues::Vec3(values), gltf::animation::Property::Translation) => {
+                translation = Vec3::from(sample_vec3(channel, values, time))
+            }
+            (ChannelValues::Vec3(values), gltf::animation::Property::Scale) => {
+                scale = Vec3::from(sample_vec3(channel, values, time))
+            }
+            (ChannelValues::Rotation(values), gltf::animation::Property::Rotation) => {
+                rotation = Quat::from_array(sample_rotation(channel, values, time))
+            }
+            _ => {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    "animation property does not match its values",
+                ));
+            }
+        }
+    }
+    FixedMat4::from_mat4(Mat4::from_scale_rotation_translation(
+        scale,
+        rotation,
+        translation,
+    ))
+}
+fn sample_index(channel: &Channel, time: f32) -> (usize, usize, f32) {
+    let next = channel.times.partition_point(|key| *key <= time);
+    if next == 0 {
+        return (0, 0, 0.0);
+    }
+    if next == channel.times.len() {
+        let last = next - 1;
+        return (last, last, 0.0);
+    }
+    let previous = next - 1;
+    let factor = (time - channel.times[previous]) / (channel.times[next] - channel.times[previous]);
+    (previous, next, factor)
+}
+pub(super) fn sample_vec3(channel: &Channel, values: &[[f32; 3]], time: f32) -> [f32; 3] {
+    let (from, to, factor) = sample_index(channel, time);
+    if channel.interpolation == gltf::animation::Interpolation::Step {
+        return values[from];
+    }
+    Vec3::from(values[from])
+        .lerp(Vec3::from(values[to]), factor)
+        .to_array()
+}
+fn sample_rotation(channel: &Channel, values: &[[f32; 4]], time: f32) -> [f32; 4] {
+    let (from, to, factor) = sample_index(channel, time);
+    if channel.interpolation == gltf::animation::Interpolation::Step {
+        return values[from];
+    }
+    Quat::from_array(values[from])
+        .slerp(Quat::from_array(values[to]), factor)
+        .to_array()
+}