winch(x64): Add support for return and unreachable

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;

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 {

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 {

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 })
+    }
 }

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 {

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);
 }

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);
 }