Reusing the same vector for all the nodes.

Author: lvella

src/loader/dag/const_collapse.rs

@@ -1,6 +1,5 @@
 use std::{cell::Cell, collections::HashMap};
 
-use itertools::Itertools;
 use wasmparser::Operator;
 
 use crate::loader::{
@@ -12,23 +11,25 @@ use crate::loader::{
 /// This is an optional optimization pass that collapses constants into
 /// the node that uses them, if the ISA supports it.
 ///
-/// For example, an `i32.add` instruction that takes a constant
-/// `5` as input could be collapsed into an `i32.add_imm` instruction
-/// that takes `5` as an immediate operand.
+/// For example, if the output ISA is RISC-V, in an `i32.add` instruction that
+/// has a constant node `5` as input, the value `5` could be collapsed into an
+/// immediate for `addi` instruction.
 ///
 /// This pass only happens if `Settings::get_const_collapse_processor()` is
 /// implemented by the user.
 ///
 /// Returns the number of constants collapsed.
 pub fn constant_collapse<'a, S: Settings<'a>>(settings: &S, dag: &mut Dag) -> usize {
     if let Some(processor) = settings.get_const_collapse_processor() {
-        recursive_constant_collapse(&processor, &mut dag.nodes, HashMap::new())
+        let mut finder = ConstantInputFinder::new();
+        recursive_constant_collapse(&mut finder, &processor, &mut dag.nodes, HashMap::new())
     } else {
         0
     }
 }
 
 fn recursive_constant_collapse(
+    finder: &mut ConstantInputFinder,
     processor: &impl Fn(&Operator, &[MaybeConstant]),
     nodes: &mut [dag::Node],
     const_block_inputs: HashMap<u32, WasmValue>,
@@ -43,23 +44,24 @@ fn recursive_constant_collapse(
 
         match &mut node.operation {
             dag::Operation::WASMOp(operator) => {
-                let (has_constant, input_constants) =
-                    find_constant_inputs(&node.inputs, prev_nodes, const_block_inputs.clone());
+                let (has_constant, input_constants) = finder.find_constant_inputs(
+                    &node.inputs,
+                    prev_nodes,
+                    const_block_inputs.clone(),
+                );
                 if has_constant {
                     // Call the user-provided processor to potentially modify the input constants.
-                    processor(operator, &input_constants);
+                    processor(operator, input_constants);
 
                     // Collapse any inputs that were marked for collapsing.
-                    for (input, maybe_const) in
-                        node.inputs.iter_mut().zip(input_constants.into_iter())
-                    {
+                    for (input, maybe_const) in node.inputs.iter_mut().zip(input_constants) {
                         if let MaybeConstant::ReferenceConstant {
                             value,
                             must_collapse,
                         } = maybe_const
                             && must_collapse.get()
                         {
-                            *input = NodeInput::Constant(value);
+                            *input = NodeInput::Constant(value.clone());
                             num_collapsed += 1;
                         }
                     }
@@ -72,27 +74,34 @@ fn recursive_constant_collapse(
                 // In order to collapse constants that come from the block inputs,
                 // we need to know which block inputs are constants.
                 let const_inputs = if matches!(kind, BlockKind::Block) {
-                    find_constant_inputs(&node.inputs, prev_nodes, const_block_inputs.clone())
+                    finder
+                        .find_constant_inputs(&node.inputs, prev_nodes, const_block_inputs.clone())
                         .1
-                        .into_iter()
+                        .iter()
                         .enumerate()
                         .filter_map(|(idx, maybe_const)| match maybe_const {
                             MaybeConstant::ReferenceConstant { value, .. } => {
-                                Some((idx as u32, value))
+                                Some((idx as u32, value.clone()))
+                            }
+                            MaybeConstant::CollapsedConstant(value) => {
+                                Some((idx as u32, value.clone()))
                             }
-                            MaybeConstant::CollapsedConstant(value) => Some((idx as u32, value)),
                             MaybeConstant::NonConstant => None,
                         })
                         .collect()
                 } else {
                     // Loops have multiple entry points, so we can't be sure a constant input
-                    // is always constant. Besides, an optimized WASM wouldn't have constant
-                    // inputs to loops anyway.
+                    // at entry will be always constant. Besides, an optimized WASM wouldn't
+                    // have constant inputs to loops anyway.
                     HashMap::new()
                 };
 
-                num_collapsed +=
-                    recursive_constant_collapse(processor, &mut sub_dag.nodes, const_inputs);
+                num_collapsed += recursive_constant_collapse(
+                    finder,
+                    processor,
+                    &mut sub_dag.nodes,
+                    const_inputs,
+                );
             }
             _ => {
                 // Constant collapse is not supported for other node types.
@@ -103,17 +112,28 @@ fn recursive_constant_collapse(
     num_collapsed
 }
 
-fn find_constant_inputs(
-    inputs: &[NodeInput],
-    nodes: &mut [dag::Node],
-    const_block_inputs: HashMap<u32, WasmValue>,
-) -> (bool, Vec<MaybeConstant>) {
-    let mut has_constant = false;
+struct ConstantInputFinder {
+    /// This buffer will be reused to avoid allocations on each node.
+    buffer: Vec<MaybeConstant>,
+}
 
-    // Check if any of the inputs are constants.
-    let input_constants = inputs
-        .iter()
-        .map(|input| {
+impl ConstantInputFinder {
+    fn new() -> Self {
+        Self { buffer: Vec::new() }
+    }
+
+    /// Finds which of the given inputs are constants.
+    fn find_constant_inputs<'a>(
+        &'a mut self,
+        inputs: &[NodeInput],
+        nodes: &mut [dag::Node],
+        const_block_inputs: HashMap<u32, WasmValue>,
+    ) -> (bool, &'a mut [MaybeConstant]) {
+        let mut has_constant = false;
+
+        // Check if any of the inputs are constants.
+        self.buffer.clear();
+        self.buffer.extend(inputs.iter().map(|input| {
             let origin = match input {
                 NodeInput::Reference(value_origin) => value_origin,
                 NodeInput::Constant(wasm_value) => {
@@ -138,8 +158,8 @@ fn find_constant_inputs(
                     must_collapse: Cell::new(false),
                 }
             })
-        })
-        .collect_vec();
+        }));
 
-    (has_constant, input_constants)
+        (has_constant, &mut self.buffer)
+    }
 }