From 3ec547cdcaaa14488327d8f1b5f7736278c4178d Mon Sep 17 00:00:00 2001
From: Jim Blandy <jimb@red-bean.com>
Date: Wed, 15 Nov 2023 18:40:27 -0800
Subject: [PATCH] [naga] Preserve spans when compacting Arenas.

When compacting a module, properly adjust spans along with `Arena`
contents.
---
 CHANGELOG.md              |  2 ++
 naga/src/arena.rs         | 21 ++++++++++++++++++---
 naga/tests/wgsl-errors.rs | 28 ++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 84773843ad..4bdc068661 100644
--- a/CHANGELOG.md
+++ b/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)
diff --git a/naga/src/arena.rs b/naga/src/arena.rs
index a9df066654..54c92e849a 100644
--- a/naga/src/arena.rs
+++ b/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);
     }
 }
 
diff --git a/naga/tests/wgsl-errors.rs b/naga/tests/wgsl-errors.rs
index e32bd0baa8..6530082776 100644
--- a/naga/tests/wgsl-errors.rs
+++ b/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:#?}");
+    }
+}