fix mesh shader handles not remapped during module composition

src/compose/mod.rs

@@ -1834,8 +1834,13 @@ impl Composer {
                 early_depth_test: ep.early_depth_test,
                 workgroup_size: ep.workgroup_size,
                 workgroup_size_overrides: ep.workgroup_size_overrides,
-                mesh_info: ep.mesh_info.clone(),
-                task_payload: ep.task_payload,
+                mesh_info: ep
+                    .mesh_info
+                    .as_ref()
+                    .map(|info| derived.remap_mesh_info(info)),
+                task_payload: ep.task_payload.map(|tp| derived.import_global(&tp)),
+                //mesh_info: ep.mesh_info.clone(),
+                //task_payload: ep.task_payload,
                 incoming_ray_payload: ep.incoming_ray_payload,
             });
         }

src/compose/test.rs

@@ -61,6 +61,41 @@ mod test {
         output_eq!(wgsl, "tests/expected/simple_compose.txt");
     }
 
+    #[test]
+    fn mesh_shader_compose() {
+        let mut composer = Composer::default().with_capabilities(Capabilities::MESH_SHADER);
+
+        let module = composer
+            .make_naga_module(NagaModuleDescriptor {
+                source: include_str!("tests/mesh_shader/shader.wgsl"),
+                file_path: "tests/mesh_shader/shader.wgsl",
+                ..Default::default()
+            })
+            .unwrap();
+
+        let info = composer.create_validator().validate(&module).unwrap();
+
+        #[allow(unused)]
+        let wgsl = naga::back::wgsl::write_string(
+            &module,
+            &info,
+            naga::back::wgsl::WriterFlags::EXPLICIT_TYPES,
+        )
+        .unwrap();
+
+        //println!("{}", wgsl);
+        //let mut f = std::fs::File::create("mesh_shader_compose.txt").unwrap();
+        //f.write_all(wgsl.as_bytes()).unwrap();
+        //drop(f);
+
+        assert!(module
+            .entry_points
+            .iter()
+            .any(|ep| ep.stage == naga::ShaderStage::Mesh));
+
+        output_eq!(wgsl, "tests/expected/mesh_shader_compose.txt");
+    }
+
     #[test]
     fn big_shaderdefs() {
         let mut composer = Composer::default();

src/compose/tests/expected/mesh_shader_compose.txt

@@ -0,0 +1,66 @@
+enable wgpu_mesh_shader;
+
+struct TaskPayload {
+    colorMask: vec4<f32>,
+    visible: bool,
+}
+
+struct VertexOutput {
+    @builtin(position) position: vec4<f32>,
+    @location(0) color: vec4<f32>,
+}
+
+struct PrimitiveOutput {
+    @builtin(triangle_indices) indices: vec3<u32>,
+    @builtin(cull_primitive) cull: bool,
+    @location(1) @per_primitive colorMask: vec4<f32>,
+}
+
+struct PrimitiveInput {
+    @location(1) @per_primitive colorMask: vec4<f32>,
+}
+
+struct MeshOutput {
+    @builtin(vertices) vertices: array<VertexOutput, 3>,
+    @builtin(primitives) primitives: array<PrimitiveOutput, 1>,
+    @builtin(vertex_count) vertex_count: u32,
+    @builtin(primitive_count) primitive_count: u32,
+}
+
+var<task_payload> taskPayload: TaskPayload;
+var<workgroup> workgroupData: f32;
+var<workgroup> mesh_output: MeshOutput;
+
+@task @payload(taskPayload) @workgroup_size(1, 1, 1) 
+fn ts_main() -> @builtin(mesh_task_size) vec3<u32> {
+    workgroupData = 1f;
+    taskPayload.colorMask = vec4<f32>(1f, 1f, 0f, 1f);
+    taskPayload.visible = true;
+    return vec3<u32>(1u, 1u, 1u);
+}
+
+@mesh(mesh_output) @workgroup_size(1, 1, 1) @payload(taskPayload) 
+fn ms_main() {
+    mesh_output.vertex_count = 3u;
+    mesh_output.primitive_count = 1u;
+    workgroupData = 2f;
+    mesh_output.vertices[0].position = vec4<f32>(0f, 1f, 0f, 1f);
+    let _e23: vec4<f32> = taskPayload.colorMask;
+    mesh_output.vertices[0].color = (vec4<f32>(0f, 1f, 0f, 1f) * _e23);
+    mesh_output.vertices[1].position = vec4<f32>(-1f, -1f, 0f, 1f);
+    let _e45: vec4<f32> = taskPayload.colorMask;
+    mesh_output.vertices[1].color = (vec4<f32>(0f, 0f, 1f, 1f) * _e45);
+    mesh_output.vertices[2].position = vec4<f32>(1f, -1f, 0f, 1f);
+    let _e67: vec4<f32> = taskPayload.colorMask;
+    mesh_output.vertices[2].color = (vec4<f32>(1f, 0f, 0f, 1f) * _e67);
+    mesh_output.primitives[0].indices = vec3<u32>(0u, 1u, 2u);
+    let _e88: bool = taskPayload.visible;
+    mesh_output.primitives[0].cull = !(_e88);
+    mesh_output.primitives[0].colorMask = vec4<f32>(1f, 0f, 1f, 1f);
+    return;
+}
+
+@fragment 
+fn fs_main(vertex: VertexOutput, primitive: PrimitiveInput) -> @location(0) vec4<f32> {
+    return (vertex.color * primitive.colorMask);
+}

src/compose/tests/mesh_shader/shader.wgsl

@@ -0,0 +1,81 @@
+//wgpu example mesh shader
+//https://github.com/gfx-rs/wgpu/blob/v28.0.0/examples/features/src/mesh_shader/shader.wgsl
+enable wgpu_mesh_shader;
+
+const positions = array(
+    vec4(0., 1., 0., 1.),
+    vec4(-1., -1., 0., 1.),
+    vec4(1., -1., 0., 1.)
+);
+const colors = array(
+    vec4(0., 1., 0., 1.),
+    vec4(0., 0., 1., 1.),
+    vec4(1., 0., 0., 1.)
+);
+
+struct TaskPayload {
+    colorMask: vec4<f32>,
+    visible: bool,
+}
+struct VertexOutput {
+    @builtin(position) position: vec4<f32>,
+    @location(0) color: vec4<f32>,
+}
+struct PrimitiveOutput {
+    @builtin(triangle_indices) indices: vec3<u32>,
+    @builtin(cull_primitive) cull: bool,
+    @per_primitive @location(1) colorMask: vec4<f32>,
+}
+struct PrimitiveInput {
+    @per_primitive @location(1) colorMask: vec4<f32>,
+}
+
+var<task_payload> taskPayload: TaskPayload;
+var<workgroup> workgroupData: f32;
+
+@task
+@payload(taskPayload)
+@workgroup_size(1)
+fn ts_main() -> @builtin(mesh_task_size) vec3<u32> {
+    workgroupData = 1.0;
+    taskPayload.colorMask = vec4(1.0, 1.0, 0.0, 1.0);
+    taskPayload.visible = true;
+    return vec3(1, 1, 1);
+}
+
+struct MeshOutput {
+    @builtin(vertices) vertices: array<VertexOutput, 3>,
+    @builtin(primitives) primitives: array<PrimitiveOutput, 1>,
+    @builtin(vertex_count) vertex_count: u32,
+    @builtin(primitive_count) primitive_count: u32,
+}
+
+var<workgroup> mesh_output: MeshOutput;
+
+@mesh(mesh_output)
+@payload(taskPayload)
+@workgroup_size(1)
+fn ms_main() {
+    mesh_output.vertex_count = 3;
+    mesh_output.primitive_count = 1;
+    workgroupData = 2.0;
+
+    mesh_output.vertices[0].position = positions[0];
+    mesh_output.vertices[0].color = colors[0] * taskPayload.colorMask;
+
+    mesh_output.vertices[1].position = positions[1];
+    mesh_output.vertices[1].color = colors[1] * taskPayload.colorMask;
+
+    mesh_output.vertices[2].position = positions[2];
+    mesh_output.vertices[2].color = colors[2] * taskPayload.colorMask;
+
+    mesh_output.primitives[0].indices = vec3<u32>(0, 1, 2);
+    mesh_output.primitives[0].cull = !taskPayload.visible;
+    mesh_output.primitives[0].colorMask = vec4<f32>(1.0, 0.0, 1.0, 1.0);
+}
+
+@fragment
+fn fs_main(vertex: VertexOutput, primitive: PrimitiveInput) -> @location(0) vec4<f32> {
+    return vertex.color * primitive.colorMask;
+}
+

src/derive.rs

@@ -878,6 +878,24 @@ impl<'a> DerivedModule<'a> {
         }
     }
 
+    // remap mesh stage info handles into our derived context
+    pub fn remap_mesh_info(&mut self, info: &naga::MeshStageInfo) -> naga::MeshStageInfo {
+        naga::MeshStageInfo {
+            topology: info.topology,
+            max_vertices: info.max_vertices,
+            max_vertices_override: info
+                .max_vertices_override
+                .map(|e| self.import_global_expression(e)),
+            max_primitives: info.max_primitives,
+            max_primitives_override: info
+                .max_primitives_override
+                .map(|e| self.import_global_expression(e)),
+            vertex_output_type: self.import_type(&info.vertex_output_type),
+            primitive_output_type: self.import_type(&info.primitive_output_type),
+            output_variable: self.import_global(&info.output_variable),
+        }
+    }
+
     // import a function defined in the source shader context.
     // func name may be already defined, the returned handle will refer to the new function.
     // the previously defined function will still be valid.