winch(x64): Add support for return and unreachable

This change adds support for the `return` and `unreachable` instructions. This change builds on top of the control flow building blocks introduced in bytecodealliance#6603
saulecabrera · Jun 21, 2023 · 3135c34 · 3135c34
1 parent b05a09c
commit 3135c34
Show file tree

Hide file tree

Showing 55 changed files with 1,195 additions and 3 deletions.
diff --git a/fuzz/fuzz_targets/differential.rs b/fuzz/fuzz_targets/differential.rs
@@ -362,7 +362,9 @@ fn winch_supports_module(module: &[u8]) -> bool {
                         | Block { .. }
                         | Loop { .. }
                         | Br { .. }
-                        | BrIf { .. } => {}
+                        | BrIf { .. }
+                        | Unreachable { .. }
+                        | Return { .. } => {}
                         _ => {
                             supported = false;
                             break 'main;

diff --git a/winch/codegen/src/codegen/control.rs b/winch/codegen/src/codegen/control.rs
@@ -351,7 +351,7 @@ impl ControlStackFrame {
     /// Returns the exit label of the current control stack frame. Note that
     /// this is similar to [`ControlStackFrame::label`], with the only difference that it
     /// returns `None` for `Loop` since its label doesn't represent an exit.
-    fn exit_label(&self) -> Option<&MachLabel> {
+    pub fn exit_label(&self) -> Option<&MachLabel> {
         use ControlStackFrame::*;
 
         match self {

diff --git a/winch/codegen/src/isa/aarch64/masm.rs b/winch/codegen/src/isa/aarch64/masm.rs
@@ -266,6 +266,10 @@ impl Masm for MacroAssembler {
     fn jmp(&mut self, _target: MachLabel) {
         todo!()
     }
+
+    fn unreachable(&mut self) {
+        todo!()
+    }
 }
 
 impl MacroAssembler {

diff --git a/winch/codegen/src/isa/x64/asm.rs b/winch/codegen/src/isa/x64/asm.rs
@@ -842,4 +842,9 @@ impl Assembler {
     pub fn jmp(&mut self, target: MachLabel) {
         self.emit(Inst::JmpKnown { dst: target });
     }
+
+    /// Emit a trap instruction.
+    pub fn trap(&mut self, code: TrapCode) {
+        self.emit(Inst::Ud2 { trap_code: code })
+    }
 }
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
@@ -14,7 +14,8 @@ use crate::{
 };
 use crate::{isa::reg::Reg, masm::CalleeKind};
 use cranelift_codegen::{
-    isa::x64::settings as x64_settings, settings, Final, MachBufferFinalized, MachLabel,
+    ir::TrapCode, isa::x64::settings as x64_settings, settings, Final, MachBufferFinalized,
+    MachLabel,
 };
 
 /// x64 MacroAssembler.
@@ -483,6 +484,10 @@ impl Masm for MacroAssembler {
             context.free_gpr(tmp);
         }
     }
+
+    fn unreachable(&mut self) {
+        self.asm.trap(TrapCode::UnreachableCodeReached)
+    }
 }
 
 impl MacroAssembler {

diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
@@ -346,4 +346,7 @@ pub(crate) trait MacroAssembler {
 
     /// Emits and unconditional jump to the given label.
     fn jmp(&mut self, target: MachLabel);
+
+    /// Emit an unreachable code trap.
+    fn unreachable(&mut self);
 }
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
@@ -103,6 +103,8 @@ macro_rules! def_unsupported {
     (emit Loop $($rest:tt)*) => {};
     (emit Br $($rest:tt)*) => {};
     (emit BrIf $($rest:tt)*) => {};
+    (emit Return $($rest:tt)*) => {};
+    (emit Unreachable $($rest:tt)*) => {};
 
     (emit $unsupported:tt $($rest:tt)*) => {$($rest)*};
 }
@@ -572,6 +574,30 @@ where
         self.context.free_gpr(top);
     }
 
+    fn visit_return(&mut self) {
+        // Grab the outermost frame, which is the function's body frame. We
+        // don't rely on `Self::control_at` since this frame is implicit and we
+        // know that it should exist at index 0.
+        let outermost = &mut self.control_frames[0];
+        self.context.pop_abi_results(outermost.result(), self.masm);
+        // The outermost should always be a block and therefore,
+        // should always have an exit label.
+        self.masm.jmp(*outermost.exit_label().unwrap());
+        // Set the frame as branch target so that
+        // we can bind the function's exit label.
+        outermost.set_as_target();
+        self.context.reachable = false;
+    }
+
+    fn visit_unreachable(&mut self) {
+        self.masm.unreachable();
+        self.context.reachable = false;
+        // Set the implicit outermost frame as target to perform the necessary
+        // stack clean up.
+        let outermost = &mut self.control_frames[0];
+        outermost.set_as_target();
+    }
+
     wasmparser::for_each_operator!(def_unsupported);
 }
 

diff --git a/winch/filetests/filetests/x64/return/as_block_first.wat b/winch/filetests/filetests/x64/return/as_block_first.wat
@@ -0,0 +1,23 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-block-first")
+   (block (return) (call $dummy))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_block_last.wat b/winch/filetests/filetests/x64/return/as_block_last.wat
@@ -0,0 +1,26 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-block-last")
+   (block (nop) (call $dummy) (return))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    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:	 4883c408             	add	rsp, 8
+;;   1d:	 5d                   	pop	rbp
+;;   1e:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_block_mid.wat b/winch/filetests/filetests/x64/return/as_block_mid.wat
@@ -0,0 +1,26 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-block-mid")
+   (block (call $dummy) (return) (call $dummy))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    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:	 4883c408             	add	rsp, 8
+;;   1d:	 5d                   	pop	rbp
+;;   1e:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_block_value.wat b/winch/filetests/filetests/x64/return/as_block_value.wat
@@ -0,0 +1,27 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-block-value") (result i32)
+   (block (result i32) (nop) (call $dummy) (return (i32.const 2)))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    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:	 48c7c002000000       	mov	rax, 2
+;;   20:	 4883c408             	add	rsp, 8
+;;   24:	 5d                   	pop	rbp
+;;   25:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_br_if_cond.wat b/winch/filetests/filetests/x64/return/as_br_if_cond.wat
@@ -0,0 +1,23 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-br-if-cond")
+   (block (br_if 0 (return)))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_br_value.wat b/winch/filetests/filetests/x64/return/as_br_value.wat
@@ -0,0 +1,24 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-br-value") (result i32)
+   (block (result i32) (br 0 (return (i32.const 9))))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 48c7c009000000       	mov	rax, 9
+;;   13:	 4883c408             	add	rsp, 8
+;;   17:	 5d                   	pop	rbp
+;;   18:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_call_fist.wat b/winch/filetests/filetests/x64/return/as_call_fist.wat
@@ -0,0 +1,28 @@
+;;! target = "x86_64"
+(module
+  (func $f (param i32 i32 i32) (result i32) (i32.const -1))
+  (func (export "as-call-first") (result i32)
+    (call $f (return (i32.const 12)) (i32.const 2) (i32.const 3))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec18             	sub	rsp, 0x18
+;;    8:	 897c2414             	mov	dword ptr [rsp + 0x14], edi
+;;    c:	 89742410             	mov	dword ptr [rsp + 0x10], esi
+;;   10:	 8954240c             	mov	dword ptr [rsp + 0xc], edx
+;;   14:	 4c89742404           	mov	qword ptr [rsp + 4], r14
+;;   19:	 48c7c0ffffffff       	mov	rax, 0xffffffffffffffff
+;;   20:	 4883c418             	add	rsp, 0x18
+;;   24:	 5d                   	pop	rbp
+;;   25:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 48c7c00c000000       	mov	rax, 0xc
+;;   13:	 4883c408             	add	rsp, 8
+;;   17:	 5d                   	pop	rbp
+;;   18:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_call_last.wat b/winch/filetests/filetests/x64/return/as_call_last.wat
@@ -0,0 +1,28 @@
+;;! target = "x86_64"
+(module
+  (func $f (param i32 i32 i32) (result i32) (i32.const -1))
+  (func (export "as-call-last") (result i32)
+    (call $f (i32.const 1) (i32.const 2) (return (i32.const 14)))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec18             	sub	rsp, 0x18
+;;    8:	 897c2414             	mov	dword ptr [rsp + 0x14], edi
+;;    c:	 89742410             	mov	dword ptr [rsp + 0x10], esi
+;;   10:	 8954240c             	mov	dword ptr [rsp + 0xc], edx
+;;   14:	 4c89742404           	mov	qword ptr [rsp + 4], r14
+;;   19:	 48c7c0ffffffff       	mov	rax, 0xffffffffffffffff
+;;   20:	 4883c418             	add	rsp, 0x18
+;;   24:	 5d                   	pop	rbp
+;;   25:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 48c7c00e000000       	mov	rax, 0xe
+;;   13:	 4883c408             	add	rsp, 8
+;;   17:	 5d                   	pop	rbp
+;;   18:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_call_mid.wat b/winch/filetests/filetests/x64/return/as_call_mid.wat
@@ -0,0 +1,28 @@
+;;! target = "x86_64"
+(module
+  (func $f (param i32 i32 i32) (result i32) (i32.const -1))
+  (func (export "as-call-mid") (result i32)
+    (call $f (i32.const 1) (return (i32.const 13)) (i32.const 3))
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec18             	sub	rsp, 0x18
+;;    8:	 897c2414             	mov	dword ptr [rsp + 0x14], edi
+;;    c:	 89742410             	mov	dword ptr [rsp + 0x10], esi
+;;   10:	 8954240c             	mov	dword ptr [rsp + 0xc], edx
+;;   14:	 4c89742404           	mov	qword ptr [rsp + 4], r14
+;;   19:	 48c7c0ffffffff       	mov	rax, 0xffffffffffffffff
+;;   20:	 4883c418             	add	rsp, 0x18
+;;   24:	 5d                   	pop	rbp
+;;   25:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 48c7c00d000000       	mov	rax, 0xd
+;;   13:	 4883c408             	add	rsp, 8
+;;   17:	 5d                   	pop	rbp
+;;   18:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_func_first.wat b/winch/filetests/filetests/x64/return/as_func_first.wat
@@ -0,0 +1,24 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+  (func (export "as-func-first") (result i32)
+   (return (i32.const 1)) (i32.const 2)
+  )
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 48c7c001000000       	mov	rax, 1
+;;   13:	 4883c408             	add	rsp, 8
+;;   17:	 5d                   	pop	rbp
+;;   18:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/return/as_func_last.wat b/winch/filetests/filetests/x64/return/as_func_last.wat
@@ -0,0 +1,12 @@
+;;! target = "x86_64"
+(module
+  (func $dummy)
+)
+
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883c408             	add	rsp, 8
+;;   10:	 5d                   	pop	rbp
+;;   11:	 c3                   	ret