From d7fdcfd791f011e33fb9720e7451e56410411e4f Mon Sep 17 00:00:00 2001
From: metamuffin <metamuffin@disroot.org>
Date: Fri, 7 Feb 2025 15:25:00 +0100
Subject: fix armature duplication

---
 world/src/main.rs | 92 +++++++++++++++++++++++++++++++------------------------
 world/src/mesh.rs | 19 ++++++++++--
 2 files changed, 68 insertions(+), 43 deletions(-)

(limited to 'world/src')

diff --git a/world/src/main.rs b/world/src/main.rs
index 67747c8..37ea88a 100644
--- a/world/src/main.rs
+++ b/world/src/main.rs
@@ -31,7 +31,7 @@ use rand::random;
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use std::{
     borrow::Cow,
-    collections::HashMap,
+    collections::{BTreeMap, HashMap},
     fs::File,
     io::{Cursor, Read, Write},
     marker::PhantomData,
@@ -126,6 +126,55 @@ fn main() -> Result<()> {
             traverse(&mut nodes, node, root_affine);
         }
 
+        let mut skin_index_to_arm_index = BTreeMap::new();
+        let armature = {
+            let mut joint_index_to_arm_index = BTreeMap::new();
+            let mut name = Vec::new();
+            let mut parent_pre_map = Vec::new();
+            let mut transform = Vec::new();
+
+            for skin in gltf.skins() {
+                for (j_ind, j) in skin.joints().enumerate() {
+                    let a_ind = match joint_index_to_arm_index.get(&j.index()) {
+                        Some(i) => *i,
+                        None => {
+                            let a_ind = name.len();
+                            name.push(j.name().unwrap_or("").to_string());
+                            transform.push(transform_to_affine(j.transform()));
+                            parent_pre_map.push(
+                                gltf.nodes()
+                                    .find(|n| n.children().any(|c| c.index() == j.index()))
+                                    .map(|n| n.index()),
+                            );
+
+                            joint_index_to_arm_index.insert(j.index(), a_ind);
+                            a_ind
+                        }
+                    };
+                    skin_index_to_arm_index.insert((skin.index(), j_ind as u16), a_ind as u32);
+                }
+            }
+
+            let parent = parent_pre_map
+                .into_iter()
+                .enumerate()
+                .map(|(i, p)| {
+                    p.and_then(|i| joint_index_to_arm_index.get(&i).copied())
+                        .unwrap_or_else(|| i) as u16
+                })
+                .collect::<Vec<_>>();
+
+            if !parent.is_empty() {
+                Some(store.set(&ArmaturePart {
+                    name: Some(name),
+                    parent: Some(parent),
+                    transform: Some(transform),
+                })?)
+            } else {
+                None
+            }
+        };
+
         let mut prefab = nodes
             .par_iter()
             .map(|(trans, node)| {
@@ -142,6 +191,7 @@ fn main() -> Result<()> {
                         &mut prefab,
                         &args,
                         &texture_cache,
+                        &skin_index_to_arm_index,
                     )?;
                 }
                 let (position, _, _) = node.transform().decomposed();
@@ -182,45 +232,7 @@ fn main() -> Result<()> {
                 },
             )?;
 
-        prefab.armature = gltf
-            .skins()
-            .map(|skin| {
-                if let Some(root) = skin.skeleton() {
-                    fn traverse(
-                        ars: &mut (&mut Vec<String>, &mut Vec<Affine3A>, &mut Vec<u16>),
-                        trans: Affine3A,
-                        parent: Option<u16>,
-                        node: &Node,
-                    ) {
-                        let trans = trans * transform_to_affine(node.transform());
-                        let ind = ars.0.len() as u16;
-                        ars.0.push(node.name().unwrap_or("").to_owned());
-                        ars.1.push(trans);
-                        ars.2.push(parent.unwrap_or(ind));
-                        for c in node.children() {
-                            traverse(ars, trans, Some(ind), &c);
-                        }
-                    }
-                    let mut name = Vec::new();
-                    let mut transform = Vec::new();
-                    let mut parent = Vec::new();
-                    traverse(
-                        &mut (&mut name, &mut transform, &mut parent),
-                        Affine3A::IDENTITY,
-                        None,
-                        &root,
-                    );
-                    let armature = ArmaturePart {
-                        name: Some(name),
-                        parent: Some(parent),
-                        transform: Some(transform),
-                    };
-                    store.set(&armature)
-                } else {
-                    store.set(&ArmaturePart::default())
-                }
-            })
-            .collect::<Result<Vec<_>>>()?;
+        prefab.armature = armature.into_iter().collect();
 
         if let Some(skybox) = &args.skybox {
             let mut buf = Vec::new();
diff --git a/world/src/mesh.rs b/world/src/mesh.rs
index 96976ce..fa5dc4d 100644
--- a/world/src/mesh.rs
+++ b/world/src/mesh.rs
@@ -18,7 +18,10 @@ use crate::{Args, TextureCache, load_texture};
 use anyhow::Result;
 use gltf::{Mesh, Node, buffer::Data};
 use log::{debug, info, warn};
-use std::path::Path;
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    path::Path,
+};
 use weareshared::{
     Affine3A, Vec3A,
     resources::{MeshPart, Prefab},
@@ -36,6 +39,7 @@ pub fn import_mesh(
     prefab: &mut Prefab,
     args: &Args,
     texture_cache: &TextureCache,
+    joint_index_map: &BTreeMap<(usize, u16), u32>,
 ) -> Result<()> {
     for p in mesh.primitives() {
         let name = mesh.name().or(node.name()).map(|e| e.to_owned());
@@ -81,8 +85,17 @@ pub fn import_mesh(
         let va_joint_index = reader
             .read_joints(0)
             .map(|iter| {
-                let a = iter.into_u16().collect::<Vec<_>>();
+                let si = node.skin().unwrap().index();
+                let a = iter
+                    .into_u16()
+                    .map(|x| x.map(|x| joint_index_map[&(si, x)]))
+                    .collect::<Vec<_>>();
                 debug!("{} vertex joint indecies", a.len());
+                eprintln!(
+                    "index {:?} {:?}",
+                    node.name(),
+                    a.iter().flatten().collect::<BTreeSet<_>>()
+                );
                 if a.len() != num_vertex {
                     warn!("joint index count does not vertex count")
                 }
@@ -343,7 +356,7 @@ pub fn import_mesh(
             None
         };
 
-        let armature = node.skin().map(|s| s.index() as u32);
+        let armature = node.skin().map(|s| 0);
 
         let mesh = store.set(&MeshPart {
             name,
-- 
cgit v1.2.3-70-g09d2