[naga] Preserve spans when compacting Arenas.

When compacting a module, properly adjust spans along with `Arena`
contents.

CHANGELOG.md

@@ -73,6 +73,8 @@ Passing an owned value `window` to `Surface` will return a `Surface<'static>`. S
 
 - When evaluating const-expressions and generating SPIR-V, properly handle `Compose` expressions whose operands are `Splat` expressions. Such expressions are created and marked as constant by the constant evaluator. By @jimblandy in [#4695](https://github.com/gfx-rs/wgpu/pull/4695).
 
+- Preserve the source spans for constants and expressions correctly across module compaction. By @jimblandy in [#4696](https://github.com/gfx-rs/wgpu/pull/4696).
+
 ### Examples
 
 - remove winit dependency from hello-compute example by @psvri in [#4699](https://github.com/gfx-rs/wgpu/pull/4699)

naga/src/arena.rs

@@ -430,11 +430,26 @@ impl<T> Arena<T> {
         P: FnMut(Handle<T>, &mut T) -> bool,
     {
         let mut index = 0;
+        let mut retained = 0;
         self.data.retain_mut(|elt| {
+            let handle = Handle::new(Index::new(index as u32 + 1).unwrap());
+            let keep = predicate(handle, elt);
+
+            // Since `predicate` needs mutable access to each element,
+            // we can't feasibly call it twice, so we have to compact
+            // spans by hand in parallel as part of this iteration.
+            #[cfg(feature = "span")]
+            if keep {
+                self.span_info[retained] = self.span_info[index];
+                retained += 1;
+            }
+
             index += 1;
-            let handle = Handle::new(Index::new(index).unwrap());
-            predicate(handle, elt)
-        })
+            keep
+        });
+
+        #[cfg(feature = "span")]
+        self.span_info.truncate(retained);
     }
 }
 

naga/tests/wgsl-errors.rs

@@ -1992,3 +1992,31 @@ fn binding_array_non_struct() {
         })
     }
 }
+
+#[cfg(feature = "span")]
+#[test]
+fn compaction_preserves_spans() {
+    let source = r#"
+        const a: i32 = -(-(-(-42i)));
+        const b: vec2<u32> = vec2<u32>(42u, 43i);
+    "#; //                   ^^^^^^^^^^^^^^^^^^^ correct error span: 68..87
+    let mut module = naga::front::wgsl::parse_str(source).expect("source ought to parse");
+    naga::compact::compact(&mut module);
+    let err = naga::valid::Validator::new(
+        naga::valid::ValidationFlags::all(),
+        naga::valid::Capabilities::default(),
+    )
+    .validate(&module)
+    .expect_err("source ought to fail validation");
+
+    // Ideally this would all just be a `matches!` with a big pattern,
+    // but the `Span` API is full of opaque structs.
+    let mut spans = err.spans();
+    let first_span = spans.next().expect("error should have at least one span").0;
+    if !matches!(
+        first_span.to_range(),
+        Some(std::ops::Range { start: 68, end: 87 })
+    ) {
+        panic!("Error message has wrong span:\n\n{err:#?}");
+    }
+}