From cf9c57e279170387f9053a47c1b4a132b8637f35 Mon Sep 17 00:00:00 2001
From: Jef <jackefransham@gmail.com>
Date: Tue, 26 Mar 2019 10:29:55 +0100
Subject: [PATCH] Fix `br_if` calling convention handling

---
 src/backend.rs       | 38 ++++++++++++++++++++++++++++++++++----
 src/function_body.rs | 26 +++++++++++++++-----------
 2 files changed, 49 insertions(+), 15 deletions(-)

diff --git a/src/backend.rs b/src/backend.rs
index 7ca5090a15..b3502d3da0 100644
--- a/src/backend.rs
+++ b/src/backend.rs
@@ -2391,6 +2391,38 @@ impl<'this, M: ModuleContext> Context<'this, M> {
         self.set_stack_depth_preserve_flags(cc.stack_depth);
     }
 
+    pub fn serialize_block_args_preserve_flags(
+        &mut self,
+        cc: &BlockCallingConvention,
+        other_to_drop: Option<RangeInclusive<u32>>,
+    ) -> BlockCallingConvention {
+        self.do_pass_block_args(cc);
+
+        let mut out_args = cc.arguments.clone();
+
+        out_args.reverse();
+
+        if let Some(to_drop) = other_to_drop {
+            for _ in to_drop {
+                let val = self.pop();
+                // TODO: We can use stack slots for values already on the stack but we
+                //       don't refcount stack slots right now
+                let loc = CCLoc::Reg(self.into_temp_reg(None, val));
+
+                out_args.push(loc);
+            }
+        }
+
+        out_args.reverse();
+
+        self.set_stack_depth_preserve_flags(cc.stack_depth);
+
+        BlockCallingConvention {
+            stack_depth: cc.stack_depth,
+            arguments: out_args,
+        }
+    }
+
     /// Puts all stack values into "real" locations so that they can i.e. be set to different
     /// values on different iterations of a loop
     pub fn serialize_args(&mut self, count: u32) -> BlockCallingConvention {
@@ -4343,9 +4375,9 @@ impl<'this, M: ModuleContext> Context<'this, M> {
 
         self.push_function_returns(rets);
 
-        self.set_stack_depth(depth);
-
         if preserve_vmctx {
+            self.set_stack_depth(depth);
+
             dynasm!(self.asm
                 ; pop Rq(VMCTX)
             );
@@ -4702,8 +4734,6 @@ impl<'this, M: ModuleContext> Context<'this, M> {
         }
 
         self.push_function_returns(return_types);
-
-        self.set_stack_depth(depth);
     }
 
     /// Call a function with the given index
diff --git a/src/function_body.rs b/src/function_body.rs
index 8565b48116..abcecbf885 100644
--- a/src/function_body.rs
+++ b/src/function_body.rs
@@ -315,17 +315,21 @@ where
                     let else_block_should_serialize_args = else_block.should_serialize_args();
 
                     match (
-                        &mut then_block.calling_convention,
-                        &mut else_block.calling_convention,
+                        (&mut then_block.calling_convention, &then.to_drop),
+                        (&mut else_block.calling_convention, &else_.to_drop),
                     ) {
-                        (Some(Left(ref cc)), ref mut other @ None)
-                        | (ref mut other @ None, Some(Left(ref cc))) => {
-                            assert_eq!(then.to_drop, else_.to_drop);
-                            **other = Some(Left(cc.clone()));
-
-                            ctx.pass_block_args_preserve_flags(cc);
+                        ((Some(Left(ref cc)), to_drop), ref mut other @ (None, _))
+                        | (ref mut other @ (None, _), (Some(Left(ref cc)), to_drop)) => {
+                            let mut cc = ctx.serialize_block_args_preserve_flags(cc, to_drop.clone());
+                            if let Some(to_drop) = other.1 {
+                                drop_elements(&mut cc.arguments, to_drop.clone());
+                            }
+                            *other.0 = Some(Left(cc));
                         }
-                        (ref mut then_cc @ None, ref mut else_cc @ None) => {
+                        (
+                            (ref mut then_cc @ None, then_to_drop),
+                            (ref mut else_cc @ None, else_to_drop),
+                        ) => {
                             let max_params = then_block.params.max(else_block.params);
                             let cc = if then_block_should_serialize_args {
                                 Left(ctx.serialize_args(max_params))
@@ -337,7 +341,7 @@ where
 
                             **then_cc = {
                                 let mut cc = cc.clone();
-                                if let Some(to_drop) = then.to_drop.clone() {
+                                if let Some(to_drop) = then_to_drop.clone() {
                                     match &mut cc {
                                         Left(cc) => drop_elements(&mut cc.arguments, to_drop),
                                         Right(cc) => drop_elements(&mut cc.stack, to_drop),
@@ -347,7 +351,7 @@ where
                             };
                             **else_cc = {
                                 let mut cc = cc;
-                                if let Some(to_drop) = else_.to_drop.clone() {
+                                if let Some(to_drop) = else_to_drop.clone() {
                                     match &mut cc {
                                         Left(cc) => drop_elements(&mut cc.arguments, to_drop),
                                         Right(cc) => drop_elements(&mut cc.stack, to_drop),