From 2a47a77cdc795253ab42184e49f796f3ae96d78b Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Thu, 3 Nov 2022 00:20:58 -0700
Subject: [PATCH] WIP.

---
 src/ir/display.rs       |   3 +
 src/ir/module.rs        |   2 +
 src/op_traits.rs        | 138 +++++++++++++++++++++-------------------
 src/passes/basic_opt.rs |  53 +++++++++++++--
 src/scoped_map.rs       |  19 ++++--
 5 files changed, 139 insertions(+), 76 deletions(-)

diff --git a/src/ir/display.rs b/src/ir/display.rs
index ab13671..95fae3a 100644
--- a/src/ir/display.rs
+++ b/src/ir/display.rs
@@ -119,6 +119,9 @@ impl<'a> Display for FunctionBodyDisplay<'a> {
                     ValueDef::PickOutput(val, idx, ty) => {
                         writeln!(f, "{}    {} = {}.{} # {}", self.1, inst, val, idx, ty)?;
                     }
+                    ValueDef::Alias(val) => {
+                        writeln!(f, "{}    {} = {}", self.1, inst, val)?;
+                    }
                     _ => unreachable!(),
                 }
             }
diff --git a/src/ir/module.rs b/src/ir/module.rs
index 77814f9..9730827 100644
--- a/src/ir/module.rs
+++ b/src/ir/module.rs
@@ -209,6 +209,8 @@ impl<'a> Module<'a> {
         for func_decl in module.funcs.values_mut() {
             if let Some(body) = func_decl.body_mut() {
                 crate::passes::rpo::run(body);
+                let cfg = crate::cfg::CFGInfo::new(body);
+                crate::passes::basic_opt::gvn(body, &cfg);
                 crate::passes::resolve_aliases::run(body);
             }
         }
diff --git a/src/op_traits.rs b/src/op_traits.rs
index 71a3954..0d50ae5 100644
--- a/src/op_traits.rs
+++ b/src/op_traits.rs
@@ -456,26 +456,26 @@ pub enum SideEffect {
     All,
 }
 
-pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
+pub fn op_effects(op: &Operator) -> Cow<'static, [SideEffect]> {
     use SideEffect::*;
 
     match op {
-        &Operator::Unreachable => Ok(Cow::Borrowed(&[Trap])),
-        &Operator::Nop => Ok(Cow::Borrowed(&[])),
+        &Operator::Unreachable => Cow::Borrowed(&[Trap]),
+        &Operator::Nop => Cow::Borrowed(&[]),
 
-        &Operator::Call { .. } => Ok(Cow::Borrowed(&[All])),
-        &Operator::CallIndirect { .. } => Ok(Cow::Borrowed(&[All])),
-        &Operator::Return => Ok(Cow::Borrowed(&[Return])),
-        &Operator::LocalSet { local_index, .. } => Ok(vec![WriteLocal(local_index)].into()),
-        &Operator::LocalGet { local_index, .. } => Ok(vec![ReadLocal(local_index)].into()),
+        &Operator::Call { .. } => Cow::Borrowed(&[All]),
+        &Operator::CallIndirect { .. } => Cow::Borrowed(&[All]),
+        &Operator::Return => Cow::Borrowed(&[Return]),
+        &Operator::LocalSet { local_index, .. } => vec![WriteLocal(local_index)].into(),
+        &Operator::LocalGet { local_index, .. } => vec![ReadLocal(local_index)].into(),
         &Operator::LocalTee { local_index, .. } => {
-            Ok(vec![ReadLocal(local_index), WriteLocal(local_index)].into())
+            vec![ReadLocal(local_index), WriteLocal(local_index)].into()
         }
 
-        &Operator::Select => Ok(Cow::Borrowed(&[])),
-        &Operator::TypedSelect { .. } => Ok(Cow::Borrowed(&[])),
-        &Operator::GlobalGet { global_index, .. } => Ok(vec![ReadGlobal(global_index)].into()),
-        &Operator::GlobalSet { global_index, .. } => Ok(vec![WriteGlobal(global_index)].into()),
+        &Operator::Select => Cow::Borrowed(&[]),
+        &Operator::TypedSelect { .. } => Cow::Borrowed(&[]),
+        &Operator::GlobalGet { global_index, .. } => vec![ReadGlobal(global_index)].into(),
+        &Operator::GlobalSet { global_index, .. } => vec![WriteGlobal(global_index)].into(),
 
         Operator::I32Load { .. }
         | Operator::I32Load8S { .. }
@@ -490,7 +490,7 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::I64Load32S { .. }
         | Operator::I64Load32U { .. }
         | Operator::F32Load { .. }
-        | Operator::F64Load { .. } => Ok(Cow::Borrowed(&[Trap, ReadMem])),
+        | Operator::F64Load { .. } => Cow::Borrowed(&[Trap, ReadMem]),
 
         Operator::I32Store { .. }
         | Operator::I64Store { .. }
@@ -500,12 +500,12 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::I32Store16 { .. }
         | Operator::I64Store8 { .. }
         | Operator::I64Store16 { .. }
-        | Operator::I64Store32 { .. } => Ok(Cow::Borrowed(&[Trap, WriteMem])),
+        | Operator::I64Store32 { .. } => Cow::Borrowed(&[Trap, WriteMem]),
 
         Operator::I32Const { .. }
         | Operator::I64Const { .. }
         | Operator::F32Const { .. }
-        | Operator::F64Const { .. } => Ok(Cow::Borrowed(&[])),
+        | Operator::F64Const { .. } => Cow::Borrowed(&[]),
 
         Operator::I32Eqz
         | Operator::I32Eq
@@ -540,7 +540,7 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::F64Lt
         | Operator::F64Gt
         | Operator::F64Le
-        | Operator::F64Ge => Ok(Cow::Borrowed(&[])),
+        | Operator::F64Ge => Cow::Borrowed(&[]),
 
         Operator::I32Clz
         | Operator::I32Ctz
@@ -555,10 +555,10 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::I32ShrS
         | Operator::I32ShrU
         | Operator::I32Rotl
-        | Operator::I32Rotr => Ok(Cow::Borrowed(&[])),
+        | Operator::I32Rotr => Cow::Borrowed(&[]),
 
         Operator::I32DivS | Operator::I32DivU | Operator::I32RemS | Operator::I32RemU => {
-            Ok(Cow::Borrowed(&[Trap]))
+            Cow::Borrowed(&[Trap])
         }
 
         Operator::I64Clz
@@ -574,10 +574,10 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::I64ShrS
         | Operator::I64ShrU
         | Operator::I64Rotl
-        | Operator::I64Rotr => Ok(Cow::Borrowed(&[])),
+        | Operator::I64Rotr => Cow::Borrowed(&[]),
 
         Operator::I64DivS | Operator::I64DivU | Operator::I64RemS | Operator::I64RemU => {
-            Ok(Cow::Borrowed(&[Trap]))
+            Cow::Borrowed(&[Trap])
         }
 
         Operator::F32Abs
@@ -593,7 +593,7 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::F32Div
         | Operator::F32Min
         | Operator::F32Max
-        | Operator::F32Copysign => Ok(Cow::Borrowed(&[])),
+        | Operator::F32Copysign => Cow::Borrowed(&[]),
 
         Operator::F64Abs
         | Operator::F64Neg
@@ -608,55 +608,59 @@ pub fn op_effects(op: &Operator) -> Result<Cow<'static, [SideEffect]>> {
         | Operator::F64Div
         | Operator::F64Min
         | Operator::F64Max
-        | Operator::F64Copysign => Ok(Cow::Borrowed(&[])),
+        | Operator::F64Copysign => Cow::Borrowed(&[]),
 
-        Operator::I32WrapI64 => Ok(Cow::Borrowed(&[])),
-        Operator::I32TruncF32S => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I32TruncF32U => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I32TruncF64S => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I32TruncF64U => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I64ExtendI32S => Ok(Cow::Borrowed(&[])),
-        Operator::I64ExtendI32U => Ok(Cow::Borrowed(&[])),
-        Operator::I64TruncF32S => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I64TruncF32U => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I64TruncF64S => Ok(Cow::Borrowed(&[Trap])),
-        Operator::I64TruncF64U => Ok(Cow::Borrowed(&[Trap])),
-        Operator::F32ConvertI32S => Ok(Cow::Borrowed(&[])),
-        Operator::F32ConvertI32U => Ok(Cow::Borrowed(&[])),
-        Operator::F32ConvertI64S => Ok(Cow::Borrowed(&[])),
-        Operator::F32ConvertI64U => Ok(Cow::Borrowed(&[])),
-        Operator::F32DemoteF64 => Ok(Cow::Borrowed(&[])),
-        Operator::F64ConvertI32S => Ok(Cow::Borrowed(&[])),
-        Operator::F64ConvertI32U => Ok(Cow::Borrowed(&[])),
-        Operator::F64ConvertI64S => Ok(Cow::Borrowed(&[])),
-        Operator::F64ConvertI64U => Ok(Cow::Borrowed(&[])),
-        Operator::F64PromoteF32 => Ok(Cow::Borrowed(&[])),
-        Operator::I32Extend8S => Ok(Cow::Borrowed(&[])),
-        Operator::I32Extend16S => Ok(Cow::Borrowed(&[])),
-        Operator::I64Extend8S => Ok(Cow::Borrowed(&[])),
-        Operator::I64Extend16S => Ok(Cow::Borrowed(&[])),
-        Operator::I64Extend32S => Ok(Cow::Borrowed(&[])),
-        Operator::I32TruncSatF32S => Ok(Cow::Borrowed(&[])),
-        Operator::I32TruncSatF32U => Ok(Cow::Borrowed(&[])),
-        Operator::I32TruncSatF64S => Ok(Cow::Borrowed(&[])),
-        Operator::I32TruncSatF64U => Ok(Cow::Borrowed(&[])),
-        Operator::I64TruncSatF32S => Ok(Cow::Borrowed(&[])),
-        Operator::I64TruncSatF32U => Ok(Cow::Borrowed(&[])),
-        Operator::I64TruncSatF64S => Ok(Cow::Borrowed(&[])),
-        Operator::I64TruncSatF64U => Ok(Cow::Borrowed(&[])),
-        Operator::F32ReinterpretI32 => Ok(Cow::Borrowed(&[])),
-        Operator::F64ReinterpretI64 => Ok(Cow::Borrowed(&[])),
-        Operator::I32ReinterpretF32 => Ok(Cow::Borrowed(&[])),
-        Operator::I64ReinterpretF64 => Ok(Cow::Borrowed(&[])),
-        Operator::TableGet { table_index, .. } => Ok(vec![ReadTable(*table_index), Trap].into()),
-        Operator::TableSet { table_index, .. } => Ok(vec![WriteTable(*table_index), Trap].into()),
-        Operator::TableGrow { table_index, .. } => Ok(vec![WriteTable(*table_index), Trap].into()),
-        Operator::TableSize { table_index, .. } => Ok(vec![ReadTable(*table_index)].into()),
-        Operator::MemorySize { .. } => Ok(Cow::Borrowed(&[ReadMem])),
-        Operator::MemoryGrow { .. } => Ok(Cow::Borrowed(&[WriteMem, Trap])),
+        Operator::I32WrapI64 => Cow::Borrowed(&[]),
+        Operator::I32TruncF32S => Cow::Borrowed(&[Trap]),
+        Operator::I32TruncF32U => Cow::Borrowed(&[Trap]),
+        Operator::I32TruncF64S => Cow::Borrowed(&[Trap]),
+        Operator::I32TruncF64U => Cow::Borrowed(&[Trap]),
+        Operator::I64ExtendI32S => Cow::Borrowed(&[]),
+        Operator::I64ExtendI32U => Cow::Borrowed(&[]),
+        Operator::I64TruncF32S => Cow::Borrowed(&[Trap]),
+        Operator::I64TruncF32U => Cow::Borrowed(&[Trap]),
+        Operator::I64TruncF64S => Cow::Borrowed(&[Trap]),
+        Operator::I64TruncF64U => Cow::Borrowed(&[Trap]),
+        Operator::F32ConvertI32S => Cow::Borrowed(&[]),
+        Operator::F32ConvertI32U => Cow::Borrowed(&[]),
+        Operator::F32ConvertI64S => Cow::Borrowed(&[]),
+        Operator::F32ConvertI64U => Cow::Borrowed(&[]),
+        Operator::F32DemoteF64 => Cow::Borrowed(&[]),
+        Operator::F64ConvertI32S => Cow::Borrowed(&[]),
+        Operator::F64ConvertI32U => Cow::Borrowed(&[]),
+        Operator::F64ConvertI64S => Cow::Borrowed(&[]),
+        Operator::F64ConvertI64U => Cow::Borrowed(&[]),
+        Operator::F64PromoteF32 => Cow::Borrowed(&[]),
+        Operator::I32Extend8S => Cow::Borrowed(&[]),
+        Operator::I32Extend16S => Cow::Borrowed(&[]),
+        Operator::I64Extend8S => Cow::Borrowed(&[]),
+        Operator::I64Extend16S => Cow::Borrowed(&[]),
+        Operator::I64Extend32S => Cow::Borrowed(&[]),
+        Operator::I32TruncSatF32S => Cow::Borrowed(&[]),
+        Operator::I32TruncSatF32U => Cow::Borrowed(&[]),
+        Operator::I32TruncSatF64S => Cow::Borrowed(&[]),
+        Operator::I32TruncSatF64U => Cow::Borrowed(&[]),
+        Operator::I64TruncSatF32S => Cow::Borrowed(&[]),
+        Operator::I64TruncSatF32U => Cow::Borrowed(&[]),
+        Operator::I64TruncSatF64S => Cow::Borrowed(&[]),
+        Operator::I64TruncSatF64U => Cow::Borrowed(&[]),
+        Operator::F32ReinterpretI32 => Cow::Borrowed(&[]),
+        Operator::F64ReinterpretI64 => Cow::Borrowed(&[]),
+        Operator::I32ReinterpretF32 => Cow::Borrowed(&[]),
+        Operator::I64ReinterpretF64 => Cow::Borrowed(&[]),
+        Operator::TableGet { table_index, .. } => vec![ReadTable(*table_index), Trap].into(),
+        Operator::TableSet { table_index, .. } => vec![WriteTable(*table_index), Trap].into(),
+        Operator::TableGrow { table_index, .. } => vec![WriteTable(*table_index), Trap].into(),
+        Operator::TableSize { table_index, .. } => vec![ReadTable(*table_index)].into(),
+        Operator::MemorySize { .. } => Cow::Borrowed(&[ReadMem]),
+        Operator::MemoryGrow { .. } => Cow::Borrowed(&[WriteMem, Trap]),
     }
 }
 
