From 677d54e46deedc17840a19968477cc9db2e4b1e3 Mon Sep 17 00:00:00 2001
From: Jef <jackefransham@gmail.com>
Date: Wed, 27 Feb 2019 17:59:27 +0100
Subject: [PATCH] Experimental br_table support

---
 src/backend.rs       | 110 ++++++++++++++++++++++++++++++++++++++++---
 src/function_body.rs |  61 +++++++++++++++++++++++-
 src/microwasm.rs     |  52 ++++++++++++++++----
 src/tests.rs         |  46 ++++++++++++++++--
 4 files changed, 248 insertions(+), 21 deletions(-)

diff --git a/src/backend.rs b/src/backend.rs
index 41debc84ba..93d3aaff8d 100644
--- a/src/backend.rs
+++ b/src/backend.rs
@@ -1,6 +1,6 @@
 #![allow(dead_code)] // for now
 
-use microwasm::{SignlessType, Type, F32, F64, I32, I64};
+use microwasm::{BrTarget, SignlessType, Type, F32, F64, I32, I64};
 
 use self::registers::*;
 use dynasmrt::x64::Assembler;
@@ -631,6 +631,7 @@ pub enum MemoryAccessMode {
 #[derive(Default)]
 struct Labels {
     trap: Option<Label>,
+    ret: Option<Label>,
     neg_const_f32: Option<Label>,
     neg_const_f64: Option<Label>,
 }
@@ -1564,6 +1565,73 @@ impl<M: ModuleContext> Context<'_, M> {
         );
     }
 
