encode staged in OpEncoder::pin_label

Author: Robbepop

crates/wasmi/src/engine/translator/func/encoder.rs

@@ -262,10 +262,11 @@ impl OpEncoder {
     /// # Panics
     ///
     /// If there is a staged [`Op`].
-    pub fn pin_label(&mut self, lref: LabelRef) {
-        assert!(self.staged.is_none());
+    pub fn pin_label(&mut self, lref: LabelRef) -> Result<(), Error> {
+        self.try_encode_staged()?;
         let next_pos = Pos::from(self.ops.next_pos());
         self.labels.pin_label(lref, next_pos);
+        Ok(())
     }
 
     /// Pins the [`Label`] at `lref` to the current encoded bytestream position if unpinned.

crates/wasmi/src/engine/translator/func/mod.rs

@@ -694,11 +694,6 @@ impl FuncTranslator {
         !can_avoid_copies
     }
 
-    /// Pins the `label` to the next [`Pos<Op>`].
-    fn pin_label(&mut self, label: LabelRef) {
-        self.instrs.pin_label(label);
-    }
-
     /// Convert the [`Operand`] at `depth` into an [`Operand::Temp`] by copying if necessary.
     ///
     /// # Note
@@ -1005,15 +1000,11 @@ impl FuncTranslator {
             }
             self.push_frame_results(&frame)?;
         }
-        self.instrs.pin_label(frame.label());
+        self.instrs.pin_label(frame.label())?;
         self.reachable |= frame.is_branched_to();
         if self.reachable && self.stack.is_control_empty() {
             self.encode_return(consume_fuel_instr)?;
         }
-        if frame.is_branched_to() {
-            // No need to reset `last_instr` if there was no branch to the end of a Wasm `block`.
-            self.instrs.try_encode_staged()?;
-        }
         Ok(())
     }
 
@@ -1062,10 +1053,8 @@ impl FuncTranslator {
             self.copy_branch_params(&frame, consume_fuel_instr)?;
         }
         self.push_frame_results(&frame)?;
-        self.instrs.pin_label(frame.label());
+        self.instrs.pin_label(frame.label())?;
         self.reachable = true;
-        // Need to reset `last_instr` since end of `if` is a control flow boundary.
-        self.instrs.try_encode_staged()?;
         Ok(())
     }
 
@@ -1094,10 +1083,8 @@ impl FuncTranslator {
             self.copy_branch_params(&frame, consume_fuel_instr)?;
         }
         self.push_frame_results(&frame)?;
-        self.instrs.pin_label(frame.label());
+        self.instrs.pin_label(frame.label())?;
         self.reachable = reachable;
-        // Need to reset `last_instr` since end of `else` is a control flow boundary.
-        self.instrs.try_encode_staged()?;
         Ok(())
     }
 
@@ -1114,13 +1101,8 @@ impl FuncTranslator {
             }
             self.push_frame_results(&frame)?;
         }
-        self.instrs.pin_label(frame.label());
+        self.instrs.pin_label(frame.label())?;
         self.reachable = end_is_reachable || frame.is_branched_to();
-        if frame.is_branched_to() {
-            // No need to reset `last_instr` if there was no branch to the
-            // end of a Wasm `if` where only `then` or `else` is reachable.
-            self.instrs.try_encode_staged()?;
-        }
         Ok(())
     }
 

crates/wasmi/src/engine/translator/func/visit.rs

@@ -123,12 +123,10 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
         let continue_label = self.instrs.new_label();
         let consume_fuel = self.stack.consume_fuel_instr();
         self.move_operands_to_temp(usize::from(len_params), consume_fuel)?;
-        self.pin_label(continue_label);
+        self.instrs.pin_label(continue_label)?;
         let consume_fuel = self.instrs.encode_consume_fuel()?;
         self.stack
             .push_loop(block_ty, continue_label, consume_fuel)?;
-        // Need to reset `last_instr` because a loop header is a control flow boundary.
-        self.instrs.try_encode_staged()?;
         Ok(())
     }
 
@@ -191,8 +189,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
                 self.encode_br(frame.label())?;
             }
             // Start of `else` block:
-            self.instrs.try_encode_staged()?;
-            self.instrs.pin_label(else_label);
+            self.instrs.pin_label(else_label)?;
         }
         let consume_fuel_instr = self.instrs.encode_consume_fuel()?;
         self.reachable = frame.is_else_reachable();
@@ -273,7 +270,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
             self.encode_copies(branch_results, len_branch_params, consume_fuel_instr)?;
         }
         self.encode_br(label)?;
-        self.instrs.pin_label(skip_label);
+        self.instrs.pin_label(skip_label)?;
         Ok(())
     }