fix(codegen): use label for control offsets

codegen/src/codegen/dispatcher.rs

@@ -75,7 +75,7 @@ impl Dispatcher {
         self.asm.increment_sp(1)?;
 
         // Prepare the `PC` of the callee function.
-        self.table.call(self.asm.pc_offset(), func);
+        self.table.call(self.asm.pc(), func);
 
         if last {
             self.asm._swap1()?;

codegen/src/masm/mod.rs

@@ -80,7 +80,7 @@ impl MacroAssembler {
     }
 
     /// Get the current program counter offset.
-    pub fn pc_offset(&self) -> u16 {
+    pub fn pc(&self) -> u16 {
         self.asm.buffer().len() as u16
     }
 

codegen/src/visitor/call.rs

@@ -78,10 +78,8 @@ impl Function {
         let base_offset = 5 + ((params + reserved) * 0x20).saturating_sub(0xff) / 0x20;
 
         // Move the PC before the parameters in the stack.
-        self.table.offset(
-            self.masm.pc_offset(),
-            base_offset as u16 + 4 * (*params as u16),
-        );
+        self.table
+            .offset(self.masm.pc(), base_offset as u16 + 4 * (*params as u16));
         self.masm.increment_sp(1)?;
 
         // Adjust the stack to place the PC before the parameters.
@@ -95,7 +93,7 @@ impl Function {
         }
 
         // Register the call index in the jump table.
-        self.table.call(self.masm.pc_offset(), index);
+        self.table.call(self.masm.pc(), index);
 
         // Jump to the callee function.
         self.masm._jump()?;

codegen/src/visitor/control.rs

@@ -15,7 +15,7 @@ impl Function {
         // push an `If` frame to the control stack
         let frame = ControlStackFrame::new(
             ControlStackFrameType::If(false),
-            self.masm.pc_offset(),
+            self.masm.pc(),
             self.masm.sp(),
             blockty,
         );
@@ -35,7 +35,7 @@ impl Function {
     pub fn _block(&mut self, blockty: BlockType) -> Result<()> {
         let frame = ControlStackFrame::new(
             ControlStackFrameType::Block,
-            self.masm.pc_offset(),
+            self.masm.pc(),
             self.masm.sp(),
             blockty,
         );
@@ -50,7 +50,7 @@ impl Function {
     pub fn _loop(&mut self, blockty: BlockType) -> Result<()> {
         let frame = ControlStackFrame::new(
             ControlStackFrameType::Loop,
-            self.masm.pc_offset(),
+            self.masm.pc(),
             self.masm.sp(),
             blockty,
         );
@@ -68,7 +68,7 @@ impl Function {
         // push an `Else` frame to the control stack.
         let frame = ControlStackFrame::new(
             ControlStackFrameType::Else,
-            self.masm.pc_offset(),
+            self.masm.pc(),
             self.masm.sp(),
             last_frame.result(),
         );
@@ -78,7 +78,7 @@ impl Function {
 
         // mark else as the jump destination of the if block.
         self.table
-            .label(last_frame.original_pc_offset, self.masm.pc_offset());
+            .label(last_frame.original_pc_offset, self.masm.pc());
         self.masm._jumpdest()?;
 
         Ok(())
@@ -92,8 +92,7 @@ impl Function {
         tracing::trace!("select");
         self.masm._iszero()?;
         self.masm.increment_sp(1)?;
-        self.table
-            .label(self.masm.pc_offset(), self.masm.pc_offset() + 4);
+        self.table.label(self.masm.pc(), self.masm.pc() + 2);
         self.masm._jumpi()?;
         self.masm._drop()?;
         self.masm._jumpdest()?;
@@ -114,7 +113,7 @@ impl Function {
     /// Conditional branch to a given label in an enclosing construct.
     pub fn _br_if(&mut self, depth: u32) -> Result<()> {
         let label = self.control.label_from_depth(depth)?;
-        self.table.label(self.masm.pc_offset(), label);
+        self.table.label(self.masm.pc(), label);
         self.masm.asm.increment_sp(1)?;
         self.masm._jumpi()?;
 
@@ -173,13 +172,12 @@ impl Function {
             ControlStackFrameType::Block => self.masm._jumpdest(),
             ControlStackFrameType::Loop => Ok(()),
             _ => {
-                self.table
-                    .label(frame.original_pc_offset, self.masm.pc_offset());
+                self.table.label(frame.original_pc_offset, self.masm.pc());
 
                 // TODO: Check the stack output and make decisions
                 // how to handle the results.
 
-                // Emit JUMPDEST after at the end of the control flow.
+                // Emit JUMPDEST at the end of the control flow.
                 self.masm._jumpdest()
             }
         }