From c53d7daa2d11596b6880394dcd855263661dbaf0 Mon Sep 17 00:00:00 2001
From: Fredrik Fornwall <fredrik@fornwall.net>
Date: Thu, 10 Aug 2023 22:03:15 +0200
Subject: [PATCH] [wgsl-in] Fix error message for invalid texture{Load,Store}()
 on 2d_array

Fixes https://github.com/gfx-rs/naga/issues/2431
---
 src/front/wgsl/lower/mod.rs | 10 ++++++++--
 src/front/wgsl/tests.rs     | 28 ++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/front/wgsl/lower/mod.rs b/src/front/wgsl/lower/mod.rs
index cf52d6ce80..84c0d993e1 100644
--- a/src/front/wgsl/lower/mod.rs
+++ b/src/front/wgsl/lower/mod.rs
@@ -1937,7 +1937,10 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
 
                             let (_, arrayed) = ctx.image_data(image, image_span)?;
                             let array_index = arrayed
-                                .then(|| self.expression(args.next()?, ctx.reborrow()))
+                                .then(|| {
+                                    args.min_args += 1;
+                                    self.expression(args.next()?, ctx.reborrow())
+                                })
                                 .transpose()?;
 
                             let value = self.expression(args.next()?, ctx.reborrow())?;
@@ -1968,7 +1971,10 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
 
                             let (class, arrayed) = ctx.image_data(image, image_span)?;
                             let array_index = arrayed
-                                .then(|| self.expression(args.next()?, ctx.reborrow()))
+                                .then(|| {
+                                    args.min_args += 1;
+                                    self.expression(args.next()?, ctx.reborrow())
+                                })
                                 .transpose()?;
 
                             let level = class
diff --git a/src/front/wgsl/tests.rs b/src/front/wgsl/tests.rs
index 31b5890707..3aec0acf2d 100644
--- a/src/front/wgsl/tests.rs
+++ b/src/front/wgsl/tests.rs
@@ -481,3 +481,31 @@ fn parse_alias() {
     )
     .unwrap();
 }
+
+#[test]
+fn parse_texture_load_store_expecting_four_args() {
+    for (func, texture) in [
+        (
+            "textureStore",
+            "texture_storage_2d_array<rg11b10float, write>",
+        ),
+        ("textureLoad", "texture_2d_array<i32>"),
+    ] {
+        let error = parse_str(&format!(
+            "
+            @group(0) @binding(0) var tex_los_res: {texture};
+            @compute
+            @workgroup_size(1)
+            fn main(@builtin(global_invocation_id) id: vec3<u32>) {{
+                var color = vec4(1, 1, 1, 1);
+                {func}(tex_los_res, id, color);
+            }}
+            "
+        ))
+        .unwrap_err();
+        assert_eq!(
+            error.message(),
+            "wrong number of arguments: expected 4, found 3"
+        );
+    }
+}