Implement module compaction.

Add a new Naga feature, `"compact"`, which adds a new function
`naga::compact::compact`, which removes unused expressions, types, and
constants from a `Module`.

Author: jimblandy

Cargo.toml

@@ -38,6 +38,7 @@ wgsl-out = []
 hlsl-out = []
 span = ["codespan-reporting", "termcolor"]
 validate = []
+compact = []
 
 [[bench]]
 name = "criterion"

cli/Cargo.toml

@@ -10,9 +10,28 @@ keywords = ["shader", "SPIR-V", "GLSL", "MSL"]
 license = "MIT OR Apache-2.0"
 
 [dependencies]
-naga = { version = "0.13", path = "../", features = ["validate", "span", "wgsl-in", "wgsl-out", "glsl-in", "glsl-out", "spv-in", "spv-out", "msl-out", "hlsl-out", "dot-out", "serialize", "deserialize"] }
 bincode = "1"
 log = "0.4"
 codespan-reporting = "0.11"
 env_logger = "0.10"
 argh = "0.1.5"
+
+[dependencies.naga]
+version = "0.13"
+path = "../"
+features = [
+    "validate",
+    "compact",
+    "span",
+    "wgsl-in",
+    "wgsl-out",
+    "glsl-in",
+    "glsl-out",
+    "spv-in",
+    "spv-out",
+    "msl-out",
+    "hlsl-out",
+    "dot-out",
+    "serialize",
+    "deserialize"
+]

cli/src/bin/naga.rs

@@ -79,6 +79,20 @@ struct Args {
     #[argh(switch, short = 'g')]
     generate_debug_symbols: bool,
 
+    /// compact the module's IR and revalidate.
+    ///
+    /// Output files will reflect the compacted IR. If you want to see the IR as
+    /// it was before compaction, use the `--before-compaction` option.
+    #[argh(switch)]
+    compact: bool,
+
+    /// write the module's IR before compaction to the given file.
+    ///
+    /// This implies `--compact`. Like any other output file, the filename
+    /// extension determines the form in which the module is written.
+    #[argh(option)]
+    before_compaction: Option<String>,
+
     /// show version
     #[argh(switch)]
     version: bool,
@@ -288,7 +302,7 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
         !params.keep_coordinate_space,
     );
 
-    let (module, input_text) = match Path::new(&input_path)
+    let (mut module, input_text) = match Path::new(&input_path)
         .extension()
         .ok_or(CliError("Input filename has no extension"))?
         .to_str()
@@ -391,12 +405,13 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
                 caps & !missing
             });
 
