Fix type resolution issue in GL push constants (#9116)

naga/src/back/glsl/mod.rs

@@ -446,11 +446,18 @@ pub struct ImmediateItem {
     ///
     pub access_path: String,
     /// Type of the uniform. This will only ever be a scalar, vector, or matrix.
-    pub ty: Handle<crate::Type>,
-    /// The offset in the immediate data memory block this uniform maps to.
     ///
-    /// The size of the uniform can be derived from the type.
+    /// Stored as a [`TypeInner`] rather than a [`Handle<Type>`] because
+    /// `process_overrides` may compact the module, renumbering type handles.
+    /// Leaf types don't reference other types, so a `TypeInner` is self-contained.
+    ///
+    /// [`TypeInner`]: crate::TypeInner
+    /// [`Handle<Type>`]: Handle
+    pub ty: TypeInner,
+    /// The offset in the immediate data memory block this uniform maps to.
     pub offset: u32,
+    /// Size of this uniform in bytes.
+    pub size_bytes: u32,
 }
 
 /// Helper structure that generates a number

naga/src/back/glsl/writer.rs

@@ -4588,7 +4588,8 @@ impl<'a, W: Write> Writer<'a, W> {
                 items.push(ImmediateItem {
                     access_path: name,
                     offset: *offset,
-                    ty,
+                    ty: self.module.types[ty].inner.clone(),
+                    size_bytes: layout.size,
                 });
                 *offset += layout.size;
             }

tests/tests/wgpu-gpu/main.rs

@@ -10,6 +10,7 @@ mod regression {
     pub mod issue_6317;
     pub mod issue_6467;
     pub mod issue_6827;
+    pub mod issue_9115;
 }
 
 mod bgra8unorm_storage;
@@ -136,6 +137,7 @@ fn all_tests() -> Vec<wgpu_test::GpuTestInitializer> {
     regression::issue_6317::all_tests(&mut tests);
     regression::issue_6467::all_tests(&mut tests);
     regression::issue_6827::all_tests(&mut tests);
+    regression::issue_9115::all_tests(&mut tests);
     render_pass_ownership::all_tests(&mut tests);
     render_target::all_tests(&mut tests);
     resource_descriptor_accessor::all_tests(&mut tests);

tests/tests/wgpu-gpu/regression/issue_9115.rs

@@ -0,0 +1,171 @@
+use wgpu::util::DeviceExt;
+use wgpu_test::{
+    gpu_test, image::ReadbackBuffers, GpuTestConfiguration, GpuTestInitializer, TestParameters,
+    TestingContext,
+};
+
+pub fn all_tests(vec: &mut Vec<GpuTestInitializer>) {
+    vec.push(IMMEDIATES_WITH_UNIFORM_IN_SINGLE_MODULE);
+}
+
+/// On the GLES backend, using immediates in both vertex and fragment shaders from a single
+/// shader module, while also having a uniform buffer with a struct type, caused a panic.
+/// The backend incorrectly encountered the uniform's struct type when processing immediates.
+///
+/// See <https://github.com/gfx-rs/wgpu/issues/9115>.
+#[gpu_test]
+static IMMEDIATES_WITH_UNIFORM_IN_SINGLE_MODULE: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(
+        TestParameters::default()
+            .features(wgpu::Features::IMMEDIATES)
+            .limits(wgpu::Limits {
+                max_immediate_size: 32,
+                ..Default::default()
+            }),
+    )
+    .run_async(immediates_with_uniform_in_single_module);
+
+#[repr(C)]
+#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
+struct Immediates {
+    position: [f32; 2],
+    size: [f32; 2],
+    color: [f32; 4],
+}
+
+async fn immediates_with_uniform_in_single_module(ctx: TestingContext) {
+    let shader = ctx
+        .device
+        .create_shader_module(wgpu::include_wgsl!("issue_9115.wgsl"));
+
+    let globals_buffer = ctx
+        .device
+        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: Some("globals"),
+            contents: bytemuck::cast_slice(&[1.0_f32, 1.0]),
+            usage: wgpu::BufferUsages::UNIFORM,
+        });
+
+    let bgl = ctx
+        .device
+        .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+            label: Some("bgl"),
+            entries: &[wgpu::BindGroupLayoutEntry {
+                binding: 0,
+                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
+                ty: wgpu::BindingType::Buffer {
+                    ty: wgpu::BufferBindingType::Uniform,
+                    has_dynamic_offset: false,
+                    min_binding_size: None,
+                },
+                count: None,
+            }],
+        });
+
+    let bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
+        label: Some("bg"),
+        layout: &bgl,
+        entries: &[wgpu::BindGroupEntry {
+            binding: 0,
+            resource: globals_buffer.as_entire_binding(),
+        }],
+    });
+
+    let pll = ctx
+        .device
+        .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+            label: Some("pll"),
+            bind_group_layouts: &[Some(&bgl)],
+            immediate_size: size_of::<Immediates>() as u32,
+        });
+
+    let pipeline = ctx
+        .device
+        .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("pipeline"),
+            layout: Some(&pll),
+            vertex: wgpu::VertexState {
+                module: &shader,
+                entry_point: Some("vs_main"),
+                compilation_options: Default::default(),
+                buffers: &[],
+            },
+            fragment: Some(wgpu::FragmentState {
+                module: &shader,
+                entry_point: Some("fs_main"),
+                compilation_options: Default::default(),
+                targets: &[Some(wgpu::ColorTargetState {
+                    format: wgpu::TextureFormat::Rgba8Unorm,
+                    blend: None,
+                    write_mask: wgpu::ColorWrites::ALL,
+                })],
+            }),
+            primitive: wgpu::PrimitiveState::default(),
+            depth_stencil: None,
+            multisample: wgpu::MultisampleState::default(),
+            multiview_mask: None,
+            cache: None,
+        });
+
+    let texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
+        label: Some("texture"),
+        size: wgpu::Extent3d {
+            width: 1,
+            height: 1,
+            depth_or_array_layers: 1,
+        },
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu::TextureDimension::D2,
+        format: wgpu::TextureFormat::Rgba8Unorm,
+        usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::RENDER_ATTACHMENT,
+        view_formats: &[],
+    });
+
+    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+    let mut encoder = ctx
+        .device
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
+            label: Some("encoder"),
+        });
+
+    {
+        let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+            label: Some("rpass"),
+            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                view: &view,
+                depth_slice: None,
+                resolve_target: None,
+                ops: wgpu::Operations {
+                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
+                    store: wgpu::StoreOp::Store,
+                },
+            })],
+            depth_stencil_attachment: None,
+            timestamp_writes: None,
+            occlusion_query_set: None,
+            multiview_mask: None,
+        });
+
+        rpass.set_pipeline(&pipeline);
+        rpass.set_bind_group(0, &bg, &[]);
+        rpass.set_immediates(
+            0,
+            bytemuck::cast_slice(&[Immediates {
+                position: [0.0, 0.0],
+                size: [1.0, 1.0],
+                color: [0.0, 1.0, 0.0, 1.0],
+            }]),
+        );
+        rpass.draw(0..3, 0..1);
+    }
+
+    let buffers = ReadbackBuffers::new(&ctx.device, &texture);
+    buffers.copy_from(&ctx.device, &mut encoder, &texture);
+    ctx.queue.submit([encoder.finish()]);
+
+    // The fragment shader outputs immediates.color, which is green (0, 1, 0, 1).
+    let expected: [u8; 4] = [0, 255, 0, 255];
+    buffers.assert_buffer_contents(&ctx, &expected).await;
+}