+pub fn is_pure(op: &Operator) -> bool {
+    op_effects(op).is_empty()
+}
+
 impl std::fmt::Display for Operator {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
diff --git a/src/passes/basic_opt.rs b/src/passes/basic_opt.rs
index 37562c3..ebfc9cc 100644
--- a/src/passes/basic_opt.rs
+++ b/src/passes/basic_opt.rs
@@ -1,10 +1,55 @@
 //! Basic optimizations: GVN and constant-propagation/folding.
 
-use crate::ir::*;
-use crate::passes::dom_pass;
-use crate::scoped_map::ScopedMap;
 use crate::cfg::CFGInfo;
+use crate::ir::*;
+use crate::op_traits::is_pure;
+use crate::passes::dom_pass::{dom_pass, DomtreePass};
+use crate::scoped_map::ScopedMap;
 
 pub fn gvn(body: &mut FunctionBody, cfg: &CFGInfo) {
-    let mut map: ScopedMap<ValueDef, Value> = ScopedMap::new();
+    dom_pass::<GVNPass>(body, cfg, &mut GVNPass::default());
+}
+
+#[derive(Clone, Debug, Default)]
+struct GVNPass {
+    map: ScopedMap<ValueDef, Value>,
+}
+
+impl DomtreePass for GVNPass {
+    fn enter(&mut self, block: Block, body: &mut FunctionBody) {
+        self.map.push_level();
+        self.optimize(block, body);
+    }
+
+    fn leave(&mut self, _block: Block, _body: &mut FunctionBody) {
+        self.map.pop_level();
+    }
+}
+
+fn value_is_pure(value: Value, body: &FunctionBody) -> bool {
+    match body.values[value] {
+        ValueDef::Operator(op, ..) if is_pure(&op) => true,
+        _ => false,
+    }
+}
+
+impl GVNPass {
+    fn optimize(&mut self, block: Block, body: &mut FunctionBody) {
+        let mut i = 0;
+        while i < body.blocks[block].insts.len() {
+            let inst = body.blocks[block].insts[i];
+            i += 1;
+            if value_is_pure(inst, body) {
+                let mut value = body.values[inst].clone();
+                value.update_uses(|val| *val = body.resolve_alias(*val));
+                if let Some(value) = self.map.get(&value) {
+                    body.set_alias(inst, *value);
+                    i -= 1;
+                    body.blocks[block].insts.remove(i);
+                } else {
+                    self.map.insert(value, inst);
+                }
+            }
+        }
+    }
 }
diff --git a/src/scoped_map.rs b/src/scoped_map.rs
index 5d9b264..6238e9d 100644
--- a/src/scoped_map.rs
+++ b/src/scoped_map.rs
@@ -1,21 +1,30 @@
 //! Scoped hashmap.
 
 use fxhash::FxHashMap;
+use std::fmt::Debug;
 use std::hash::Hash;
 
-pub struct ScopedMap<K: Hash + Eq, V> {
+#[derive(Clone, Debug)]
+pub struct ScopedMap<K: Hash + Eq + Clone + Debug, V: Clone + Debug> {
     map: FxHashMap<K, ScopedMapEntry<V>>,
     gen: u32,
     gen_by_level: Vec<u32>,
 }
 
-struct ScopedMapEntry<V> {
+impl<K: Hash + Eq + Clone + Debug, V: Clone + Debug> std::default::Default for ScopedMap<K, V> {
+    fn default() -> Self {
+        ScopedMap::new()
+    }
+}
+
+#[derive(Clone, Debug)]
+struct ScopedMapEntry<V: Clone + Debug> {
     gen: u32,
     level: u32,
     value: V,
 }
 
-impl<K: Hash + Eq, V> ScopedMap<K, V> {
+impl<K: Hash + Eq + Clone + Debug, V: Clone + Debug> ScopedMap<K, V> {
     pub fn new() -> ScopedMap<K, V> {
         ScopedMap {
             map: FxHashMap::default(),
@@ -37,8 +46,8 @@ impl<K: Hash + Eq, V> ScopedMap<K, V> {
         self.map.insert(
             k,
             ScopedMapEntry {
-                gen: self.gen,
-                level: self.gen_by_level.len() as u32,
+                gen: *self.gen_by_level.last().unwrap(),
+                level: (self.gen_by_level.len() - 1) as u32,
                 value: v,
             },
         );