Fix Wasm translator bug: end of toplevel frame is branched-to only for fallthrough returns.

cranelift/wasm/src/code_translator.rs

@@ -514,7 +514,9 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
         Operator::Return => {
             let (return_count, br_destination) = {
                 let frame = &mut state.control_stack[0];
-                frame.set_branched_to_exit();
+                if environ.return_mode() == ReturnMode::FallthroughReturn {
+                    frame.set_branched_to_exit();
+                }
                 let return_count = frame.num_return_values();
                 (return_count, frame.br_destination())
             };

cranelift/wasm/src/environ/dummy.rs

@@ -10,6 +10,7 @@ use crate::environ::{
     WasmFuncType, WasmResult,
 };
 use crate::func_translator::FuncTranslator;
+use crate::state::FuncTranslationState;
 use crate::translation_utils::{
     DataIndex, DefinedFuncIndex, ElemIndex, FuncIndex, Global, GlobalIndex, Memory, MemoryIndex,
     Table, TableIndex, TypeIndex,
@@ -25,7 +26,7 @@ use cranelift_frontend::FunctionBuilder;
 use std::boxed::Box;
 use std::string::String;
 use std::vec::Vec;
-use wasmparser::{FuncValidator, FunctionBody, ValidatorResources, WasmFeatures};
+use wasmparser::{FuncValidator, FunctionBody, Operator, ValidatorResources, WasmFeatures};
 
 /// Compute a `ir::ExternalName` for a given wasm function index.
 fn get_func_name(func_index: FuncIndex) -> ir::ExternalName {
@@ -111,6 +112,31 @@ impl DummyModuleInfo {
     }
 }
 
+/// State for tracking and checking reachability at each operator. Used for unit testing with the
+/// `DummyEnvironment`.
+#[derive(Clone)]
+pub struct ExpectedReachability {
+    /// Before- and after-reachability
+    reachability: Vec<(bool, bool)>,
+    before_idx: usize,
+    after_idx: usize,
+}
+
+impl ExpectedReachability {
+    fn check_before(&mut self, reachable: bool) {
+        assert_eq!(reachable, self.reachability[self.before_idx].0);
+        self.before_idx += 1;
+    }
+    fn check_after(&mut self, reachable: bool) {
+        assert_eq!(reachable, self.reachability[self.after_idx].1);
+        self.after_idx += 1;
+    }
+    fn check_end(&self) {
+        assert_eq!(self.before_idx, self.reachability.len());
+        assert_eq!(self.after_idx, self.reachability.len());
+    }
+}
+
 /// This `ModuleEnvironment` implementation is a "naïve" one, doing essentially nothing and
 /// emitting placeholders when forced to. Don't try to execute code translated for this
 /// environment, essentially here for translation debug purposes.
@@ -135,6 +161,9 @@ pub struct DummyEnvironment {
 
     /// Function names.
     function_names: SecondaryMap<FuncIndex, String>,
+
+    /// Expected reachability data (before/after for each op) to assert. This is used for testing.
+    expected_reachability: Option<ExpectedReachability>,
 }
 
 impl DummyEnvironment {
@@ -148,13 +177,18 @@ impl DummyEnvironment {
             debug_info,
             module_name: None,
             function_names: SecondaryMap::new(),
+            expected_reachability: None,
         }
     }
 
     /// Return a `DummyFuncEnvironment` for translating functions within this
     /// `DummyEnvironment`.
     pub fn func_env(&self) -> DummyFuncEnvironment {
-        DummyFuncEnvironment::new(&self.info, self.return_mode)
+        DummyFuncEnvironment::new(
+            &self.info,
+            self.return_mode,
+            self.expected_reachability.clone(),
+        )
     }
 
     fn get_func_type(&self, func_index: FuncIndex) -> TypeIndex {
@@ -171,20 +205,39 @@ impl DummyEnvironment {
     pub fn get_func_name(&self, func_index: FuncIndex) -> Option<&str> {
         self.function_names.get(func_index).map(String::as_ref)
     }
+
+    /// Test reachability bits before and after every opcode during translation, as provided by the
+    /// `FuncTranslationState`. This is generally used only for unit tests. This is applied to
+    /// every function in the module (so is likely only useful for test modules with one function).
+    pub fn test_expected_reachability(&mut self, reachability: Vec<(bool, bool)>) {
+        self.expected_reachability = Some(ExpectedReachability {
+            reachability,
+            before_idx: 0,
+            after_idx: 0,
+        });
+    }
 }
 
 /// The `FuncEnvironment` implementation for use by the `DummyEnvironment`.
 pub struct DummyFuncEnvironment<'dummy_environment> {
     pub mod_info: &'dummy_environment DummyModuleInfo,
 
     return_mode: ReturnMode,
+
+    /// Expected reachability data (before/after for each op) to assert. This is used for testing.
+    expected_reachability: Option<ExpectedReachability>,
 }
 
 impl<'dummy_environment> DummyFuncEnvironment<'dummy_environment> {
-    pub fn new(mod_info: &'dummy_environment DummyModuleInfo, return_mode: ReturnMode) -> Self {
+    pub fn new(
+        mod_info: &'dummy_environment DummyModuleInfo,
+        return_mode: ReturnMode,
+        expected_reachability: Option<ExpectedReachability>,
+    ) -> Self {
         Self {
             mod_info,
             return_mode,
+            expected_reachability,
         }
     }
 
@@ -307,6 +360,41 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
         }))
     }
 
