[spv-out] Properly reset Writers between module writes.

A single `Writer` can be reused to convert any number of Naga IR modules to
SPIR-V. This lets us reuse the heap allocations, but makes us responsible for
resetting the state of the `Writer` fully between modules. The old code
neglected to clear several different fields: logical_layout (all fields),
lookup_type, and cached_constants.

This commit uses Rust struct literal expressions to construct the re-initialized
`Writer`, which makes the compiler check that each field has been handled.

It also introduces a policy limiting how much storage will be retained, so that
processing the occasional whale doesn't leave us with leviathan buffers
permanently.

src/back/spv/layout.rs

@@ -25,6 +25,18 @@ impl PhysicalLayout {
     }
 }
 
+impl super::recyclable::Recyclable for PhysicalLayout {
+    fn recycle(self) -> Self {
+        PhysicalLayout {
+            magic_number: self.magic_number,
+            version: self.version,
+            generator: self.generator,
+            instruction_schema: self.instruction_schema,
+            bound: 0,
+        }
+    }
+}
+
 impl LogicalLayout {
     pub(super) fn in_words(&self, sink: &mut impl Extend<Word>) {
         sink.extend(self.capabilities.iter().cloned());
@@ -41,6 +53,24 @@ impl LogicalLayout {
     }
 }
 
+impl super::recyclable::Recyclable for LogicalLayout {
+    fn recycle(self) -> Self {
+        Self {
+            capabilities: self.capabilities.recycle(),
+            extensions: self.extensions.recycle(),
+            ext_inst_imports: self.ext_inst_imports.recycle(),
+            memory_model: self.memory_model.recycle(),
+            entry_points: self.entry_points.recycle(),
+            execution_modes: self.execution_modes.recycle(),
+            debugs: self.debugs.recycle(),
+            annotations: self.annotations.recycle(),
+            declarations: self.declarations.recycle(),
+            function_declarations: self.function_declarations.recycle(),
+            function_definitions: self.function_definitions.recycle(),
+        }
+    }
+}
+
 impl Instruction {
     pub(super) fn new(op: Op) -> Self {
         Instruction {

src/back/spv/mod.rs

@@ -5,6 +5,7 @@ mod helpers;
 mod index;
 mod instructions;
 mod layout;
+mod recyclable;
 mod writer;
 
 pub use spirv::Capability;
@@ -15,6 +16,7 @@ use spirv::Word;
 use std::ops;
 use thiserror::Error;
 
+#[derive(Clone)]
 struct PhysicalLayout {
     magic_number: Word,
     version: Word,
@@ -239,6 +241,13 @@ impl ops::IndexMut<Handle<crate::Expression>> for CachedExpressions {
         id
     }
 }
+impl recyclable::Recyclable for CachedExpressions {
+    fn recycle(self) -> Self {
+        CachedExpressions {
+            ids: self.ids.recycle(),
+        }
+    }
+}
 
 struct GlobalVariable {
     /// Actual ID of the variable.

src/back/spv/recyclable.rs

@@ -0,0 +1,74 @@
+//! Reusing collections' previous allocations.
+
+/// A value that can be reset to its initial state, retaining its current allocations.
+///
+/// Naga attempts to lower the cost of SPIR-V generation by allowing clients to
+/// reuse the same `Writer` for multiple Module translations. Reusing a `Writer`
+/// means that the `Vec`s, `HashMap`s, and other heap-allocated structures the
+/// `Writer` uses internally begin the translation with heap-allocated buffers
+/// ready to use.
+///
+/// But this approach introduces the risk of `Writer` state leaking from one
+/// module to the next. When a developer adds fields to `Writer` or its internal
+/// types, they must remember to reset their contents between modules.
+///
+/// One trick to ensure that every field has been accounted for is to use Rust's
+/// struct literal syntax to construct a new, reset value. If a developer adds a
+/// field, but neglects to update the reset code, the compiler will complain
+/// that a field is missing from the literal. This trait's `recycle` method
+/// takes `self` by value, and returns `Self` by value, encouraging the use of
+/// struct literal expressions in its implementation.
+pub trait Recyclable {
+    /// Clear `self`, retaining its current memory allocations.
+    ///
+    /// Shrink the buffer if it's currently much larger than was actually used.
+    /// This prevents a module with exceptionally large allocations from causing
+    /// the `Writer` to retain more memory than it needs indefinitely.
+    fn recycle(self) -> Self;
+}
+
+// Stock values for various collections.
+
+/// Maximum extra capacity that a recycled vector is allowed to have. If the
+/// actual capacity is larger, we re-allocate the vector storage with lower
+/// capacity.
+const MAX_EXTRA_CAPACITY_PERCENT: usize = 200;
+
+/// Minimum extra capacity to keep when re-allocating the vector storage.
+const MIN_EXTRA_CAPACITY_PERCENT: usize = 20;
+
+/// Minimum sensible length to consider for re-allocation.
+const MIN_LENGTH: usize = 16;
+
+impl<T> Recyclable for Vec<T> {
+    fn recycle(mut self) -> Self {
+        let extra_capacity = (self.capacity() - self.len()) * 100 / self.len().max(MIN_LENGTH);
+
+        if extra_capacity > MAX_EXTRA_CAPACITY_PERCENT {
+            //TODO: use `shrink_to` when it's stable
+            self = Vec::with_capacity(self.len() + self.len() * MIN_EXTRA_CAPACITY_PERCENT / 100);
+        } else {
+            self.clear();
+        }
+
+        self
+    }
+}
+
+impl<K, V, S: Clone> Recyclable for std::collections::HashMap<K, V, S> {
+    fn recycle(mut self) -> Self {
+        let extra_capacity = (self.capacity() - self.len()) * 100 / self.len().max(MIN_LENGTH);
+
+        if extra_capacity > MAX_EXTRA_CAPACITY_PERCENT {
+            //TODO: use `shrink_to` when it's stable
+            self = Self::with_capacity_and_hasher(
+                self.len() + self.len() * MIN_EXTRA_CAPACITY_PERCENT / 100,
+                self.hasher().clone(),
+            );
+        } else {
+            self.clear();
+        }
+
+        self
+    }
+}

src/back/spv/writer.rs

@@ -118,9 +118,6 @@ impl Writer {
             return Err(Error::UnsupportedVersion(major, minor));
         }
         let raw_version = ((major as u32) << 16) | ((minor as u32) << 8);
-        let mut id_gen = IdGenerator::default();
-        let gl450_ext_inst_id = id_gen.next();
-        let void_type = id_gen.next();
 
         let (capabilities, forbidden_caps) = match options.capabilities {
             Some(ref caps) => (caps.clone(), None),
@@ -132,6 +129,10 @@ impl Writer {
             }
         };
 
+        let mut id_gen = IdGenerator::default();
+        let gl450_ext_inst_id = id_gen.next();
+        let void_type = id_gen.next();
+
         Ok(Writer {
             physical_layout: PhysicalLayout::new(raw_version),
             logical_layout: LogicalLayout::default(),
@@ -156,6 +157,56 @@ impl Writer {
         })
     }
 
+    /// Reset `Writer` to its initial state, retaining any allocations.
+    ///
+    /// Why not just implement `Recyclable` for `Writer`? By design,
+    /// `Recyclable::recycle` requires ownership of the value, not just
+    /// `&mut`; see the trait documentation. But we need to use this method
+    /// from functions like `Writer::write`, which only have `&mut Writer`.
+    /// Workarounds include unsafe code (`std::ptr::read`, then `write`, ugh)
+    /// or something like a `Default` impl that returns an oddly-initialized
+    /// `Writer`, which is worse.
+    fn reset(&mut self) {
+        use super::recyclable::Recyclable;
+        use std::mem::take;
+
+        let mut id_gen = IdGenerator::default();
+        let gl450_ext_inst_id = id_gen.next();
+        let void_type = id_gen.next();
+
+        // Every field of the old writer that is not determined by the `Options`
+        // passed to `Writer::new` should be reset somehow.
+        let fresh = Writer {
+            // Copied from the old Writer:
+            flags: self.flags,
+            index_bounds_check_policy: self.index_bounds_check_policy,
+            capabilities: take(&mut self.capabilities),
+            forbidden_caps: take(&mut self.forbidden_caps),
+
+            // Initialized afresh:
+            id_gen,
+            void_type,
+            gl450_ext_inst_id,
+
+            // Recycled:
+            physical_layout: self.physical_layout.clone().recycle(),
+            logical_layout: take(&mut self.logical_layout).recycle(),
+            debugs: take(&mut self.debugs).recycle(),
+            annotations: take(&mut self.annotations).recycle(),
+            lookup_type: take(&mut self.lookup_type).recycle(),
+            lookup_function: take(&mut self.lookup_function).recycle(),
+            lookup_function_type: take(&mut self.lookup_function_type).recycle(),
+            lookup_function_call: take(&mut self.lookup_function_call).recycle(),
+            constant_ids: take(&mut self.constant_ids).recycle(),
+            cached_constants: take(&mut self.cached_constants).recycle(),
+            global_variables: take(&mut self.global_variables).recycle(),
+            cached: take(&mut self.cached).recycle(),
+            temp_list: take(&mut self.temp_list).recycle(),
+        };
+
+        *self = fresh;
+    }
+
     fn check(&mut self, capabilities: &[spirv::Capability]) -> Result<(), Error> {
         if capabilities.is_empty()
             || capabilities
@@ -2783,7 +2834,6 @@ impl Writer {
                 .push(Instruction::source(spirv::SourceLanguage::GLSL, 450));
         }
 
-        self.constant_ids.clear();
         self.constant_ids.resize(ir_module.constants.len(), 0);
         // first, output all the scalar constants
         for (handle, constant) in ir_module.constants.iter() {
@@ -2826,7 +2876,6 @@ impl Writer {
         debug_assert_eq!(self.constant_ids.iter().position(|&id| id == 0), None);
 
         // now write all globals
-        self.global_variables.clear();
         for (_, var) in ir_module.global_variables.iter() {
             let (instruction, id) = self.write_global_variable(ir_module, var)?;
             instruction.to_words(&mut self.logical_layout.declarations);
@@ -2884,9 +2933,7 @@ impl Writer {
         info: &ModuleInfo,
         words: &mut Vec<Word>,
     ) -> Result<(), Error> {
-        self.lookup_function.clear();
-        self.lookup_function_type.clear();
-        self.lookup_function_call.clear();
+        self.reset();
 
         self.write_logical_layout(ir_module, info)?;
         self.write_physical_layout();