fix armature duplication

6 files changed, 92 insertions, 46 deletions

diff --git a/doc/protocol.md b/doc/protocol.md
index 4c05cf2..a3b8364 100644
--- a/doc/protocol.md
+++ b/doc/protocol.md
@@ -25,7 +25,7 @@ encoded text.
 0x03 add(ob: Obj, prefab: Res<Prefab>)
 0x04 remove(ob: Obj)
 0x05 position(ob: Obj, pos: Vec3, rot: Vec3)
-0x06 pose(ob: Obj, params: Vec<f32>)
+0x06 pose(ob: Obj, params: Vec<(u16, Affine)>)
 0x07 parent(parent: Obj, child: Obj)
 0x08 sound(ob: Obj, data: Vec<u8>)
 0x09 prefab_index(res: Res)
diff --git a/doc/resources.md b/doc/resources.md
index 67fca4a..344101c 100644
--- a/doc/resources.md
+++ b/doc/resources.md
@@ -68,7 +68,7 @@ white except normals are zero.
 | `va_transmission`    | `Res<[Float]>`      |                    |
 | `va_emission`        | `Res<[Vec3]>`       |                    |
 | `va_joint_weight`    | `Res<[[Float; 4]]>` |                    |
-| `va_joint_index`     | `Res<[[u16; 4]]>`   |                    |
+| `va_joint_index`     | `Res<[[u32; 4]]>`   |                    |
 | `tex_normal`         | `Res<Texture>`      | Use color channels |
 | `tex_roughness`      | `Res<Texture>`      | Use green channel  |
 | `tex_metallic`       | `Res<Texture>`      | Use blue channel   |
diff --git a/shared/src/helper.rs b/shared/src/helper.rs
index 873c9d8..d6fbe66 100644
--- a/shared/src/helper.rs
+++ b/shared/src/helper.rs
@@ -180,6 +180,27 @@ impl ReadWrite for Vec<[u32; 3]> {
             .collect())
     }
 }
+impl ReadWrite for Vec<[u32; 4]> {
+    fn write(&self, w: &mut dyn Write) -> Result<()> {
+        for e in self {
+            w.write_all(&e[0].to_le_bytes())?;
+            w.write_all(&e[1].to_le_bytes())?;
+            w.write_all(&e[2].to_le_bytes())?;
+            w.write_all(&e[3].to_le_bytes())?;
+        }
+        Ok(())
+    }
+    fn read(r: &mut dyn Read) -> Result<Self> {
+        let mut buf = Vec::new();
+        r.read_to_end(&mut buf)?;
+        Ok(buf
+            .into_iter()
+            .array_chunks::<{ size_of::<u32>() }>()
+            .map(u32::from_le_bytes)
+            .array_chunks::<4>()
+            .collect())
+    }
+}
 impl ReadWrite for Vec<[u16; 4]> {
     fn write(&self, w: &mut dyn Write) -> Result<()> {
         for e in self {
diff --git a/shared/src/resources.rs b/shared/src/resources.rs
index 654bfc6..b3d9414 100644
--- a/shared/src/resources.rs
+++ b/shared/src/resources.rs
@@ -74,7 +74,7 @@ pub struct MeshPart {
     pub va_transmission: Option<Resource<Vec<f32>>>,
     pub va_alpha: Option<Resource<Vec<f32>>>,
     pub va_emission: Option<Resource<Vec<Vec3A>>>,
-    pub va_joint_index: Option<Resource<Vec<[u16; 4]>>>,
+    pub va_joint_index: Option<Resource<Vec<[u32; 4]>>>,
     pub va_joint_weight: Option<Resource<Vec<[f32; 4]>>>,
     pub tex_normal: Option<Resource<Image<'static>>>,
     pub tex_roughness: Option<Resource<Image<'static>>>,
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,