From f483b964bb2243768ab14802c9af9824aba8361c Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Sat, 13 Nov 2021 03:38:47 -0800
Subject: [PATCH] WIP.

---
 src/frontend.rs | 53 ++++++++++++++++++++++++++++++++++++++++++++-----
 src/ir.rs       | 28 ++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/src/frontend.rs b/src/frontend.rs
index 3a8221c..afa32d5 100644
--- a/src/frontend.rs
+++ b/src/frontend.rs
@@ -6,7 +6,9 @@ use crate::ir::*;
 use crate::op_traits::{op_inputs, op_outputs};
 use anyhow::{bail, Result};
 use log::trace;
-use wasmparser::{ImportSectionEntryType, Operator, Parser, Payload, Type, TypeDef};
+use wasmparser::{
+    ImportSectionEntryType, Operator, Parser, Payload, Type, TypeDef, TypeOrFuncType,
+};
 
 pub fn wasm_to_ir<'a>(bytes: &'a [u8]) -> Result<Module<'a>> {
     let mut module = Module::default();
@@ -158,12 +160,29 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
             | Operator::LocalSet { .. }
             | Operator::LocalTee { .. } => self.emit(op.clone())?,
 
-            Operator::End => {
-                if self.ctrl_stack.is_empty() {
+            Operator::End => match self.ctrl_stack.pop() {
+                None => {
                     self.emit(Operator::Return)?;
-                } else {
-                    bail!("Unsupported End");
                 }
+                Some(Frame::Block { out, .. }) | Some(Frame::Loop { out, .. }) => {
+                    // No need to manipulate stack: assuming the input
+                    // Wasm was validated properly, the `results`
+                    // values must be on the top of the stack now, and
+                    // they will remain so once we exit the block.
+                    self.emit_branch(out);
+                    self.cur_block = out;
+                }
+                _ => bail!("unsupported block type"),
+            },
+
+            Operator::Block { ty } => {
+                let (params, results) = self.block_params_and_results(ty);
+                let out = self.create_block();
+                self.ctrl_stack.push(Frame::Block {
+                    out,
+                    params,
+                    results,
+                });
             }
 
             _ => bail!("Unsupported operator: {:?}", op),
@@ -172,6 +191,30 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
         Ok(())
     }
 
+    fn create_block(&mut self) -> BlockId {
+        let id = self.body.blocks.len() as BlockId;
+        self.body.blocks.push(Block::default());
+        id
+    }
+
+    fn block_params_and_results(&self, ty: TypeOrFuncType) -> (Vec<Type>, Vec<Type>) {
+        match ty {
+            TypeOrFuncType::Type(ret_ty) => (vec![], vec![ret_ty]),
+            TypeOrFuncType::FuncType(sig_idx) => {
+                let sig = &self.module.signatures[sig_idx as SignatureId];
+                (
+                    Vec::from(sig.params.clone()),
+                    Vec::from(sig.returns.clone()),
+                )
+            }
+        }
+    }
+
+    fn emit_branch(&mut self, target: BlockId) {
+        let block = self.cur_block;
+        self.body.blocks[block].terminator = Terminator::Br { target };
+    }
+
     fn emit(&mut self, op: Operator<'a>) -> Result<()> {
         let block = self.cur_block;
         let inst = self.body.blocks[block].insts.len() as InstId;
diff --git a/src/ir.rs b/src/ir.rs
index c0f18bd..f886efb 100644
--- a/src/ir.rs
+++ b/src/ir.rs
@@ -52,6 +52,7 @@ pub enum ValueKind {
 pub struct Block<'a> {
     pub params: Vec<Type>,
     pub insts: Vec<Inst<'a>>,
+    pub terminator: Terminator<'a>,
 }
 
 #[derive(Clone, Debug)]
@@ -66,3 +67,30 @@ pub enum Operand<'a> {
     Value(ValueId),
     Sub(Box<Inst<'a>>),
 }
+
+#[derive(Clone, Debug)]
+pub enum Terminator<'a> {
+    Br {
+        target: BlockId,
+    },
+    CondBr {
+        cond: Operand<'a>,
+        if_true: BlockId,
+        if_false: BlockId,
+    },
+    Select {
+        value: Operand<'a>,
+        targets: Vec<BlockId>,
+        default: BlockId,
+    },
+    Return {
+        values: Vec<Operand<'a>>,
+    },
+    None,
+}
+
+impl<'a> std::default::Default for Terminator<'a> {
+    fn default() -> Self {
+        Terminator::None
+    }
+}