+    fn before_translate_operator(
+        &mut self,
+        _op: &Operator,
+        _builder: &mut FunctionBuilder,
+        state: &FuncTranslationState,
+    ) -> WasmResult<()> {
+        if let Some(ref mut r) = &mut self.expected_reachability {
+            r.check_before(state.reachable());
+        }
+        Ok(())
+    }
+
+    fn after_translate_operator(
+        &mut self,
+        _op: &Operator,
+        _builder: &mut FunctionBuilder,
+        state: &FuncTranslationState,
+    ) -> WasmResult<()> {
+        if let Some(ref mut r) = &mut self.expected_reachability {
+            r.check_after(state.reachable());
+        }
+        Ok(())
+    }
+
+    fn after_translate_function(
+        &mut self,
+        _builder: &mut FunctionBuilder,
+        _state: &FuncTranslationState,
+    ) -> WasmResult<()> {
+        if let Some(ref mut r) = &mut self.expected_reachability {
+            r.check_end();
+        }
+        Ok(())
+    }
+
     fn translate_call_indirect(
         &mut self,
         mut pos: FuncCursor,
@@ -746,7 +834,11 @@ impl<'data> ModuleEnvironment<'data> for DummyEnvironment {
         self.func_bytecode_sizes
             .push(body.get_binary_reader().bytes_remaining());
         let func = {
-            let mut func_environ = DummyFuncEnvironment::new(&self.info, self.return_mode);
+            let mut func_environ = DummyFuncEnvironment::new(
+                &self.info,
+                self.return_mode,
+                self.expected_reachability.clone(),
+            );
             let func_index =
                 FuncIndex::new(self.get_num_func_imports() + self.info.function_bodies.len());
             let name = get_func_name(func_index);

cranelift/wasm/tests/wasm_testsuite.rs

@@ -94,3 +94,86 @@ fn handle_module(data: Vec<u8>, flags: &Flags, return_mode: ReturnMode) {
             .unwrap();
     }
 }
+
+#[test]
+fn reachability_is_correct() {
+    let tests = vec![
+        (
+            ReturnMode::NormalReturns,
+            r#"
+        (module (func (param i32)
+         (loop
+          (block
+           local.get 0
+           br_if 0
+           br 1))))"#,
+            vec![
+                (true, true),  // Loop
+                (true, true),  // Block
+                (true, true),  // LocalGet
+                (true, true),  // BrIf
+                (true, false), // Br
+                (false, true), // End
+                (true, true),  // End
+                (true, true),  // End
+            ],
+        ),
+        (
+            ReturnMode::NormalReturns,
+            r#"
+        (module (func (param i32)
+         (loop
+          (block
+           br 1
+           nop))))"#,
+            vec![
+                (true, true),   // Loop
+                (true, true),   // Block
+                (true, false),  // Br
+                (false, false), // Nop
+                (false, false), // Nop
+                (false, false), // Nop
+                (false, false), // End
+            ],
+        ),
+        (
+            ReturnMode::NormalReturns,
+            r#"
+        (module (func (param i32) (result i32)
+          i32.const 1
+          return
+          i32.const 42))"#,
+            vec![
+                (true, true),   // I32Const
+                (true, false),  // Return
+                (false, false), // I32Const
+                (false, false), // End
+            ],
+        ),
+        (
+            ReturnMode::FallthroughReturn,
+            r#"
+        (module (func (param i32) (result i32)
+         i32.const 1
+         return
+         i32.const 42))"#,
+            vec![
+                (true, true),   // I32Const
+                (true, false),  // Return
+                (false, false), // I32Const
+                (false, true),  // End
+            ],
+        ),
+    ];
+
+    for (return_mode, wat, expected_reachability) in tests {
+        println!("testing wat:\n{}", wat);
+        let flags = Flags::new(settings::builder());
+        let triple = triple!("riscv64");
+        let isa = isa::lookup(triple).unwrap().finish(flags.clone());
+        let mut env = DummyEnvironment::new(isa.frontend_config(), return_mode, false);
+        env.test_expected_reachability(expected_reachability);
+        let data = wat::parse_str(wat).unwrap();
+        translate_module(data.as_ref(), &mut env).unwrap();
+    }
+}