Winch: cleanup stack in br_if in non-fallthrough case (bytecodeallian…

…ce#7590) * Winch: cleanup stack in br_if in non-fallthrough case * Remove unnecessary refetch of sp_offsets * Refactoring based on PR feedback * Have SPOffset implement Ord
alexcrichton · Nov 28, 2023 · 4d22446 · 4d22446
1 parent 1c7a07b
commit 4d22446
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 13 deletions.
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
@@ -24,7 +24,7 @@ pub(crate) enum RemKind {
 }
 
 /// Representation of the stack pointer offset.
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, PartialOrd, Ord)]
 pub struct SPOffset(u32);
 
 impl SPOffset {

diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
@@ -1189,14 +1189,31 @@ where
             self.context.pop_to_reg(self.masm, None)
         };
 
-        self.masm.branch(
-            IntCmpKind::Ne,
-            top.reg.into(),
-            top.reg.into(),
-            *frame.label(),
-            OperandSize::S32,
-        );
+        // Emit instructions to balance the machine stack if the frame has
+        // a different offset.
+        let current_sp_offset = self.masm.sp_offset();
+        let (_, frame_sp_offset) = frame.base_stack_len_and_sp();
+        let (label, cmp, needs_cleanup) = if current_sp_offset > frame_sp_offset {
+            (self.masm.get_label(), IntCmpKind::Eq, true)
+        } else {
+            (*frame.label(), IntCmpKind::Ne, false)
+        };
+
+        self.masm
+            .branch(cmp, top.reg.into(), top.reg.into(), label, OperandSize::S32);
         self.context.free_reg(top);
+
+        if needs_cleanup {
+            // Emit instructions to balance the stack and jump if not falling
+            // through.
+            self.masm.ensure_sp_for_jump(frame_sp_offset);
+            self.masm.jmp(*frame.label());
+
+            // Restore sp_offset to what it was for falling through and emit
+            // fallthrough label.
+            self.masm.reset_stack_pointer(current_sp_offset);
+            self.masm.bind(label);
+        }
     }
 
     fn visit_br_table(&mut self, targets: BrTable<'a>) {

diff --git a/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat b/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat
@@ -0,0 +1,42 @@
+;;! target = "x86_64"
+
+(module
+  (func (export "") 
+    call 1
+    call 1
+    br_if 0
+    drop
+  )
+  (func (;1;) (result i32)
+    i32.const 1
+  )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883ec08             	sub	rsp, 8
+;;   10:	 e800000000           	call	0x15
+;;   15:	 4883c408             	add	rsp, 8
+;;   19:	 4883ec04             	sub	rsp, 4
+;;   1d:	 890424               	mov	dword ptr [rsp], eax
+;;   20:	 4883ec04             	sub	rsp, 4
+;;   24:	 e800000000           	call	0x29
+;;   29:	 4883c404             	add	rsp, 4
+;;   2d:	 85c0                 	test	eax, eax
+;;   2f:	 0f8409000000         	je	0x3e
+;;   35:	 4883c404             	add	rsp, 4
+;;   39:	 e904000000           	jmp	0x42
+;;   3e:	 4883c404             	add	rsp, 4
+;;   42:	 4883c408             	add	rsp, 8
+;;   46:	 5d                   	pop	rbp
+;;   47:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 b801000000           	mov	eax, 1
+;;   11:	 4883c408             	add	rsp, 8
+;;   15:	 5d                   	pop	rbp
+;;   16:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/call/reg_on_stack.wat b/winch/filetests/filetests/x64/call/reg_on_stack.wat
@@ -35,8 +35,10 @@
 ;;   56:	 8b0424               	mov	eax, dword ptr [rsp]
 ;;   59:	 4883c404             	add	rsp, 4
 ;;   5d:	 85c9                 	test	ecx, ecx
-;;   5f:	 0f8502000000         	jne	0x67
-;;   65:	 0f0b                 	ud2	
-;;   67:	 4883c410             	add	rsp, 0x10
-;;   6b:	 5d                   	pop	rbp
-;;   6c:	 c3                   	ret	
+;;   5f:	 0f8409000000         	je	0x6e
+;;   65:	 4883c404             	add	rsp, 4
+;;   69:	 e902000000           	jmp	0x70
+;;   6e:	 0f0b                 	ud2	
+;;   70:	 4883c410             	add	rsp, 0x10
+;;   74:	 5d                   	pop	rbp
+;;   75:	 c3                   	ret