+    /// If `default` is `None` then the default is just continuing execution
+    pub fn br_table<I>(
+        &mut self,
+        targets: I,
+        default: Option<BrTarget<Label>>,
+        mut pass_args: impl FnOnce(&mut Self),
+    ) where
+        I: IntoIterator<Item = BrTarget<Label>>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let mut targets = targets.into_iter();
+        let count = targets.len();
+
+        let mut selector = self.pop();
+
+        pass_args(self);
+
+        if count == 0 {
+            if let Some(default) = default {
+                match default {
+                    BrTarget::Label(label) => self.br(label),
+                    BrTarget::Return => {
+                        dynasm!(self.asm
+                            ; ret
+                        );
+                    }
+                }
+            }
+        } else if let Some(imm) = selector.imm_i32() {
+            if let Some(target) = targets.nth(imm as _).or(default) {
+                match target {
+                    BrTarget::Label(label) => self.br(label),
+                    BrTarget::Return => {
+                        dynasm!(self.asm
+                            ; ret
+                        );
+                    }
+                }
+            }
+        } else {
+            let selector_reg = self.into_reg(GPRType::Rq, selector);
+            selector = ValueLocation::Reg(selector_reg);
+
+            // TODO: Jump table (wrestling with dynasm to implement it is too much work)
+            for (i, target) in targets.enumerate() {
+                let label = self.target_to_label(target);
+                dynasm!(self.asm
+                    ; cmp Rq(selector_reg.rq().unwrap()), i as i32
+                    ; je =>label.0
+                );
+            }
+
+            if let Some(def) = default {
+                match def {
+                    BrTarget::Label(label) => dynasm!(self.asm
+                        ; jmp =>label.0
+                    ),
+                    BrTarget::Return => dynasm!(self.asm
+                        ; ret
+                    ),
+                }
+            }
+        }
+
+        self.free_value(selector);
+    }
+
     fn set_stack_depth_preserve_flags(&mut self, depth: StackDepth) {
         if self.block_state.depth.0 < depth.0 {
             // TODO: We need to preserve ZF on `br_if` so we use `push`/`pop` but that isn't
@@ -1604,7 +1672,12 @@ impl<M: ModuleContext> Context<'_, M> {
 
     pub fn pass_block_args(&mut self, cc: &CallingConvention) {
         let args = &cc.arguments;
-        for (remaining, &dst) in args.iter().enumerate().rev() {
+        for (remaining, &dst) in args
+            .iter()
+            .enumerate()
+            .rev()
+            .take(self.block_state.stack.len())
+        {
             if let CCLoc::Reg(r) = dst {
                 if !self.block_state.regs.is_free(r)
                     && *self.block_state.stack.last().unwrap() != ValueLocation::Reg(r)
@@ -2482,11 +2555,9 @@ impl<M: ModuleContext> Context<'_, M> {
     }
 
     fn align(&mut self, align_to: u32) {
-        while self.asm.offset().0 % align_to as usize != 0 {
-            dynasm!(self.asm
-                ; .byte 0
-            );
-        }
+        dynasm!(self.asm
+            ; .align align_to as usize
+        );
     }
 
     /// Writes the function epilogue (right now all this does is add the trap label that the
@@ -2500,6 +2571,13 @@ impl<M: ModuleContext> Context<'_, M> {
             );
         }
 
+        if let Some(l) = self.labels.ret {
+            self.define_label(l);
+            dynasm!(self.asm
+                ; ret
+            );
+        }
+
         if let Some(l) = self.labels.neg_const_f32 {
             self.align(16);
             self.define_label(l);
@@ -2529,6 +2607,13 @@ impl<M: ModuleContext> Context<'_, M> {
         );
     }
 
+    fn target_to_label(&mut self, target: BrTarget<Label>) -> Label {
+        match target {
+            BrTarget::Label(label) => label,
+            BrTarget::Return => self.ret_label(),
+        }
+    }
+
     #[must_use]
     fn trap_label(&mut self) -> Label {
         if let Some(l) = self.labels.trap {
@@ -2540,6 +2625,17 @@ impl<M: ModuleContext> Context<'_, M> {
         label
     }
 
+    #[must_use]
+    fn ret_label(&mut self) -> Label {
+        if let Some(l) = self.labels.ret {
+            return l;
+        }
+
+        let label = self.create_label();
+        self.labels.ret = Some(label);
+        label
+    }
+
     #[must_use]
     fn neg_const_f32_label(&mut self) -> Label {
         if let Some(l) = self.labels.neg_const_f32 {
diff --git a/src/function_body.rs b/src/function_body.rs
index 762b0c3552..07fa1c06bf 100644
--- a/src/function_body.rs
+++ b/src/function_body.rs
@@ -36,7 +36,7 @@ where
 {
     let ty = session.module_context.func_type(func_idx);
 
-    if false {
+    if true {
         let mut microwasm = vec![];
 
         let microwasm_conv = MicrowasmConv::new(
@@ -285,7 +285,9 @@ where
                             **then_cc = cc.clone();
                             **else_cc = cc;
                         }
-                        _ => unimplemented!(),
+                        _ => unimplemented!(
+                            "Can't pass different params to different sides of `br_if` yet"
+                        ),
                     }
                 };
 
@@ -303,6 +305,61 @@ where
                     other => unimplemented!("{:#?}", other),
                 }
             }
+            Operator::BrTable(BrTable { targets, default }) => {
+                use itertools::Itertools;
+
+                let (def, params) = {
+                    let def = blocks.get(&default).unwrap();
+                    (
+                        if def.is_next {
+                            None
+                        } else {
+                            Some(def.label)
+                        },
+                        def.params.clone()
+                    )
+                };
+
+                let target_labels = targets.iter()
+                    .map(|target| blocks.get(target).unwrap().label)
+                    .collect::<Vec<_>>();
+
+                ctx.br_table(target_labels, def, |ctx| {
+                    let mut cc = None;
+                    let mut max_num_callers = Some(0);
+
+                    for target in targets.iter().chain(std::iter::once(&default)).unique() {
+                        let block = blocks.get_mut(target).unwrap();
+                        block.actual_num_callers += 1;
+
+                        if block.calling_convention.is_some() {
+                            assert!(cc.is_none(), "Can't pass different params to different elements of `br_table` yet");
+                            cc = block.calling_convention.clone();
+                        }
+
+                        if let Some(max) = max_num_callers {
+                            max_num_callers = block.num_callers.map(|n| max.max(n));
+                        }
+                    }
+
+                    if let Some(Left(cc)) = &cc {
+                        ctx.pass_block_args(cc);
+                    }
+       
+                    let cc = cc.unwrap_or_else(||
+                        if max_num_callers == Some(1) {
+                            Right(ctx.virtual_calling_convention())
+                        } else {
+                            Left(ctx.serialize_args(params))
+                        }
+                    );
+
+                    for target in targets.iter().chain(std::iter::once(&default)).unique() {
+                        let block = blocks.get_mut(target).unwrap();
+                        block.calling_convention = Some(cc.clone());
+                    }
+                });
+            }
             Operator::Swap { depth } => ctx.swap(depth),
             Operator::Pick { depth } => ctx.pick(depth),
             Operator::Eq(I32) => ctx.i32_eq(),
diff --git a/src/microwasm.rs b/src/microwasm.rs
index 89d2a1a402..0b85325d72 100644
--- a/src/microwasm.rs
+++ b/src/microwasm.rs
@@ -310,7 +310,8 @@ impl TryFrom<wasmparser::Type> for SignlessType {
 
 #[derive(Debug, Clone)]
 pub struct BrTable<L> {
-    targets: Vec<L>,
+    pub targets: Vec<BrTarget<L>>,
+    pub default: BrTarget<L>,
 }
 
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
@@ -371,8 +372,8 @@ impl fmt::Display for BrTarget<&str> {
 pub enum Operator<Label> {
     /// Explicit trap instruction
     Unreachable,
-    /// Start a new block. It is an error if the previous block has not been closed by emitting a `Br` or
-    /// `BrTable`.
+    /// Define metadata for a block - its label, its signature, whether it has backwards callers etc. It
+    /// is an error to branch to a block that has yet to be defined.
     Block {
         label: Label,
         // TODO: Do we need this?
@@ -381,6 +382,8 @@ pub enum Operator<Label> {
         has_backwards_callers: bool,
         num_callers: Option<u32>,
     },
+    /// Start a new block. It is an error if the previous block has not been closed by emitting a `Br` or
+    /// `BrTable`.
     Label(Label),
     /// Unconditionally break to a new block. This the parameters off the stack and passes them into
     /// the new block. Any remaining elements on the stack are discarded.
@@ -398,10 +401,10 @@ pub enum Operator<Label> {
     },
     /// Pop a value off the top of the stack, jump to `table[value.min(table.len() - 1)]`. All elements
     /// in the table must have the same parameters.
-    BrTable {
+    BrTable(
         /// The table of labels to jump to - the index should be clamped to the length of the table
-        table: BrTable<Label>,
-    },
+        BrTable<Label>,
+    ),
     /// Call a function
     Call {
         function_index: u32,
@@ -643,6 +646,18 @@ where
             }
             Operator::Br { target } => write!(f, "br {}", target),
             Operator::BrIf { then, else_ } => write!(f, "br_if {}, {}", then, else_),
+            Operator::BrTable(BrTable { targets, default }) => {
+                write!(f, "br_table [")?;
+                let mut iter = targets.iter();
+                if let Some(p) = iter.next() {
+                    write!(f, "{}", p)?;
+                    for p in iter {
+                        write!(f, ", {}", p)?;
+                    }
+                }
+
+                write!(f, "], {}", default)
+            },
             Operator::Call { function_index } => write!(f, "call {}", function_index),
             Operator::CallIndirect { .. } => write!(f, "call_indirect"),
             Operator::Drop(range) => {
@@ -798,7 +813,9 @@ impl ControlFrame {
         match self.kind {
             ControlFrameKind::Loop => BrTarget::Label((self.id, NameTag::Header)),
             ControlFrameKind::Function => BrTarget::Return,
-            _ => BrTarget::Label((self.id, NameTag::End)),
+            ControlFrameKind::Block { .. } | ControlFrameKind::If { .. } => {
+                BrTarget::Label((self.id, NameTag::End))
+            }
         }
     }
 
@@ -1593,7 +1610,26 @@ where
                     ]
                 }
             }
-            WasmOperator::BrTable { .. } => unimplemented!("{:?}", op),
+            WasmOperator::BrTable { table } => {
+                self.unreachable = true;
+                let (entries, default) = match table.read_table() {
+                    Ok(o) => o,
+                    Err(e) => return Some(Err(e)),
+                };
+                let targets = entries
+                    .into_iter()
+                    .map(|depth| {
+                        let block = self.nth_block_mut(*depth as _);
+                        block.mark_branched_to();
+                        block.br_target()
+                    })
+                    .collect();
+                let default = self.nth_block_mut(default as _);
+                default.mark_branched_to();
+                let default = default.br_target();
+
+                smallvec![Operator::BrTable(BrTable { targets, default })]
+            }
             WasmOperator::Return => {
                 self.unreachable = true;
 
diff --git a/src/tests.rs b/src/tests.rs
index 6e3c771fb9..5abf11ab12 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1182,7 +1182,44 @@ fn fib_opt() {
 }
 
 #[test]
-fn just_storage() {
+fn br_table() {
+    const CODE: &str = r"
+(module
+  (func
+    (block (br_table 0 0 0 (i32.const 0)) (call $dummy))
+  )
+  (func
+    (block (call $dummy) (br_table 0 0 0 (i32.const 0)) (call $dummy))
+  )
+  (func
+    (block (nop) (call $dummy) (br_table 0 0 0 (i32.const 0)))
+  )
+  (func $dummy)
+)
+";
+
+    let translated = translate_wat(CODE);
+    translated.disassemble();
+
+    println!("as-block-first");
+    assert_eq!(
+        translated.execute_func::<_, ()>(0, ()),
+        Ok(()),
+    );
+    println!("as-block-mid");
+    assert_eq!(
+        translated.execute_func::<_, ()>(1, ()),
+        Ok(()),
+    );
+    println!("as-block-last");
+    assert_eq!(
+        translated.execute_func::<_, ()>(2, ()),
+        Ok(()),
+    );
+}
+
+#[test]
+fn storage() {
     const CODE: &str = r#"
 (module
   (memory 1 1)
@@ -1283,8 +1320,8 @@ fn nested_storage_calls() {
     assert_eq!(translated.execute_func::<(), i32>(0, ()), Ok(1));
 }
 
-// TODO: Signature mismatches correctly fail, but we can't add a test
-//       for that until we implement traps properly.
+// TODO: Signature mismatches correctly fail at time of writing this comment,
+//       but we can't add a test for that until we implement traps properly.
 #[test]
 fn call_indirect() {
     const CODE: &str = r#"
@@ -1377,7 +1414,7 @@ macro_rules! test_select {
 
                 fn lit(cond: bool, then: $ty, else_: $ty) -> bool {
                     let icond: i32 = if cond { 1 } else { 0 };
-                                            let translated = translate_wat(&format!("
+                                                    let translated = translate_wat(&format!("
                             (module (func (param {ty}) (param {ty}) (result {ty})
                                 (select (get_local 0) (get_local 1) (i32.const {val}))))
                         ",
@@ -1778,3 +1815,4 @@ fn sieve() {
 
     translate(&wabt::wat2wasm(CODE).unwrap()).unwrap();
 }
+