Macros and dataflow analysis framework

src/analysis/liveness.rs

@@ -0,0 +1,18 @@
+//! Liveness analysis.
+
+use crate::{backward_pass, dataflow_use_def};
+use crate::{ir::*, pass::*};
+
+backward_pass!(
+    Liveness,
+    UnionBitSet,
+    dataflow_use_def!(
+        UnionBitSet,
+        use: |u, lattice| {
+            lattice.add(u.index() as usize);
+        },
+        def: |d, lattice| {
+            lattice.remove(d.index() as usize);
+        }
+    )
+);

src/analysis/mod.rs

@@ -0,0 +1,4 @@
+//! Analyses.
+
+pub mod liveness;
+use liveness::*;

src/backend/location.rs

@@ -2,8 +2,30 @@
 
 use crate::cfg::CFGInfo;
 use crate::ir::*;
+use fxhash::FxHashMap;
 use wasmparser::Type;
 
+/// Pass to compute reference counts for every value.
+struct UseCounts {
+    counts: FxHashMap<Value, usize>,
+}
+
+impl UseCounts {
+    fn new(f: &FunctionBody) -> UseCounts {
+        let mut counts = FxHashMap::default();
+        for block in &f.blocks {
+            block.visit_uses(|value| {
+                *counts.entry(value).or_insert(0) += 1;
+            });
+        }
+        UseCounts { counts }
+    }
+
+    fn count(&self, value: Value) -> usize {
+        *self.counts.get(&value).unwrap_or(&0)
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Location {
     // Store in a local.

src/frontend.rs

@@ -114,12 +114,12 @@ fn handle_payload<'a>(
     Ok(())
 }
 
-fn parse_body<'a, 'b>(
+fn parse_body<'a>(
-    module: &'b Module<'a>,
+    module: &'a Module,
     my_sig: SignatureId,
-    body: wasmparser::FunctionBody<'a>,
+    body: wasmparser::FunctionBody,
-) -> Result<FunctionBody<'a>> {
+) -> Result<FunctionBody> {
-    let mut ret: FunctionBody<'a> = FunctionBody::default();
+    let mut ret: FunctionBody = FunctionBody::default();
 
     for &param in &module.signatures[my_sig].params[..] {
         ret.locals.push(param);
@@ -168,7 +168,7 @@ fn parse_body<'a, 'b>(
 struct FunctionBodyBuilder<'a, 'b> {
     module: &'b Module<'a>,
     my_sig: SignatureId,
-    body: &'b mut FunctionBody<'a>,
+    body: &'b mut FunctionBody,
     cur_block: Option<BlockId>,
     ctrl_stack: Vec<Frame>,
     op_stack: Vec<(Type, Value)>,
@@ -245,7 +245,7 @@ impl Frame {
 }
 
 impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
-    fn new(module: &'b Module<'a>, my_sig: SignatureId, body: &'b mut FunctionBody<'a>) -> Self {
+    fn new(module: &'b Module<'a>, my_sig: SignatureId, body: &'b mut FunctionBody) -> Self {
         body.blocks.push(Block::default());
         let mut ret = Self {
             module,
@@ -757,6 +757,7 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
     fn create_block(&mut self) -> BlockId {
         let id = self.body.blocks.len() as BlockId;
         self.body.blocks.push(Block::default());
+        self.body.blocks[id].id = id;
         id
     }
 
@@ -928,11 +929,9 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
                 self.op_stack.push((output_ty, Value::inst(block, inst, i)));
             }
 
-            self.body.blocks[block].insts.push(Inst {
+            self.body.blocks[block]
-                operator: op,
+                .insts
-                n_outputs,
+                .push(Inst::make(&op, n_outputs, input_operands));
-                inputs: input_operands,
-            });
         } else {
             let _ = self.pop_n(inputs.len());
             for ty in outputs {

src/ir.rs

@@ -16,19 +16,19 @@ pub const INVALID_BLOCK: BlockId = usize::MAX;
 #[derive(Clone, Debug, Default)]
 pub struct Module<'a> {
     pub orig_bytes: &'a [u8],
-    pub funcs: Vec<FuncDecl<'a>>,
+    pub funcs: Vec<FuncDecl>,
     pub signatures: Vec<FuncType>,
     pub globals: Vec<Type>,
     pub tables: Vec<Type>,
 }
 
 #[derive(Clone, Debug)]
-pub enum FuncDecl<'a> {
+pub enum FuncDecl {
     Import(SignatureId),
-    Body(SignatureId, FunctionBody<'a>),
+    Body(SignatureId, FunctionBody),
 }
 
-impl<'a> FuncDecl<'a> {
+impl FuncDecl {
     pub fn sig(&self) -> SignatureId {
         match self {
             FuncDecl::Import(sig) => *sig,
@@ -38,81 +38,81 @@ impl<'a> FuncDecl<'a> {
 }
 
 #[derive(Clone, Debug, Default)]
-pub struct FunctionBody<'a> {
+pub struct FunctionBody {
     pub arg_values: Vec<Value>,
     pub locals: Vec<Type>,
-    pub blocks: Vec<Block<'a>>,
+    pub blocks: Vec<Block>,
     pub types: FxHashMap<Value, Type>,
 }
 
 #[derive(Clone, Debug, Default)]
-pub struct Block<'a> {
+pub struct Block {
     pub id: BlockId,
     pub params: Vec<Type>,
-    pub insts: Vec<Inst<'a>>,
+    pub insts: Vec<Inst>,
     pub terminator: Terminator,
 }
 
-impl<'a> Block<'a> {
+impl Block {
     pub fn successors(&self) -> Vec<BlockId> {
         self.terminator.successors()
     }
 
-    pub fn values<'b>(&'b self) -> impl Iterator<Item = Value> + 'b {
+    pub fn defs<'b>(&'b self) -> impl Iterator<Item = Value> + 'b {
         let block = self.id;
-        self.insts
+        let param_values = (0..self.params.len()).map(move |i| Value::blockparam(block, i));
+        let inst_values = self
+            .insts
             .iter()
             .enumerate()
             .map(move |(inst_id, inst)| {
                 (0..inst.n_outputs).map(move |i| Value::inst(block, inst_id, i))
             })
-            .flatten()
+            .flatten();
+        param_values.chain(inst_values)
     }
 
-    pub fn visit_values<F: Fn(&Value)>(&self, f: F) {
+    pub fn visit_uses<F: FnMut(Value)>(&self, mut f: F) {
         for inst in &self.insts {
-            for input in &inst.inputs {
+            for &input in &inst.inputs {
                 f(input);
             }
         }
-        match &self.terminator {
+        self.terminator.visit_uses(f);
-            &Terminator::CondBr { ref cond, .. } => f(cond),
-            &Terminator::Select { ref value, .. } => f(value),
-            &Terminator::Return { ref values, .. } => {
-                for value in values {
-                    f(value);
-                }
-            }
-            _ => {}
-        }
     }
 
-    pub fn update_values<F: Fn(&mut Value)>(&mut self, f: F) {
+    pub fn update_uses<F: FnMut(&mut Value)>(&mut self, mut f: F) {
         for inst in &mut self.insts {
             for input in &mut inst.inputs {
                 f(input);
             }
         }
-        match &mut self.terminator {
+        self.terminator.update_uses(f);
-            &mut Terminator::CondBr { ref mut cond, .. } => f(cond),
-            &mut Terminator::Select { ref mut value, .. } => f(value),
-            &mut Terminator::Return { ref mut values, .. } => {
-                for value in values {
-                    f(value);
-                }
-            }
-            _ => {}
-        }
     }
 }
 
 #[derive(Clone, Debug)]
-pub struct Inst<'a> {
+pub struct Inst {
-    pub operator: Operator<'a>,
+    pub operator: Operator<'static>,
     pub n_outputs: usize,
     pub inputs: Vec<Value>,
 }
 
+impl Inst {
+    pub fn make<'a>(operator: &Operator<'a>, n_outputs: usize, inputs: Vec<Value>) -> Self {
+        // The only operator that actually makes use of the lifetime
+        // parameter is BrTable.
+        assert!(!matches!(operator, &Operator::BrTable { .. }));
+        let operator = operator.clone();
+        let operator = unsafe { std::mem::transmute(operator) };
+        Inst {
+            operator,
+            n_outputs,
+            inputs,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Value(u64);
 
@@ -201,6 +201,16 @@ impl Value {
             _ => None,
         }
     }
+
+    pub fn index(self) -> u64 {
+        self.0
+    }
+
+    pub fn from_index(value: u64) -> Value {
+        let tag = value >> VALUE_TAG_SHIFT;
+        assert!(tag < 4);
+        Self(value)
+    }
 }
 
 impl std::fmt::Display for Value {
@@ -330,4 +340,30 @@ impl Terminator {
             Terminator::None => vec![],
         }
     }
+
+    pub fn visit_uses<F: FnMut(Value)>(&self, mut f: F) {
+        match self {
+            &Terminator::CondBr { cond, .. } => f(cond),
+            &Terminator::Select { value, .. } => f(value),
+            &Terminator::Return { ref values, .. } => {
+                for &value in values {
+                    f(value);
+                }
+            }
+            _ => {}
+        }
+    }
+
+    pub fn update_uses<F: FnMut(&mut Value)>(&mut self, mut f: F) {
+        match self {
+            &mut Terminator::CondBr { ref mut cond, .. } => f(cond),
+            &mut Terminator::Select { ref mut value, .. } => f(value),
+            &mut Terminator::Return { ref mut values, .. } => {
+                for value in values {
+                    f(value);
+                }
+            }
+            _ => {}
+        }
+    }
 }

src/lib.rs

@@ -7,6 +7,7 @@
 pub use wasm_encoder;
 pub use wasmparser;
 
+mod analysis;
 mod backend;
 mod cfg;
 mod frontend;

src/pass/dataflow.rs

@@ -6,36 +6,140 @@ use crate::ir::*;
 use crate::pass::Lattice;
 use fxhash::{FxHashMap, FxHashSet};
 use std::collections::hash_map::Entry as HashEntry;
+use std::marker::PhantomData;
 use std::{collections::VecDeque, default::Default};
 use wasmparser::Type;
 
-impl<'a> FunctionBody<'a> {
+impl FunctionBody {
-    fn insts(&self) -> impl Iterator<Item = &Inst<'a>> {
+    fn insts(&self) -> impl Iterator<Item = &Inst> {
         self.blocks.iter().map(|block| block.insts.iter()).flatten()
     }
 }
 
-pub trait DataflowFunctions<L: Lattice> {
+pub trait DataflowFunctions {
-    fn start_block(&self, _lattice: &mut L, _block: BlockId, _param_types: &[Type]) -> bool {
+    type L: Lattice;
-        false
+
-    }
+    fn start_block(&self, _lattice: &mut Self::L, _block: BlockId, _param_types: &[Type]) {}
     fn end_block(
         &self,
-        _lattce: &mut L,
+        _lattce: &mut Self::L,
         _block: BlockId,
         _next: BlockId,
         _terminator: &Terminator,
-    ) -> bool {
+    ) {
-        false
+    }
+    fn instruction(&self, _lattice: &mut Self::L, _block: BlockId, _instid: InstId, _inst: &Inst) {}
+}
+
+pub struct DataflowFunctionsImpl<L, F1, F2, F3> {
+    f1: F1,
+    f2: F2,
+    f3: F3,
+    _phantom: PhantomData<L>,
+}
+
+impl<L, F1, F2, F3> DataflowFunctionsImpl<L, F1, F2, F3> {
+    pub fn new(f1: F1, f2: F2, f3: F3) -> Self {
+        Self {
+            f1,
+            f2,
+            f3,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<L, F1, F2, F3> DataflowFunctions for DataflowFunctionsImpl<L, F1, F2, F3>
+where
+    L: Lattice,
+    F1: Fn(&mut L, BlockId, InstId, &Inst),
+    F2: Fn(&mut L, BlockId, &[Type]),
+    F3: Fn(&mut L, BlockId, BlockId, &Terminator),
+{
+    type L = L;
+    fn instruction(&self, lattice: &mut L, block: BlockId, instid: InstId, inst: &Inst) {
+        (self.f1)(lattice, block, instid, inst);
+    }
+    fn start_block(&self, lattice: &mut L, block: BlockId, params: &[Type]) {
+        (self.f2)(lattice, block, params);
+    }
+    fn end_block(&self, lattice: &mut L, block: BlockId, next: BlockId, terminator: &Terminator) {
+        (self.f3)(lattice, block, next, terminator);
+    }
+}
+
+#[macro_export]
+macro_rules! dataflow {
+    ($latticety:ty,
+     |$lattice:ident, $block:ident, $instid:ident, $inst:ident| { $($body:tt)* }) => {
+        DataflowFunctionsImpl::new(|$lattice:&mut $latticety, $block, $instid, $inst| {
+            $($body)*
+        }, |_, _, _| {}, |_, _, _, _| {})
+    };
+
+    ($latticety:ty,
+     inst: |$lattice1:ident, $block1:ident, $instid1:ident, $inst1:ident| { $($body1:tt)* },
+     start_block: |$lattice2:ident, $block2:ident, $params2:ident| { $($body2:tt)* }) => {
+        DataflowFunctionsImpl::new(|$lattice1:&mut $latticety, $block1, $instid1, $inst1| {
+            $($body1)*
+        },
+        |$lattice2, $block2, $params2| {
+            $($body2)*
+        }, |_, _, _, _| {})
+    };
+
+    ($latticety:ty,
+     inst: |$lattice1:ident, $block1:ident, $instid1:ident, $inst1:ident| { $($body1:tt)* },
+     start_block: |$lattice2:ident, $block2:ident, $params2:ident| { $($body2:tt)* },
+     end_block: |$lattice3:ident, $block3:ident, $next3:ident, $term3:ident| { $($body3:tt)* }) => {
+        DataflowFunctionsImpl::new(|$lattice1:&mut $latticety, $block1, $instid1, $inst1:&Inst| {
+            $($body1)*
+        },
+        |$lattice2:&mut $latticety, $block2, $params2:&[wasmparser::Type]| {
+            $($body2)*
+        },
+        |$lattice3:&mut $latticety, $block3, $next3, $term3:&Terminator| {
+            $($body3)*
+        })
+    };
+}
+
+#[macro_export]
+macro_rules! dataflow_use_def {
+    ($lattice:ty,
+     use: |$use:ident, $uselattice:ident| { $($usebody:tt)* },
+     def: |$def:ident, $deflattice:ident| { $($defbody:tt)* }) => {
+        {
+            $crate::dataflow!(
+                $lattice,
+                inst: |lattice, block, instid, inst| {
+                    let $deflattice = lattice;
+                    for output in 0..inst.n_outputs {
+                        let $def = $crate::ir::Value::inst(block, instid, output);
+                        $($defbody)*
+                    }
+                    let $uselattice = $deflattice;
+                    for &input in &inst.inputs {
+                        let $use = input;
+                        $($usebody)*
+                    }
+                },
+                start_block: |lattice, block, param_tys| {
+                    let $deflattice = lattice;
+                    for i in 0..param_tys.len() {
+                        let $def = $crate::ir::Value::blockparam(block, i);
+                        $($defbody)*
+                    }
+                },
+                end_block: |lattice, _block, _next, term| {
+                    let $uselattice = lattice;
+                    term.visit_uses(|u| {
+                        let $use = u;
+                        $($usebody)*
+                    });
+                }
+            )
         }
-    fn instruction<'a>(
-        &self,
-        _lattice: &mut L,
-        _block: BlockId,
-        _instid: InstId,
-        _inst: &Inst<'a>,
-    ) -> bool {
-        false
     }
 }
 
@@ -45,7 +149,7 @@ pub struct ForwardDataflow<L: Lattice> {
 }
 
 impl<L: Lattice> ForwardDataflow<L> {
-    pub fn new<'a, D: DataflowFunctions<L>>(f: &FunctionBody<'a>, d: &D) -> Self {
+    pub fn new<D: DataflowFunctions<L = L>>(f: &FunctionBody, d: &D) -> Self {
         let mut analysis = Self {
             block_in: FxHashMap::default(),
         };
@@ -53,7 +157,7 @@ impl<L: Lattice> ForwardDataflow<L> {
         analysis
     }
 
-    fn compute<'a, D: DataflowFunctions<L>>(&mut self, f: &FunctionBody<'a>, d: &D) {
+    fn compute<D: DataflowFunctions<L = L>>(&mut self, f: &FunctionBody, d: &D) {
         let mut workqueue = VecDeque::new();
         let mut workqueue_set = FxHashSet::default();
 
@@ -65,7 +169,7 @@ impl<L: Lattice> ForwardDataflow<L> {
             let mut value = self
                 .block_in
                 .entry(block)
-                .or_insert_with(|| L::top())
+                .or_insert_with(|| D::L::top())
                 .clone();
 
             d.start_block(&mut value, block, &f.blocks[block].params[..]);
@@ -79,13 +183,13 @@ impl<L: Lattice> ForwardDataflow<L> {
                 let mut value = if i + 1 < succs.len() {
                     value.clone()
                 } else {
-                    std::mem::replace(&mut value, L::top())
+                    std::mem::replace(&mut value, D::L::top())
                 };
 
                 d.end_block(&mut value, block, succ, &f.blocks[block].terminator);
 
                 let (succ_in, mut changed) = match self.block_in.entry(succ) {
-                    HashEntry::Vacant(v) => (v.insert(L::top()), true),
+                    HashEntry::Vacant(v) => (v.insert(D::L::top()), true),
                     HashEntry::Occupied(o) => (o.into_mut(), false),
                 };
                 changed |= succ_in.meet_with(&value);
@@ -105,11 +209,7 @@ pub struct BackwardDataflow<L: Lattice> {
 }
 
 impl<L: Lattice> BackwardDataflow<L> {
-    pub fn new<'a, D: DataflowFunctions<L>>(
+    pub fn new<D: DataflowFunctions<L = L>>(f: &FunctionBody, cfginfo: &CFGInfo, d: &D) -> Self {
-        f: &FunctionBody<'a>,
-        cfginfo: &CFGInfo,
-        d: &D,
-    ) -> Self {
         let mut analysis = Self {
             block_out: FxHashMap::default(),
         };
@@ -117,12 +217,7 @@ impl<L: Lattice> BackwardDataflow<L> {
         analysis
     }
 
-    fn compute<'a, D: DataflowFunctions<L>>(
+    fn compute<D: DataflowFunctions<L = L>>(&mut self, f: &FunctionBody, cfginfo: &CFGInfo, d: &D) {
-        &mut self,
-        f: &FunctionBody<'a>,
-        cfginfo: &CFGInfo,
-        d: &D,
-    ) {
         let mut workqueue = VecDeque::new();
         let mut workqueue_set = FxHashSet::default();
 
@@ -145,7 +240,7 @@ impl<L: Lattice> BackwardDataflow<L> {
             let mut value = self
                 .block_out
                 .entry(block)
-                .or_insert_with(|| L::top())
+                .or_insert_with(|| D::L::top())
                 .clone();
 
             for (instid, inst) in f.blocks[block].insts.iter().rev().enumerate() {
@@ -159,7 +254,7 @@ impl<L: Lattice> BackwardDataflow<L> {
                 let mut value = if i + 1 < preds.len() {
                     value.clone()
                 } else {
-                    std::mem::replace(&mut value, L::top())
+                    std::mem::replace(&mut value, D::L::top())
                 };
 
                 d.end_block(&mut value, pred, block, &f.blocks[pred].terminator);
@@ -178,3 +273,34 @@ impl<L: Lattice> BackwardDataflow<L> {
         }
     }
 }
+
+#[macro_export]
+macro_rules! forward_pass {
+    ($name:ident, $lattice:ident, $($dataflow:tt),*) => {
+        #[derive(Clone, Debug)]
+        pub struct $name($crate::pass::ForwardDataflow<$lattice>);
+
+        impl $name {
+            pub fn compute(f: &$crate::ir::FunctionBody) -> $name {
+                let results = $crate::pass::ForwardDataflow::new(f, $($dataflow)*);
+                Self(results)
+            }
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! backward_pass {
+    ($name:ident, $lattice:ident, $($dataflow:tt)*) => {
+        #[derive(Clone, Debug)]
+        pub struct $name($crate::pass::BackwardDataflow<$lattice>);
+
+        impl $name {
+            pub fn compute(f: &$crate::ir::FunctionBody, c: &$crate::cfg::CFGInfo) -> $name {
+                let dataflow = $($dataflow)*;
+                let results = $crate::pass::BackwardDataflow::new(f, c, &dataflow);
+                Self(results)
+            }
+        }
+    };
+}