Move return_at_end out of Settings and into the wasm FuncEnvironment. (#547)

* Move `return_at_end` out of Settings and into the wasm FuncEnvironment.

The `return_at_end` flag supports users that want to append a custom
epilogue to Cranelift-produced functions. It arranges for functions to
always return via a single return statement at the end, and users are
expected to remove this return to append their code.

This patch makes two changes:
 - First, introduce a `fallthrough_return` instruction and use that
   instead of adding a `return` at the end. That's simpler than having
   users remove the `return` themselves.

 - Second, move this setting out of the Settings and into the wasm
   FuncEnvironment. This flag isn't something the code generator uses,
   it's something that the wasm translator uses. The code generator
   needs to preserve the property, however we can give the
   `fallthrough_return` instruction properties to ensure this as needed,
   such as marking it non-cloneable.

lib/wasm/src/code_translator.rs

@@ -28,7 +28,7 @@ use cranelift_codegen::ir::{self, InstBuilder, JumpTableData, MemFlags};
 use cranelift_codegen::packed_option::ReservedValue;
 use cranelift_entity::EntityRef;
 use cranelift_frontend::{FunctionBuilder, Variable};
-use environ::{FuncEnvironment, GlobalVariable, WasmError, WasmResult};
+use environ::{FuncEnvironment, GlobalVariable, ReturnMode, WasmError, WasmResult};
 use state::{ControlStackFrame, TranslationState};
 use std::collections::{hash_map, HashMap};
 use std::vec::Vec;
@@ -340,11 +340,10 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             };
             {
                 let args = state.peekn(return_count);
-                if environ.flags().return_at_end() {
-                    builder.ins().jump(br_destination, args);
-                } else {
-                    builder.ins().return_(args);
-                }
+                match environ.return_mode() {
+                    ReturnMode::NormalReturns => builder.ins().return_(args),
+                    ReturnMode::FallthroughReturn => builder.ins().jump(br_destination, args),
+                };
             }
             state.popn(return_count);
             state.reachable = false;

lib/wasm/src/environ/dummy.rs

@@ -7,7 +7,7 @@ use cranelift_codegen::ir::types::*;
 use cranelift_codegen::ir::{self, InstBuilder};
 use cranelift_codegen::settings;
 use cranelift_entity::{EntityRef, PrimaryMap};
-use environ::{FuncEnvironment, GlobalVariable, ModuleEnvironment, WasmResult};
+use environ::{FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode, WasmResult};
 use func_translator::FuncTranslator;
 use std::string::String;
 use std::vec::Vec;
