From 63aa7e4cb6b4bc29ec1e1584f79b458dd52c6ba9 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Fri, 24 Dec 2021 13:20:30 -0800
Subject: [PATCH] working roundtrip for simple module

---
 src/backend/final.rs      |  5 +++++
 src/backend/serialize.rs  |  2 ++
 src/backend/structured.rs | 18 +++++++++++-------
 src/bin/roundtrip.rs      | 13 +++++++++++++
 src/frontend.rs           |  3 +++
 src/ir.rs                 |  1 +
 6 files changed, 35 insertions(+), 7 deletions(-)
 create mode 100644 src/bin/roundtrip.rs

diff --git a/src/backend/final.rs b/src/backend/final.rs
index a06ef09..78df26c 100644
--- a/src/backend/final.rs
+++ b/src/backend/final.rs
@@ -169,6 +169,11 @@ pub fn produce_func_wasm<FT: FuncTypeSink>(
     for operator in &body.operators {
         ctx.translate(operator, locations);
     }
+    
+    // Fixup: add unreachable just before last `end`; there must be an explicit return.
+//    assert!(matches!(wasm.operators.pop(), Some(wasm_encoder::Instruction::End)));
+//    wasm.operators.push(wasm_encoder::Instruction::Unreachable);
+//    wasm.operators.push(wasm_encoder::Instruction::End);
 
     wasm
 }
diff --git a/src/backend/serialize.rs b/src/backend/serialize.rs
index 7f1994e..96c40ae 100644
--- a/src/backend/serialize.rs
+++ b/src/backend/serialize.rs
@@ -288,6 +288,8 @@ impl<'a> SerializedBodyContext<'a> {
                         self.operators.extend(rev_ops);
                         self.operators
                             .push(SerializedOperator::BrTable { targets, default });
+                        self.operators
+                            .push(SerializedOperator::Operator(Operator::Unreachable));
                     }
                     &Terminator::Return { ref values, .. } => {
                         let mut rev_ops = vec![];
diff --git a/src/backend/structured.rs b/src/backend/structured.rs
index 523ecfd..a89811e 100644
--- a/src/backend/structured.rs
+++ b/src/backend/structured.rs
@@ -345,13 +345,17 @@ impl BlockOrder {
                     target_stack.push(target);
                 }
                 let params = f.blocks[header].params.clone();
-                let results = match fallthrough {
-                    Some(fallthrough) => f.blocks[fallthrough]
-                        .params
-                        .iter()
-                        .map(|(ty, _)| *ty)
-                        .collect(),
-                    None => vec![],
+                let results = if header == 0 {
+                    f.rets.clone()
+                } else {
+                    match fallthrough {
+                        Some(fallthrough) => f.blocks[fallthrough]
+                            .params
+                            .iter()
+                            .map(|(ty, _)| *ty)
+                            .collect(),
+                        None => vec![],
+                    }
                 };
                 if is_loop {
                     entries.push(BlockOrderEntry::StartLoop(header, params, results));
diff --git a/src/bin/roundtrip.rs b/src/bin/roundtrip.rs
new file mode 100644
index 0000000..e75d032
--- /dev/null
+++ b/src/bin/roundtrip.rs
@@ -0,0 +1,13 @@
+//! Roundtrip utility.
+
+use std::io::{Read, Write};
+use waffle::Module;
+
+fn main() {
+    let _ = env_logger::try_init();
+    let mut bytes = vec![];
+    std::io::stdin().read_to_end(&mut bytes).unwrap();
+    let module = Module::from_wasm_bytes(&bytes).unwrap();
+    let new_bytes = module.to_wasm_bytes();
+    std::io::stdout().write(&new_bytes[..]).unwrap();
+}
diff --git a/src/frontend.rs b/src/frontend.rs
index 5a0d02c..8f1254f 100644
--- a/src/frontend.rs
+++ b/src/frontend.rs
@@ -109,6 +109,9 @@ fn parse_body<'a>(
     for &param in &module.signatures[my_sig].params[..] {
         ret.locals.push(param);
     }
+    for &r in &module.signatures[my_sig].returns[..] {
+        ret.rets.push(r);
+    }
 
     let mut locals = body.get_locals_reader()?;
     for _ in 0..locals.get_count() {
diff --git a/src/ir.rs b/src/ir.rs
index f1d2892..5aa43d0 100644
--- a/src/ir.rs
+++ b/src/ir.rs
@@ -53,6 +53,7 @@ impl FuncDecl {
 
 #[derive(Clone, Debug, Default)]
 pub struct FunctionBody {
+    pub rets: Vec<Type>,
     pub locals: Vec<Type>,
     pub blocks: Vec<Block>,
     /// Sea-of-nodes representation.