tests/tests/wgpu-gpu/regression/issue_9115.wgsl

@@ -0,0 +1,35 @@
+// Regression test for https://github.com/gfx-rs/wgpu/issues/9115
+//
+// On the GLES backend, using immediates in both vertex and fragment shaders
+// from a single shader module while also having a uniform buffer causes a panic:
+//   "Unsupported uniform datatype: Struct { ... }"
+// The backend confuses the uniform struct type with the immediates struct.
+
+struct Globals {
+    inv_screen_size: vec2f,
+}
+@group(0) @binding(0)
+var<uniform> globals: Globals;
+
+struct Immediates {
+    position: vec2f,
+    size: vec2f,
+    color: vec4f,
+}
+var<immediate> immediates: Immediates;
+
+@vertex
+fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4f {
+    // Reference both globals and immediates in the vertex shader.
+    _ = globals.inv_screen_size;
+    _ = immediates.position + immediates.size;
+
+    let uv = vec2f(f32((vertex_index << 1u) & 2u), f32(vertex_index & 2u));
+    return vec4f(uv * 2.0 - 1.0, 0.0, 1.0);
+}
+
+@fragment
+fn fs_main() -> @location(0) vec4f {
+    // Accessing immediates in the fragment shader is what triggers the GLES panic.
+    return immediates.color;
+}

wgpu-hal/src/gles/device.rs

@@ -503,22 +503,14 @@ impl super::Device {
 
         let mut uniforms = ArrayVec::new();
 
-        for (stage_idx, stage_items) in immediates_items.into_iter().enumerate() {
+        for stage_items in immediates_items {
             for item in stage_items {
-                let source = &shaders[stage_idx].1.module.source;
-                let super::ShaderModuleSource::Naga(naga_module) = source else {
-                    // ImmediateItem can only be constructed given a naga module, as it requires a type handle.
-                    // Passthrough shaders will have immediates_items empty
-                    unreachable!("Passthrough shaders don't currently support immediates on GLES");
-                };
-                let type_inner = &naga_module.module.types[item.ty].inner;
-
                 let location = unsafe { gl.get_uniform_location(program, &item.access_path) };
 
                 log::trace!(
                     "immediate data item: name={}, ty={:?}, offset={}, location={:?}",
                     item.access_path,
-                    type_inner,
+                    item.ty,
                     item.offset,
                     location,
                 );
@@ -527,8 +519,8 @@ impl super::Device {
                     uniforms.push(super::ImmediateDesc {
                         location,
                         offset: item.offset,
-                        size_bytes: type_inner.size(naga_module.module.to_ctx()),
-                        ty: type_inner.clone(),
+                        size_bytes: item.size_bytes,
+                        ty: item.ty,
                     });
                 }
             }