Fix out of range mesh when assigning variant userData (google#3195)

* Quick fix to make sure the mesh index is valid when assigning variant info to userData on model load

* VariantMaterialLoaderPlugin afterRoot refactor to avoid looping over the primitives when traversing single mesh and avoid out-of-bounds errors

* Remove ts-ignore from VariantMaterialLoaderPlugin.mappingsArrayToTable definition

* Added Some tests for the fix

* Adding new glb asset to tests single mesh with primitive and variants edge case

* Remove duplicated test

.gitignore

@@ -3,4 +3,5 @@
 .DS_Store
 lib
 node_modules
-dist
+dist
+.idea

packages/model-viewer/src/test/features/scene-graph-spec.ts

@@ -23,6 +23,7 @@ import {ModelViewerGLTFInstance} from '../../three-components/gltf-instance/Mode
 import {waitForEvent} from '../../utilities.js';
 import {assetPath, rafPasses} from '../helpers.js';
 import {BasicSpecTemplate} from '../templates.js';
+import {ModelScene} from "../../three-components/ModelScene";
 
 
 
@@ -31,11 +32,18 @@ const expect = chai.expect;
 const ASTRONAUT_GLB_PATH = assetPath('models/Astronaut.glb');
 const HORSE_GLB_PATH = assetPath('models/Horse.glb');
 const CUBES_GLB_PATH = assetPath('models/cubes.gltf');  // has variants
+const MESH_PRIMITIVES_GLB_PATH = assetPath('models/MeshPrimitivesVariants.glb'); // has variants
 const CUBE_GLB_PATH = assetPath('models/cube.gltf');    // has UV coords
 const SUNRISE_IMG_PATH = assetPath('environments/spruit_sunrise_1k_LDR.jpg');
 const RIGGEDFIGURE_GLB_PATH = assetPath(
     'models/glTF-Sample-Models/2.0/RiggedFigure/glTF-Binary/RiggedFigure.glb');
 
+function getGLTFRoot(scene: ModelScene, hasBeenExportedOnce = false) {
+  // TODO: export is putting in an extra node layer, because the loader
+  // gives us a Group, but if the exporter doesn't get a Scene, then it
+  // wraps everything in an "AuxScene" node. Feels like a three.js bug.
+  return hasBeenExportedOnce ? scene.modelContainer.children[0].children[0] : scene.modelContainer.children[0];
+}
 
 suite('ModelViewerElementBase with SceneGraphMixin', () => {
   let nextId = 0;
@@ -88,13 +96,13 @@ suite('ModelViewerElementBase with SceneGraphMixin', () => {
       test('has variants', () => {
         expect(element[$scene].currentGLTF!.userData.variants.length)
             .to.be.eq(3);
-        const glTFroot = element[$scene].modelContainer.children[0];
-        expect(glTFroot.children[0].userData.variantMaterials.size).to.be.eq(3);
-        expect(glTFroot.children[1].userData.variantMaterials.size).to.be.eq(3);
+        const gltfRoot = getGLTFRoot(element[$scene]);
+        expect(gltfRoot.children[0].userData.variantMaterials.size).to.be.eq(3);
+        expect(gltfRoot.children[1].userData.variantMaterials.size).to.be.eq(3);
       });
 
       test(
-          `Setting varianName to null results in primitive
+          `Setting variantName to null results in primitive
            reverting to default/initial material`,
           async () => {
             let primitiveNode: PrimitiveNode|null = null
@@ -122,7 +130,7 @@ suite('ModelViewerElementBase with SceneGraphMixin', () => {
                 .equal('purple');
           });
 
-      test('exports and reimports the model with variants', async () => {
+      test('exports and re-imports the model with variants', async () => {
         const exported = await element.exportScene({binary: true});
         const url = URL.createObjectURL(exported);
         element.src = url;
@@ -131,12 +139,70 @@ suite('ModelViewerElementBase with SceneGraphMixin', () => {
 
         expect(element[$scene].currentGLTF!.userData.variants.length)
             .to.be.eq(3);
-        // TODO: export is putting in an extra node layer, because the loader
-        // gives us a Group, but if the exporter doesn't get a Scene, then it
-        // wraps everything in an "AuxScene" node. Feels like a three.js bug.
-        const glTFroot = element[$scene].modelContainer.children[0].children[0];
-        expect(glTFroot.children[0].userData.variantMaterials.size).to.be.eq(3);
-        expect(glTFroot.children[1].userData.variantMaterials.size).to.be.eq(3);
+        const gltfRoot = getGLTFRoot(element[$scene], true);
+        expect(gltfRoot.children[0].userData.variantMaterials.size).to.be.eq(3);
+        expect(gltfRoot.children[1].userData.variantMaterials.size).to.be.eq(3);
+      });
+    });
+
+    suite('with a loaded model containing a mesh with multiple primitives', () => {
+      setup(async () => {
+        element.src = MESH_PRIMITIVES_GLB_PATH;
+
+        await waitForEvent(element, 'load');
+        await rafPasses();
+      });
+
+      test('has variants', () => {
+        expect(element[$scene].currentGLTF!.userData.variants.length)
+            .to.be.eq(2);
+        const gltfRoot = getGLTFRoot(element[$scene]);
+        expect(gltfRoot.children[0].children[0].userData.variantMaterials.size).to.be.eq(2);
+        expect(gltfRoot.children[0].children[1].userData.variantMaterials.size).to.be.eq(2);
+        expect(gltfRoot.children[0].children[2].userData.variantMaterials.size).to.be.eq(2);
+      });
+
+      test(`Setting variantName to null results in primitive
+           reverting to default/initial material`, async () => {
+        let primitiveNode: PrimitiveNode|null = null
+        // Finds the first primitive with material 0 assigned.
+        for (const primitive of element.model![$primitivesList]) {
+          if (primitive.variantInfo != null &&
+              primitive[$initialMaterialIdx] == 0) {
+            primitiveNode = primitive;
+            return;
+          }
+        }
+
+        expect(primitiveNode).to.not.be.null;
+
+        // Switches to a new variant.
+        element.variantName = 'Inverse';
+        await waitForEvent(element, 'variant-applied');
+        expect((primitiveNode!.mesh.material as MeshStandardMaterial).name)
+            .equal('STEEL RED X');
+
+        // Switches to null variant.
+        element.variantName = null;
+        await waitForEvent(element, 'variant-applied');
+        expect((primitiveNode!.mesh.material as MeshStandardMaterial).name)
+            .equal('STEEL METALLIC');
+      });
+
+      test('exports and re-imports the model with variants', async () => {
+        const exported = await element.exportScene({binary: true});
+        const url = URL.createObjectURL(exported);
+        element.src = url;
+        await waitForEvent(element, 'load');
+        await rafPasses();
+
+        expect(element[$scene].currentGLTF!.userData.variants.length)
+            .to.be.eq(2);
+
+        const gltfRoot = getGLTFRoot(element[$scene], true);
+        expect(gltfRoot.children[0].children[0].userData.variantMaterials.size).to.be.eq(2);
+        expect(gltfRoot.children[0].children[1].userData.variantMaterials.size).to.be.eq(2);
+        expect(gltfRoot.children[0].children[2].userData.variantMaterials.size).to.be.eq(2);
       });
     });
 

packages/model-viewer/src/test/features/scene-graph/nodes/primitive-node-spec.ts

@@ -30,6 +30,7 @@ const BRAIN_STEM_GLB_PATH = assetPath(
     'models/glTF-Sample-Models/2.0/BrainStem/glTF-Binary/BrainStem.glb');
 const CUBES_GLTF_PATH = assetPath('models/cubes.gltf');
 const CUBE_GLTF_PATH = assetPath('models/cube.gltf');
+const MESH_PRIMITIVES_GLB_PATH = assetPath('models/MeshPrimitivesVariants.glb');
 const KHRONOS_TRIANGLE_GLB_PATH =
     assetPath('models/glTF-Sample-Models/2.0/Triangle/glTF/Triangle.gltf');
 
@@ -170,6 +171,69 @@ suite('scene-graph/model/mesh-primitives', () => {
     });
   });
 
+  suite('Mesh with multiple primitives each with variants', () => {
+    let model: Model;
+    setup(async () => {
+      const threeGLTF = await loadThreeGLTF(MESH_PRIMITIVES_GLB_PATH);
+      model = new Model(CorrelatedSceneGraph.from(threeGLTF));
+    });
+
+    test('Primitive count matches glTF file', async () => {
+      expect(model![$primitivesList].length).to.equal(3) ;
+    });
+
+    test('Primitives should have expected variant names', async () => {
+      expect(findPrimitivesWithVariant(model, 'Normal')).to.not.be.null;
+      expect(findPrimitivesWithVariant(model, 'Inverse')).to.not.be.null;
+    });
+
+    test('Switching to incorrect variant name', async () => {
+      const primitive = findPrimitivesWithVariant(model, 'Normal')![0];
+      const material = await primitive.enableVariant('Does not exist');
+      expect(material).to.be.null;
+    });
+
+    test('Switching to variant and then switch back', async () => {
+      const MATERIAL_NAME = 'STEEL BLACK';
+      const primitives = findPrimitivesWithVariant(model, 'Normal')!;
+      let materials = new Array<MeshStandardMaterial>();
+      for (const primitive of primitives) {
+        materials.push(
+            await primitive.enableVariant('Inverse') as
+                MeshStandardMaterial);
+      }
+
+      expect(materials).to.not.be.empty;
+      expect(materials.find((material: MeshStandardMaterial) => {
+        return material.name === MATERIAL_NAME;
+      })).to.be.undefined;
+
+      materials = new Array<MeshStandardMaterial>();
+      for (const primitive of primitives) {
+        materials.push(
+            await primitive.enableVariant('Normal') as
+                MeshStandardMaterial);
+      }
+
+      expect(materials.find((material: MeshStandardMaterial) => {
+        return material.name === MATERIAL_NAME;
+      })).to.be.ok;
+    });
+
+    test('Primitive switches to initial material', async () => {
+      const primitive = findPrimitivesWithVariant(model, 'Normal')![0];
+
+      // Gets current material.
+      const initialMaterial = await primitive.enableVariant('Normal');
+      // Switches to variant.
+      const variantMaterial = await primitive.enableVariant('Inverse');
+      expect(initialMaterial).to.not.equal(variantMaterial)
+      // Switches to initial material.
+      const resetMaterial = await primitive.enableVariant(null);
+      expect(resetMaterial).to.equal(initialMaterial);
+    });
+  });
+
   suite('Skinned Primitive Without Variant', () => {
     test('Primitive count matches glTF file', async () => {
       const threeGLTF = await loadThreeGLTF(BRAIN_STEM_GLB_PATH);

packages/model-viewer/src/three-components/CachingGLTFLoader.ts

@@ -212,8 +212,10 @@ export class CachingGLTFLoader<T extends GLTFInstanceConstructor =
                                     .then((preparedGLTF) => {
                                       progressCallback(0.9);
                                       return new GLTFInstance(preparedGLTF);
-                                    });
-
+                                    }).catch((reason => {
+                                      console.error(reason);
+                                      return new GLTFInstance();
+                                    }));
       cache.set(url, gltfInstanceLoads);
     }
 

packages/model-viewer/src/three-components/gltf-instance/VariantMaterialLoaderPlugin.ts

@@ -25,8 +25,9 @@
  * https://github.com/takahirox/three-gltf-extensions/tree/main/loaders/KHR_materials_variants
  */
 
-import {Material as ThreeMaterial} from 'three';
+import {Material as ThreeMaterial, Mesh} from 'three';
 import {GLTF, GLTFLoaderPlugin, GLTFParser} from 'three/examples/jsm/loaders/GLTFLoader.js';
+import {GLTFReference} from "three/examples/jsm/loaders/GLTFLoader";
 
 
 export interface UserDataVariantMapping {
@@ -72,8 +73,7 @@ const ensureUniqueNames = (variantNames: string[]) => {
  * @param variantNames {Array<string>} Required to be unique names
  * @return {Map}
  */
-const mappingsArrayToTable = (extensionDef:
-                                  any): Map<number, UserDataVariantMapping> => {
+const mappingsArrayToTable = (extensionDef: any): Map<number, UserDataVariantMapping> => {
   const table = new Map<number, UserDataVariantMapping>();
   for (const mapping of extensionDef.mappings) {
     for (const variant of mapping.variants) {
@@ -114,38 +114,28 @@ export default class GLTFMaterialsVariantsExtension implements
     for (const scene of gltf.scenes) {
       // Save the variants data under associated mesh.userData
       scene.traverse(object => {
-        // The following code can be simplified if parser.associations directly
-        // supports meshes.
-        const association = parser.associations.get(object);
+        const mesh = object as Mesh;
 
-        if (association == null || association.meshes == null) {
+        if (!mesh.isMesh) {
           return;
         }
 
-        const meshIndex = association.meshes;
+        const association = parser.associations.get(mesh) as GLTFReference & {primitives: number};
 
-        // Two limitations:
-        // 1. The nodeDef shouldn't have any objects (camera, light, or
-        // nodeDef.extensions object)
-        //    other than nodeDef.mesh
-        // 2. Other plugins shouldn't change any scene graph hierarchy
-        // The following code can cause error if hitting the either or both
-        // limitations If parser.associations will directly supports meshes
-        // these limitations can be removed
+        if (association == null || association.meshes == null || association.primitives == null) {
+          return;
+        }
 
-        const meshDef = json.meshes[meshIndex];
+        const meshDef = json.meshes[association.meshes];
         const primitivesDef = meshDef.primitives;
-        const meshes = 'isMesh' in object ? [object] : object.children;
-
-        for (let i = 0; i < primitivesDef.length; i++) {
-          const primitiveDef = primitivesDef[i];
-          const extensionsDef = primitiveDef.extensions;
-          if (!extensionsDef || !extensionsDef[this.name]) {
-            continue;
-          }
-          meshes[i].userData.variantMaterials =
-              mappingsArrayToTable(extensionsDef[this.name]);
+        const primitiveDef = primitivesDef[association.primitives];
+        const extensionsDef = primitiveDef.extensions;
+
+        if (!extensionsDef || !extensionsDef[this.name]) {
+          return;
         }
+
+        mesh.userData.variantMaterials = mappingsArrayToTable(extensionsDef[this.name]);
       });
     }
 

packages/space-opera/package.json

@@ -15,7 +15,7 @@
     "watch:test": "karma start",
     "watch:tsc": "tsc -w",
     "watch:rollup": "rollup -c -w",
-    "serve": "./node_modules/.bin/http-server -a 127.0.0.1 -o /demo/ -c-1",
+    "serve": "./node_modules/.bin/http-server -a 127.0.0.1 -o /editor/ -c-1",
     "pyserve": "echo This assumes 'python' is Python 2. && which python && echo Visit $(hostname):8000/index.html && python -m SimpleHTTPServer 8000",
     "dev": "npm run build -- --incremental && npm-run-all --parallel 'watch:tsc -- --preserveWatchOutput --incremental' 'watch:rollup' 'watch:test' 'serve -- -s'",
     "dev-py": "npm run build -- --incremental && npm-run-all --parallel 'watch:tsc -- --preserveWatchOutput --incremental' 'watch:rollup' 'watch:test' 'pyserve'"