1 files changed, 52 insertions, 40 deletions

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();