Add new remove_phis pass, and fix empty_blocks.

fuzz/fuzz_targets/roundtrip.rs

@@ -1,7 +1,7 @@
 #![no_main]
 use libfuzzer_sys::fuzz_target;
 
-use waffle::{FrontendError, FrontendOptions, Module};
+use waffle::{FrontendError, FrontendOptions, Fuel, Module};
 
 fuzz_target!(|module: wasm_smith::Module| {
     let _ = env_logger::try_init();
@@ -26,6 +26,6 @@ fuzz_target!(|module: wasm_smith::Module| {
             }
         };
     parsed_module.expand_all_funcs().unwrap();
-    parsed_module.per_func_body(|body| body.optimize());
+    parsed_module.per_func_body(|body| body.optimize(&mut Fuel::infinite()));
     let _ = parsed_module.to_wasm_bytes();
 });

src/ir/func.rs

@@ -134,9 +134,10 @@ impl FunctionBody {
 
     pub fn optimize(&mut self, fuel: &mut Fuel) {
         let cfg = crate::cfg::CFGInfo::new(self);
+        crate::passes::remove_phis::run(self, &cfg, fuel);
         crate::passes::basic_opt::gvn(self, &cfg, fuel);
-        crate::passes::resolve_aliases::run(self);
+        crate::passes::remove_phis::run(self, &cfg, fuel);
-        crate::passes::ssa::run(self, &cfg);
+        crate::passes::empty_blocks::run(self, fuel);
     }
 
     pub fn convert_to_max_ssa(&mut self) {

src/lib.rs

@@ -14,6 +14,7 @@ mod ir;
 mod op_traits;
 mod ops;
 pub mod passes;
+pub use passes::Fuel;
 mod scoped_map;
 
 pub use errors::*;

src/passes.rs

@@ -4,6 +4,7 @@ pub mod basic_opt;
 pub mod dom_pass;
 pub mod empty_blocks;
 pub mod maxssa;
+pub mod remove_phis;
 pub mod resolve_aliases;
 pub mod ssa;
 pub mod trace;
@@ -11,9 +12,11 @@ pub mod trace;
 #[derive(Clone, Debug)]
 pub struct Fuel {
     pub remaining: u64,
+    pub consumed: u64,
 }
 impl Fuel {
     pub fn consume(&mut self) -> bool {
+        self.consumed += 1;
         if self.remaining == u64::MAX {
             return true;
         }
@@ -26,6 +29,7 @@ impl Fuel {
     }
     pub fn infinite() -> Fuel {
         Fuel {
+            consumed: 0,
             remaining: u64::MAX,
         }
     }

src/passes/basic_opt.rs

@@ -51,7 +51,7 @@ impl<'a> GVNPass<'a> {
             i += 1;
             if value_is_pure(inst, body) {
                 let mut value = body.values[inst].clone();
-                value.update_uses(|val| *val = body.resolve_alias(*val));
+                value.update_uses(|val| *val = body.resolve_and_update_alias(*val));
 
                 if let ValueDef::Operator(op, args, ..) = &value {
                     let arg_values = args

src/passes/empty_blocks.rs

@@ -1,112 +1,56 @@
 //! Pass to remove empty blocks.
 
+use super::Fuel;
 use crate::entity::EntityRef;
-use crate::ir::{Block, BlockTarget, FunctionBody, Terminator, Value, ValueDef};
+use crate::ir::{Block, BlockTarget, FunctionBody, Terminator};
-use std::borrow::Cow;
-use std::collections::HashSet;
 
-#[derive(Clone, Debug)]
+/// Determines whether a block (i) has no blockparams, and (ii) is
-struct Forwarding {
+/// solely a jump to another block. We can remove these blocks.
-    to: Block,
+///
-    args: Vec<ForwardingArg>,
+/// Why can't we remove blocks that are solely jumps but *do* have
-}
+/// blockparams? Because They still serve a purpose in SSA: they
-
+/// define these blockparams as a join of multiple possible other
-#[derive(Clone, Copy, Debug)]
+/// definitions in preds.
-enum ForwardingArg {
+fn block_is_empty_jump(body: &FunctionBody, block: Block) -> Option<BlockTarget> {
-    BlockParam(usize),
+    // Must be empty except for terminator, and must have no
-    Value(Value),
+    // blockparams, and must have an unconditional-branch terminator.
-}
-
-impl Forwarding {
-    fn compose(a: &Forwarding, b: &Forwarding) -> Forwarding {
-        // `b` should be the target of `a.to`, but we can't assert
-        // that here. The composed target is thus `b.to`.
-        let to = b.to;
-
-        // For each arg in `b.args`, evaluate, replacing any
-        // `BlockParam` with the corresponding value from `a.args`.
-        let args = b
-            .args
-            .iter()
-            .map(|&arg| match arg {
-                ForwardingArg::BlockParam(idx) => a.args[idx].clone(),
-                ForwardingArg::Value(v) => ForwardingArg::Value(v),
-            })
-            .collect::<Vec<_>>();
-
-        Forwarding { to, args }
-    }
-}
-
-fn block_to_forwarding(body: &FunctionBody, block: Block) -> Option<Forwarding> {
-    // Must be empty except for terminator, and must have an
-    // unconditional-branch terminator.
     if body.blocks[block].insts.len() > 0 {
         return None;
     }
+    if body.blocks[block].params.len() > 0 {
+        return None;
+    }
     let target = match &body.blocks[block].terminator {
         &Terminator::Br { ref target } => target,
         _ => return None,
     };
 
-    // If conditions met, then gather ForwardingArgs.
+    Some(target.clone())
-    let args = target
-        .args
-        .iter()
-        .map(|&arg| {
-            let arg = body.resolve_alias(arg);
-            match &body.values[arg] {
-                &ValueDef::BlockParam(param_block, index, _) if param_block == block => {
-                    ForwardingArg::BlockParam(index)
-                }
-                _ => ForwardingArg::Value(arg),
-            }
-        })
-        .collect::<Vec<_>>();
-
-    Some(Forwarding {
-        to: target.block,
-        args,
-    })
 }
 
-fn rewrite_target(forwardings: &[Option<Forwarding>], target: &BlockTarget) -> Option<BlockTarget> {
+fn rewrite_target(
-    if !forwardings[target.block.index()].is_some() {
+    forwardings: &[Option<BlockTarget>],
+    target: &BlockTarget,
+) -> Option<BlockTarget> {
+    if target.args.len() > 0 {
         return None;
     }
-
+    forwardings[target.block.index()].clone()
-    let mut forwarding = Cow::Borrowed(forwardings[target.block.index()].as_ref().unwrap());
-    let mut seen = HashSet::new();
-    while forwardings[forwarding.to.index()].is_some() && seen.insert(forwarding.to.index()) {
-        forwarding = Cow::Owned(Forwarding::compose(
-            &forwarding,
-            forwardings[forwarding.to.index()].as_ref().unwrap(),
-        ));
-    }
-
-    let args = forwarding
-        .args
-        .iter()
-        .map(|arg| match arg {
-            &ForwardingArg::Value(v) => v,
-            &ForwardingArg::BlockParam(idx) => target.args[idx],
-        })
-        .collect::<Vec<_>>();
-
-    Some(BlockTarget {
-        block: forwarding.to,
-        args,
-    })
 }
 
-pub fn run(body: &mut FunctionBody) {
+pub fn run(body: &mut FunctionBody, fuel: &mut Fuel) {
+    log::trace!(
+        "empty_blocks: running on func:\n{}\n",
+        body.display_verbose("| ", None)
+    );
+
     // Identify empty blocks, and to where they should forward.
     let forwardings = body
         .blocks
         .iter()
         .map(|block| {
             if block != body.entry {
-                block_to_forwarding(body, block)
+                block_is_empty_jump(body, block)
             } else {
                 None
             }
@@ -118,11 +62,19 @@ pub fn run(body: &mut FunctionBody) {
     for block_data in body.blocks.values_mut() {
         block_data.terminator.update_targets(|target| {
             if let Some(new_target) = rewrite_target(&forwardings[..], target) {
-                *target = new_target;
+                if fuel.consume() {
+                    log::trace!("empty_blocks: replacing {:?} with {:?}", target, new_target);
+                    *target = new_target;
+                }
             }
         });
     }
 
     // Recompute preds/succs.
     body.recompute_edges();
+
+    log::trace!(
+        "empty_blocks: finished:\n{}\n",
+        body.display_verbose("| ", None)
+    );
 }

src/passes/remove_phis.rs

@@ -0,0 +1,113 @@
+//! Remove-useless-phis (blockparams) pass.
+
+use super::Fuel;
+use crate::cfg::CFGInfo;
+use crate::ir::*;
+
+fn all_equal(mut vals: impl Iterator<Item = Value>) -> Option<Value> {
+    match vals.next() {
+        Some(val) if vals.all(|other_val| other_val == val) => Some(val),
+        _ => None,
+    }
+}
+
+fn delete_indices<T: Copy>(vec: &mut Vec<T>, indices: &[usize]) {
+    let mut out = 0;
+    let mut indices_idx = 0;
+    for i in 0..vec.len() {
+        if indices_idx < indices.len() && indices[indices_idx] == i {
+            indices_idx += 1;
+            // Deleted!
+        } else {
+            if out < i {
+                vec[out] = vec[i];
+            }
+            out += 1;
+        }
+    }
+    if out < vec.len() {
+        vec.truncate(out);
+    }
+}
+
+pub fn run(func: &mut FunctionBody, cfg: &CFGInfo, fuel: &mut Fuel) {
+    // For every block, collect the arg-lists of preds. If a given
+    // blockparam has all the same values for an arg, replace the
+    // blockparam value with an alias to that one value, and then
+    // remove it from the blockparams and target-lists of all preds.
+
+    log::trace!(
+        "remove_phis: running on func:\n{}\n",
+        func.display_verbose("| ", None)
+    );
+
+    let mut deleted = vec![];
+    for &block in cfg.rpo.values() {
+        // Skip the entry block -- we can't remove any args, because
+        // there is also an implicit in-edge from the function entry
+        // with arguments.
+        if block == func.entry {
+            continue;
+        }
+
+        deleted.clear();
+
+        // Gather arg-lists from each pred's terminator.
+        let mut arglists = vec![];
+        for (i, &pred) in func.blocks[block].preds.iter().enumerate() {
+            let pos = func.blocks[block].pos_in_pred_succ[i];
+            func.blocks[pred].terminator.visit_target(pos, |target| {
+                assert_eq!(target.block, block);
+                assert_eq!(target.args.len(), func.blocks[block].params.len());
+                arglists.push(target.args.clone());
+            });
+        }
+
+        // For each arg-position, check if all args are the same. If
+        // so, rewrite value and mark index as deleted.
+        for i in 0..func.blocks[block].params.len() {
+            let blockparam = func.blocks[block].params[i].1;
+            let same = all_equal(
+                arglists
+                    .iter()
+                    .map(|arglist| func.resolve_alias(arglist[i])),
+            );
+            if let Some(val) = same {
+                if !fuel.consume() {
+                    continue;
+                }
+                if val != blockparam {
+                    log::trace!(
+                        "deleting blockparam {} from block {}: now {}",
+                        blockparam,
+                        block,
+                        val
+                    );
+                    func.values[blockparam] = ValueDef::Alias(val);
+                    deleted.push(i);
+                }
+            }
+        }
+
+        // If anything was deleted, remove the appropriate indices in
+        // the func's blockparams list and in targets' arg lists.
+        if !deleted.is_empty() {
+            delete_indices(&mut func.blocks[block].params, &deleted[..]);
+            for i in 0..func.blocks[block].preds.len() {
+                let pred = func.blocks[block].preds[i];
+                let pos = func.blocks[block].pos_in_pred_succ[i];
+                func.blocks[pred].terminator.update_target(pos, |target| {
+                    delete_indices(&mut target.args, &deleted[..]);
+                });
+            }
+
+            // Renumber blockparam values.
+            for (i, &(_, param)) in func.blocks[block].params.iter().enumerate() {
+                let ty = func.values[param].ty().unwrap();
+                func.values[param] = ValueDef::BlockParam(block, i, ty);
+            }
+        }
+    }
+
+    log::trace!("remove_phis: done:\n{}\n", func.display_verbose("| ", None));
+}

src/passes/ssa.rs

@@ -25,13 +25,34 @@ impl DefBlocks {
 pub fn run(body: &FunctionBody, cfg: &CFGInfo) {
     let def_blocks = DefBlocks::compute(body);
 
-    for (block, data) in body.blocks.entries() {
+    // Visit only reachable blocks.
+    for &block in cfg.rpo.values() {
+        let data = &body.blocks[block];
         let validate = |value| {
             let value = body.resolve_alias(value);
             let def_block = def_blocks.def_block[value];
-            assert!(cfg.dominates(def_block, block));
+            assert!(
+                cfg.dominates(def_block, block),
+                "value {} defined in block {} used in block {}: def does not dominate use",
+                value,
+                def_block,
+                block
+            );
         };
 
+        for (i, &(_, param)) in data.params.iter().enumerate() {
+            match &body.values[param] {
+                &ValueDef::BlockParam(param_block, param_idx, _) => {
+                    assert_eq!(param_block, block);
+                    assert_eq!(param_idx, i);
+                }
+                _ => panic!(
+                    "Bad blockparam value for param {} of {} ({}): {:?}",
+                    i, block, param, body.values[param]
+                ),
+            }
+        }
+
         for &inst in &data.insts {
             match &body.values[inst] {
                 &ValueDef::Operator(_, ref args, _) => {