[naga] Preserve spans when compacting Arenas.

CHANGELOG.md

@@ -65,6 +65,8 @@ Bottom level categories:
 
 - 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:#?}");
+ }
+}