feat: Fix combineSkeletons

packages/three-vrm/src/VRMUtils/combineSkeletons.ts

@@ -12,13 +12,21 @@ export function combineSkeletons(root: THREE.Object3D): void {
   const skinnedMeshes = collectSkinnedMeshes(root);
 
   // List all used skin indices for each skin index attribute
-  const attributeUsedIndexSetMap = new Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, Set<number>>();
+  /** A map: skin index attribute -> skin weight attribute -> used index set */
+  const attributeUsedIndexSetMap = new Map<
+    THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+    Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, Set<number>>
+  >();
   for (const mesh of skinnedMeshes) {
     const geometry = mesh.geometry;
+
     const skinIndexAttr = geometry.getAttribute('skinIndex');
+    const skinIndexMap = attributeUsedIndexSetMap.get(skinIndexAttr) ?? new Map();
+    attributeUsedIndexSetMap.set(skinIndexAttr, skinIndexMap);
+
     const skinWeightAttr = geometry.getAttribute('skinWeight');
     const usedIndicesSet = listUsedIndices(skinIndexAttr, skinWeightAttr);
-    attributeUsedIndexSetMap.set(skinIndexAttr, usedIndicesSet);
+    skinIndexMap.set(skinWeightAttr, usedIndicesSet);
   }
 
   // List all bones and boneInverses for each meshes
@@ -63,18 +71,16 @@ export function combineSkeletons(root: THREE.Object3D): void {
     const newBoneInverses = Array.from(boneInverseMap.values());
     const newSkeleton = new THREE.Skeleton(newBones, newBoneInverses);
 
-    const attributeProcessedSet = new Set<THREE.BufferAttribute | THREE.InterleavedBufferAttribute>();
-
-    for (const mesh of meshes) {
-      const attribute = mesh.geometry.getAttribute('skinIndex');
+    // collect skin index attributes and corresponding bone arrays
+    const skinIndexBonesPairSet = collectSkinIndexAttrs(meshes);
 
-      if (!attributeProcessedSet.has(attribute)) {
-        // remap skin index attribute
-        remapSkinIndexAttribute(attribute, mesh.skeleton.bones, newBones);
-        attributeProcessedSet.add(attribute);
-      }
+    // remap skin index attribute
+    for (const [skinIndexAttr, bones] of skinIndexBonesPairSet) {
+      remapSkinIndexAttribute(skinIndexAttr, bones, newBones);
+    }
 
-      // bind the new skeleton to the mesh
+    // bind the new skeleton to the meshes
+    for (const mesh of meshes) {
       mesh.bind(newSkeleton, new THREE.Matrix4());
     }
   }
@@ -132,18 +138,25 @@ function listUsedIndices(
  */
 function listUsedBones(
   mesh: THREE.SkinnedMesh,
-  attributeUsedIndexSetMap: Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, Set<number>>,
+  attributeUsedIndexSetMap: Map<
+    THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+    Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, Set<number>>
+  >,
 ): Map<THREE.Bone, THREE.Matrix4> {
   const boneInverseMap = new Map<THREE.Bone, THREE.Matrix4>();
 
   const skeleton = mesh.skeleton;
 
   const geometry = mesh.geometry;
   const skinIndexAttr = geometry.getAttribute('skinIndex');
-  const usedIndicesSet = attributeUsedIndexSetMap.get(skinIndexAttr);
+  const skinWeightAttr = geometry.getAttribute('skinWeight');
+  const skinIndexMap = attributeUsedIndexSetMap.get(skinIndexAttr);
+  const usedIndicesSet = skinIndexMap?.get(skinWeightAttr);
 
   if (!usedIndicesSet) {
-    throw new Error('Unreachable. attributeUsedIndexSetMap does not know the skin index attribute');
+    throw new Error(
+      'Unreachable. attributeUsedIndexSetMap does not know the skin index attribute or the skin weight attribute.',
+    );
   }
 
   for (const index of usedIndicesSet) {
@@ -222,3 +235,62 @@ function matrixEquals(a: THREE.Matrix4, b: THREE.Matrix4, tolerance?: number) {
 
   return true;
 }
+
+/**
+ * Check if the contents of two arrays are equal.
+ */
+function arrayEquals<T>(a: T[], b: T[]): boolean {
+  if (a.length !== b.length) {
+    return false;
+  }
+
+  return a.every((v, i) => v === b[i]);
+}
+
+/**
+ * Collect skin index attributes and corresponding bone arrays from the given skinned meshes.
+ * If a skin index attribute is shared among different bone sets, clone the attribute.
+ */
+function collectSkinIndexAttrs(
+  meshes: Iterable<THREE.SkinnedMesh>,
+): Set<[THREE.BufferAttribute | THREE.InterleavedBufferAttribute, THREE.Bone[]]> {
+  const skinIndexBonesPairSet = new Set<[THREE.BufferAttribute | THREE.InterleavedBufferAttribute, THREE.Bone[]]>();
+
+  // Collect skin index attributes
+  // skinIndex attribute might be shared among different bone sets
+  // If there are multiple bone sets that share the same skinIndex attribute, clone the attribute
+  const skinIndexNewSkinIndexBonesMapMap = new Map<
+    THREE.BufferAttribute | THREE.InterleavedBufferAttribute,
+    Map<THREE.BufferAttribute | THREE.InterleavedBufferAttribute, THREE.Bone[]>
+  >();
+
+  for (const mesh of meshes) {
+    let skinIndexAttr = mesh.geometry.getAttribute('skinIndex');
+
+    // Get or create a map for the skin index attribute
+    let newSkinIndexBonesMap = skinIndexNewSkinIndexBonesMapMap.get(skinIndexAttr);
+    if (newSkinIndexBonesMap == null) {
+      // Create a new map for the skin index attribute and register the bone array
+      newSkinIndexBonesMap = new Map();
+      skinIndexNewSkinIndexBonesMapMap.set(skinIndexAttr, newSkinIndexBonesMap);
+      newSkinIndexBonesMap.set(skinIndexAttr, mesh.skeleton.bones);
+      skinIndexBonesPairSet.add([skinIndexAttr, mesh.skeleton.bones]);
+      continue;
+    }
+
+    // Check if the bone set is already registered
+    for (const bones of newSkinIndexBonesMap.values()) {
+      if (arrayEquals(bones, mesh.skeleton.bones)) {
+        continue;
+      }
+    }
+
+    // There is no matching bone set, so clone the skin index attribute
+    skinIndexAttr = skinIndexAttr.clone();
+    mesh.geometry.setAttribute('skinIndex', skinIndexAttr);
+    newSkinIndexBonesMap.set(skinIndexAttr, mesh.skeleton.bones);
+    skinIndexBonesPairSet.add([skinIndexAttr, mesh.skeleton.bones]);
+  }
+
+  return skinIndexBonesPairSet;
+}