-    // validate the IR
+    // Validate the IR before compaction.
     let info = match naga::valid::Validator::new(params.validation_flags, validation_caps)
         .validate(&module)
     {
         Ok(info) => Some(info),
         Err(error) => {
+            // Validation failure is not fatal. Just report the error.
             if let Some(input) = &input_text {
                 let filename = input_path.file_name().and_then(std::ffi::OsStr::to_str);
                 emit_annotated_error(&error, filename.unwrap_or("input"), input);
@@ -406,6 +421,45 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
         }
     };
 
+    // Compact the module, if requested.
+    let info = if args.compact || args.before_compaction.is_some() {
+        // Compact only if validation succeeded. Otherwise, compaction may panic.
+        if info.is_some() {
+            // Write out the module state before compaction, if requested.
+            if let Some(ref before_compaction) = args.before_compaction {
+                write_output(&module, &info, &params, before_compaction)?;
+            }
+
+            naga::compact::compact(&mut module);
+
+            // Re-validate the IR after compaction.
+            match naga::valid::Validator::new(params.validation_flags, validation_caps)
+                .validate(&module)
+            {
+                Ok(info) => Some(info),
+                Err(error) => {
+                    // Validation failure is not fatal. Just report the error.
+                    eprintln!("Error validating compacted module:");
+                    if let Some(input) = &input_text {
+                        let filename = input_path.file_name().and_then(std::ffi::OsStr::to_str);
+                        emit_annotated_error(&error, filename.unwrap_or("input"), input);
+                    }
+                    print_err(&error);
+                    None
+                }
+            }
+        } else {
+            eprintln!("Skipping compaction due to validation failure.");
+            None
+        }
+    } else {
+        info
+    };
+
+    // If no output was requested, then report validation results and stop here.
+    //
+    // If the user asked for output, don't stop: some output formats (".txt",
+    // ".dot", ".bin") can be generated even without a `ModuleInfo`.
     if output_paths.is_empty() {
         if info.is_some() {
             println!("Validation successful");

src/arena.rs

@@ -192,12 +192,31 @@ impl<T> Iterator for Range<T> {
 }
 
 impl<T> Range<T> {
+    /// Return a range enclosing handles `first` through `last`, inclusive.
     pub fn new_from_bounds(first: Handle<T>, last: Handle<T>) -> Self {
         Self {
             inner: (first.index() as u32)..(last.index() as u32 + 1),
             marker: Default::default(),
         }
     }
+
+    /// Return the first and last handles included in `self`.
+    ///
+    /// If `self` is an empty range, there are no handles included, so
+    /// return `None`.
+    pub fn first_and_last(&self) -> Option<(Handle<T>, Handle<T>)> {
+        if self.inner.start < self.inner.end {
+            Some((
+                // `Range::new_from_bounds` expects a 1-based, start- and
+                // end-inclusive range, but `self.inner` is a zero-based,
+                // end-exclusive range.
+                Handle::new(Index::new(self.inner.start + 1).unwrap()),
+                Handle::new(Index::new(self.inner.end).unwrap()),
+            ))
+        } else {
+            None
+        }
+    }
 }
 
 /// An arena holding some kind of component (e.g., type, constant,
@@ -382,6 +401,19 @@ impl<T> Arena<T> {
             Ok(())
         }
     }
+
+    #[cfg(feature = "compact")]
+    pub(crate) fn retain_mut<P>(&mut self, mut predicate: P)
+    where
+        P: FnMut(Handle<T>, &mut T) -> bool,
+    {
+        let mut index = 0;
+        self.data.retain_mut(|elt| {
+            index += 1;
+            let handle = Handle::new(Index::new(index).unwrap());
+            predicate(handle, elt)
+        })
+    }
 }
 
 #[cfg(feature = "deserialize")]
@@ -544,6 +576,44 @@ impl<T> UniqueArena<T> {
             Span::default()
         }
     }
+
+    #[cfg(feature = "compact")]
+    pub(crate) fn drain_all(&mut self) -> UniqueArenaDrain<T> {
+        UniqueArenaDrain {
+            inner_elts: self.set.drain(..),
+            #[cfg(feature = "span")]
+            inner_spans: self.span_info.drain(..),
+            index: Index::new(1).unwrap(),
+        }
+    }
+}
+
+#[cfg(feature = "compact")]
+pub(crate) struct UniqueArenaDrain<'a, T> {
+    inner_elts: indexmap::set::Drain<'a, T>,
+    #[cfg(feature = "span")]
+    inner_spans: std::vec::Drain<'a, Span>,
+    index: Index,
+}
+
+#[cfg(feature = "compact")]
+impl<'a, T> Iterator for UniqueArenaDrain<'a, T> {
+    type Item = (Handle<T>, T, Span);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.inner_elts.next() {
+            Some(elt) => {
+                let handle = Handle::new(self.index);
+                self.index = self.index.checked_add(1).unwrap();
+                #[cfg(feature = "span")]
+                let span = self.inner_spans.next().unwrap();
+                #[cfg(not(feature = "span"))]
+                let span = Span::default();
+                Some((handle, elt, span))
+            }
+            None => None,
+        }
+    }
 }
 
 impl<T: Eq + hash::Hash> UniqueArena<T> {