@@ -105,38 +105,55 @@ pub struct DummyEnvironment {
 
     /// Vector of wasm bytecode size for each function.
     pub func_bytecode_sizes: Vec<usize>,
+
+    /// How to return from functions.
+    return_mode: ReturnMode,
 }
 
 impl DummyEnvironment {
     /// Allocates the data structures with default flags.
     pub fn with_triple(triple: Triple) -> Self {
-        Self::with_triple_flags(triple, settings::Flags::new(settings::builder()))
+        Self::with_triple_flags(
+            triple,
+            settings::Flags::new(settings::builder()),
+            ReturnMode::NormalReturns,
+        )
     }
 
-    /// Allocates the data structures with the given flags.
-    pub fn with_triple_flags(triple: Triple, flags: settings::Flags) -> Self {
+    /// Allocates the data structures with the given triple.
+    pub fn with_triple_flags(
+        triple: Triple,
+        flags: settings::Flags,
+        return_mode: ReturnMode,
+    ) -> Self {
         Self {
             info: DummyModuleInfo::with_triple_flags(triple, flags),
             trans: FuncTranslator::new(),
             func_bytecode_sizes: Vec::new(),
+            return_mode,
         }
     }
 
     /// Return a `DummyFuncEnvironment` for translating functions within this
     /// `DummyEnvironment`.
     pub fn func_env(&self) -> DummyFuncEnvironment {
-        DummyFuncEnvironment::new(&self.info)
+        DummyFuncEnvironment::new(&self.info, self.return_mode)
     }
 }
 
 /// The `FuncEnvironment` implementation for use by the `DummyEnvironment`.
 pub struct DummyFuncEnvironment<'dummy_environment> {
     pub mod_info: &'dummy_environment DummyModuleInfo,
+
+    return_mode: ReturnMode,
 }
 
 impl<'dummy_environment> DummyFuncEnvironment<'dummy_environment> {
-    pub fn new(mod_info: &'dummy_environment DummyModuleInfo) -> Self {
-        Self { mod_info }
+    pub fn new(mod_info: &'dummy_environment DummyModuleInfo, return_mode: ReturnMode) -> Self {
+        Self {
+            mod_info,
+            return_mode,
+        }
     }
 
     // Create a signature for `sigidx` amended with a `vmctx` argument after the standard wasm
@@ -321,6 +338,10 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
     ) -> WasmResult<ir::Value> {
         Ok(pos.ins().iconst(I32, -1))
     }
+
+    fn return_mode(&self) -> ReturnMode {
+        self.return_mode
+    }
 }
 
 impl<'data> ModuleEnvironment<'data> for DummyEnvironment {
@@ -433,7 +454,7 @@ impl<'data> ModuleEnvironment<'data> for DummyEnvironment {
 
     fn define_function_body(&mut self, body_bytes: &'data [u8]) -> WasmResult<()> {
         let func = {
-            let mut func_environ = DummyFuncEnvironment::new(&self.info);
+            let mut func_environ = DummyFuncEnvironment::new(&self.info, self.return_mode);
             let func_index =
                 FuncIndex::new(self.get_num_func_imports() + self.info.function_bodies.len());
             let name = get_func_name(func_index);

lib/wasm/src/environ/mod.rs

@@ -5,5 +5,5 @@ mod spec;
 
 pub use environ::dummy::DummyEnvironment;
 pub use environ::spec::{
-    FuncEnvironment, GlobalVariable, ModuleEnvironment, WasmError, WasmResult,
+    FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode, WasmError, WasmResult,
 };

lib/wasm/src/environ/spec.rs

@@ -74,6 +74,15 @@ impl WasmError {
 /// A convenient alias for a `Result` that uses `WasmError` as the error type.
 pub type WasmResult<T> = Result<T, WasmError>;
 
+/// How to return from functions.
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub enum ReturnMode {
+    /// Use normal return instructions as needed.
+    NormalReturns,
+    /// Use a single fallthrough return at the end of the function.
+    FallthroughReturn,
+}
+
 /// Environment affecting the translation of a single WebAssembly function.
 ///
 /// A `FuncEnvironment` trait object is required to translate a WebAssembly function to Cranelift
@@ -216,6 +225,11 @@ pub trait FuncEnvironment {
     fn translate_loop_header(&mut self, _pos: FuncCursor) {
         // By default, don't emit anything.
     }
+
+    /// Should the code be structured to use a single `fallthrough_return` instruction at the end
+    /// of the function body, rather than `return` instructions as needed? This is used by VMs
+    /// to append custom epilogues.
+    fn return_mode(&self) -> ReturnMode;
 }
 
 /// An object satisfying the `ModuleEnvironment` trait can be passed as argument to the

lib/wasm/src/func_translator.rs

@@ -9,7 +9,7 @@ use cranelift_codegen::entity::EntityRef;
 use cranelift_codegen::ir::{self, Ebb, InstBuilder};
 use cranelift_codegen::timing;
 use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
-use environ::{FuncEnvironment, WasmError, WasmResult};
+use environ::{FuncEnvironment, ReturnMode, WasmError, WasmResult};
 use state::TranslationState;
 use wasmparser::{self, BinaryReader};
 
@@ -209,7 +209,10 @@ fn parse_function_body<FE: FuncEnvironment + ?Sized>(
     if state.reachable {
         debug_assert!(builder.is_pristine());
         if !builder.is_unreachable() {
-            builder.ins().return_(&state.stack);
+            match environ.return_mode() {
+                ReturnMode::NormalReturns => builder.ins().return_(&state.stack),
+                ReturnMode::FallthroughReturn => builder.ins().fallthrough_return(&state.stack),
+            };
         }
     }
 

lib/wasm/src/lib.rs

@@ -59,7 +59,8 @@ mod state;
 mod translation_utils;
 
 pub use environ::{
-    DummyEnvironment, FuncEnvironment, GlobalVariable, ModuleEnvironment, WasmError, WasmResult,
+    DummyEnvironment, FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode, WasmError,
+    WasmResult,
 };
 pub use func_translator::FuncTranslator;
 pub use module_translator::translate_module;

lib/wasm/tests/wasm_testsuite.rs

@@ -6,9 +6,9 @@ extern crate wabt;
 
 use cranelift_codegen::isa;
 use cranelift_codegen::print_errors::pretty_verifier_error;
-use cranelift_codegen::settings::{self, Configurable, Flags};
+use cranelift_codegen::settings::{self, Flags};
 use cranelift_codegen::verifier;
-use cranelift_wasm::{translate_module, DummyEnvironment};
+use cranelift_wasm::{translate_module, DummyEnvironment, ReturnMode};
 use std::fs;
 use std::fs::File;
 use std::io;
@@ -35,16 +35,18 @@ fn testsuite() {
     let flags = Flags::new(settings::builder());
     for path in paths {
         let path = path.path();
-        handle_module(&path, &flags);
+        handle_module(&path, &flags, ReturnMode::NormalReturns);
     }
 }
 
 #[test]
-fn return_at_end() {
-    let mut flag_builder = settings::builder();
-    flag_builder.enable("return_at_end").unwrap();
-    let flags = Flags::new(flag_builder);
-    handle_module(Path::new("../../wasmtests/return_at_end.wat"), &flags);
+fn use_fallthrough_return() {
+    let flags = Flags::new(settings::builder());
+    handle_module(
+        Path::new("../../wasmtests/use_fallthrough_return.wat"),
+        &flags,
+        ReturnMode::FallthroughReturn,
+    );
 }
 
 fn read_file(path: &Path) -> io::Result<Vec<u8>> {
@@ -54,7 +56,7 @@ fn read_file(path: &Path) -> io::Result<Vec<u8>> {
     Ok(buf)
 }
 
-fn handle_module(path: &Path, flags: &Flags) {
+fn handle_module(path: &Path, flags: &Flags, return_mode: ReturnMode) {
     let data = match path.extension() {
         None => {
             panic!("the file extension is not wasm or wat");
@@ -73,7 +75,9 @@ fn handle_module(path: &Path, flags: &Flags) {
             None | Some(&_) => panic!("the file extension for {:?} is not wasm or wat", path),
         },
     };
-    let mut dummy_environ = DummyEnvironment::with_triple_flags(triple!("riscv64"), flags.clone());
+    let mut dummy_environ =
+        DummyEnvironment::with_triple_flags(triple!("riscv64"), flags.clone(), return_mode);
+
     translate_module(&data, &mut dummy_environ).unwrap();
 
     let isa = isa::lookup(dummy_environ.info.triple)