refactor(codegen): code relocation for code length greater than 0xff (#…

…276)

* feat(codegen): grow jump targets for label offsets on pc greater than 0xff

* fix(codegen): detect pushN from the bytes of target pc

* fix(codegen): shift label pc with relocation result

* feat(codegen): handle offset label for relocation

* feat(codegen): handle offset label for relocation

* feat(codegen): handle offset target for relocation

* fix(codegen): index checks of variables for external functions

* feat(codegen): clean code of offset relocation

codegen/src/asm.rs

@@ -117,7 +117,6 @@ impl Assembler {
     /// Mock the stack input and output for checking
     /// the stack usages.
     pub fn emit_op(&mut self, opcode: OpCode) -> Result<()> {
-        // tracing::trace!("stack length: {:?}", self.sp);
         tracing::trace!("emit opcode: {:?}", opcode);
         self.decrement_sp(opcode.stack_in() as u8)?;
         self.emit(opcode.into());

codegen/src/jump/mod.rs

@@ -1,6 +1,7 @@
 //! Jump table implementation.
 
 use crate::codegen::ExtFunc;
+use core::fmt::Display;
 pub use table::JumpTable;
 
 mod pc;
@@ -21,6 +22,17 @@ pub enum Jump {
     ExtFunc(ExtFunc),
 }
 
+impl Display for Jump {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            Jump::Offset(offset) => write!(f, "Offset(0x{offset:x})"),
+            Jump::Label(offset) => write!(f, "Label(0x{offset:x})"),
+            Jump::Func(index) => write!(f, "Func({index})"),
+            Jump::ExtFunc(_) => write!(f, "ExtFunc"),
+        }
+    }
+}
+
 impl Jump {
     /// If the target is a label.
     pub fn is_label(&self) -> bool {
@@ -31,4 +43,9 @@ impl Jump {
     pub fn is_offset(&self) -> bool {
         matches!(self, Jump::Offset(_))
     }
+
+    /// If the target is a function call
+    pub fn is_call(&self) -> bool {
+        !self.is_label() && !self.is_offset()
+    }
 }

codegen/src/jump/pc.rs

@@ -1,7 +1,7 @@
 //! Program counter handlers.
 
 use crate::{
-    jump::{relocate, Jump, JumpTable},
+    jump::{Jump, JumpTable},
     Error, Result, BUFFER_LIMIT,
 };
 
@@ -11,9 +11,44 @@ impl JumpTable {
         tracing::trace!("shift pc from 0x{start:x} with offset={offset}");
         self.shift_label_pc(start, offset)?;
         self.shift_label_target(start, offset)?;
-        self.shift_func_target(start, offset)?;
+        self.shift_func_target(start, offset)
+    }
 
-        Ok(())
+    /// Shift the target program counters.
+    ///
+    /// Calculating target pc from the offset of original pc.
+    pub fn shift_targets(&mut self) -> Result<()> {
+        let mut total_offset = 0;
+        self.jump
+            .clone()
+            .iter()
+            .try_for_each(|(original_pc, target)| -> Result<()> {
+                tracing::debug!("shift targets for {target} <- (0x{original_pc:x})");
+                let pc = original_pc + total_offset;
+                let offset = self.target_offset(
+                    target,
+                    if target.is_offset() {
+                        total_offset + original_pc
+                    } else {
+                        total_offset
+                    },
+                )?;
+
+                self.shift_target(pc, offset)?;
+                total_offset += offset;
+                Ok(())
+            })
+    }
+
+    /// Shift the program counter of targets with given ptr and offset.
+    ///
+    /// 1. shift code section.
+    /// 2. shift label targets.
+    /// 3. shift function targets.
+    pub fn shift_target(&mut self, ptr: u16, offset: u16) -> Result<()> {
+        self.code.shift(offset);
+        self.shift_label_target(ptr, offset)?;
+        self.shift_func_target(ptr, offset)
     }
 
     /// Shift program counter for labels.
@@ -23,12 +58,14 @@ impl JumpTable {
             .iter()
             .map(|(k, v)| {
                 let mut k = *k;
+                let next_label = k + offset;
                 if k > start {
                     tracing::trace!(
-                        "shift {v:x?} pc with offset={offset}: 0x{k:x}(0x{start:x}) -> 0x{:x}",
-                        k + offset
+                        "shift {v} pc with offset={offset}: 0x{k:x}(0x{start:x}) -> 0x{:x}",
+                        next_label
                     );
-                    k += offset;
+                    k = next_label;
+
                     if k > BUFFER_LIMIT as u16 {
                         return Err(Error::InvalidPC(k as usize));
                     }
@@ -41,63 +78,55 @@ impl JumpTable {
         Ok(())
     }
 
-    /// Shift the target program counters.
-    ///
-    /// Calculating target pc from the offset of original pc.
-    pub fn shift_targets(&mut self) -> Result<()> {
-        let mut total_offset = 0;
-        self.jump
-            .clone()
-            .keys()
-            .try_for_each(|original_pc| -> Result<()> {
-                let pc = original_pc + total_offset;
-                let offset = relocate::offset(pc)?;
-                total_offset += offset;
-                self.shift_target(pc, offset)
-            })
-    }
+    /// Shift program counter for functions.
+    pub fn shift_func_target(&mut self, ptr: u16, offset: u16) -> Result<()> {
+        if self.func.is_empty() {
+            tracing::trace!("No functions to shift.");
+            return Ok(());
+        }
 
-    /// Shift the program counter of targets with given ptr and offset.
-    ///
-    /// 1. shift code section.
-    /// 2. shift label targets.
-    /// 3. shift function targets.
-    pub fn shift_target(&mut self, ptr: u16, offset: u16) -> Result<()> {
-        self.code.shift(offset);
-        self.shift_label_target(ptr, offset)?;
-        self.shift_func_target(ptr, offset)
-    }
+        self.func.iter_mut().try_for_each(|(index, target)| {
+            let next_target = *target + offset;
 
-    /// Shift program counter for functions.
-    pub fn shift_func_target(&mut self, start: u16, offset: u16) -> Result<()> {
-        self.func.iter_mut().try_for_each(|(k, v)| {
-            if *v > start {
+            if *target > ptr {
                 tracing::trace!(
-                    "shift Func({k}) target with offset={offset}: 0x{v:x}(0x{start:x}) -> 0x{:x}",
-                    *v + offset
+                    "shift Func({index}) target with offset={offset}: 0x{target:x}(0x{ptr:x}) -> 0x{:x}",
+                    next_target
+                );
+
+                *target = next_target;
+            } else {
+                tracing::trace!(
+                    "shift Func({index}) target with offset=0: 0x{target:x}(0x{ptr:x}) -> 0x{target:x}"
                 );
-                *v += offset;
-                if *v > BUFFER_LIMIT as u16 {
-                    return Err(Error::InvalidPC(*v as usize));
-                }
             }
 
             Ok(())
-        })?;
-
-        Ok(())
+        })
     }
 
     /// Shift target program counter for labels.
     pub fn shift_label_target(&mut self, ptr: u16, offset: u16) -> Result<()> {
+        if self.jump.is_empty() {
+            tracing::trace!("No labels to shift.");
+            return Ok(());
+        }
+
         self.jump.iter_mut().try_for_each(|(pc, target)| {
             if let Jump::Label(label) = target {
+                let next_label = *label + offset;
+
                 if *label > ptr {
                     tracing::trace!(
-                        "shift Label(pc=0x{pc:x}) target with offset={offset} 0x{label:x}(0x{ptr:x}) -> 0x{:x}",
-                        *label + offset
+                        "shift Label(0x{pc:x}) target with offset={offset}: 0x{label:x}(0x{ptr:x}) -> 0x{:x}",
+                        next_label,
+                    );
+
+                    *label = next_label;
+                } else {
+                    tracing::trace!(
+                        "shift Label(0x{pc:x}) target with offset=0: 0x{label:x}(0x{ptr:x}) -> 0x{label:x}"
                     );
-                    *label += offset;
                 }
             }
 

codegen/src/jump/relocate.rs

@@ -1,7 +1,7 @@
 //! Program Relocations
 
 use crate::{
-    jump::{relocate, JumpTable},
+    jump::{relocate, Jump, JumpTable},
     wasm::ToLSBytes,
     Buffer, Error, Result, BUFFER_LIMIT,
 };
@@ -13,72 +13,98 @@ impl JumpTable {
     /// *WARNING*: This function should only be called once in the compiler.
     /// considering move it to the compiler.
     pub fn relocate(&mut self, buffer: &mut Buffer) -> Result<()> {
+        // pre-calculate and shift targets
         self.shift_targets()?;
         tracing::trace!("code section offset: 0x{:x}", self.code.offset());
 
+        // relocate functions
         while let Some((pc, jump)) = self.jump.pop_first() {
-            tracing::trace!("run relocation for {jump:?}");
-
-            let offset = relocate::offset(pc)?;
+            tracing::debug!("run relocation for {jump}");
             let mut target = self.target(&jump)?;
-            if jump.is_offset() {
-                target += pc;
-            }
+            self.relocate_offset(pc, jump, &mut target)?;
 
-            relocate::pc(buffer, pc, target, offset)?;
-            self.shift_label_pc(pc, offset)?;
+            let offset = relocate::pc(buffer, pc, target)?;
+            self.shift_label_pc(pc, offset as u16)?;
         }
 
         buffer.extend_from_slice(&self.code.finish());
         Ok(())
     }
-}
 
-/// Get the offset of the program counter for relocation.
-pub fn offset(original_pc: u16) -> Result<u16> {
-    let pc = original_pc;
-    let mut offset = 0;
-
-    // Update the target program counter
-    {
-        // The maximum size of the PC is 2 bytes, whatever PUSH1 or PUSH2
-        // takes 1 more byte.
-        offset += 1;
-
-        // Update the program counter for the edge cases.
-        //
-        // Start from 0xff, the lowest significant bytes of the target
-        // program counter will take 2 bytes instead of 1 byte.
-        //
-        // | PC   | PC BYTES | TARGET PC |
-        // |------|----------|-----------|
-        // | 0xfe | 1        |      0xff |
-        // | 0xff | 2        |     0x101 |
-        offset += if pc > 0xfe { 2 } else { 1 }
-    }
+    /// relocate the target of offset jump
+    fn relocate_offset(&self, pc: u16, jump: Jump, target: &mut u16) -> Result<()> {
+        if !jump.is_offset() {
+            return Ok(());
+        }
 
-    // Check PC range.
-    if pc + offset > BUFFER_LIMIT as u16 {
-        return Err(Error::InvalidPC((pc + offset) as usize));
-    }
+        // NOTE: If the target is offset the return data is
+        // the offset instead of the PC.
+        *target += pc;
+
+        // check the original pc of the offset is greater than 0xff
+        if pc > 0xff {
+            *target += 1;
+        }
 
-    Ok(offset)
+        // check if the offset of the embedded call is greater than 0xff
+        if let Some((_, next_target)) = self.jump.first_key_value() {
+            if next_target.is_call() && self.target(next_target)? > 0xff {
+                *target += 1
+            }
+        }
+
+        Ok(())
+    }
 }
 
+// /// Get the offset of the program counter for relocation.
+// pub fn offset(pc: u16) -> Result<u16> {
+//     let mut offset = 0;
+//
+//     // Update the target program counter
+//     {
+//         // The maximum size of the PC is 2 bytes, whatever PUSH1 or PUSH32
+//         // takes 1 more byte.
+//         offset += 1;
+//
+//         // Update the program counter for the edge cases.
+//         //
+//         // Start from 0xff, the lowest significant bytes of the target
+//         // program counter will take 2 bytes instead of 1 byte.
+//         //
+//         // | PC   | PC BYTES | TARGET PC |
+//         // |------|----------|-----------|
+//         // | 0xfe | 1        |      0xff |
+//         // | 0xff | 2        |    0x0101 |
+//         offset += if pc > 0xfe { 2 } else { 1 }
+//     }
+//
+//     // Check PC range.
+//     if pc + offset > BUFFER_LIMIT as u16 {
+//         return Err(Error::InvalidPC((pc + offset) as usize));
+//     }
+//
+//     Ok(offset)
+// }
+
 /// Relocate program counter to buffer.
-fn pc(buffer: &mut Buffer, original_pc: u16, target_pc: u16, offset: u16) -> Result<()> {
+fn pc(buffer: &mut Buffer, original_pc: u16, target_pc: u16) -> Result<usize> {
     let original_pc = original_pc as usize;
     let mut new_buffer: Buffer = buffer[..original_pc].into();
     let rest_buffer: Buffer = buffer[original_pc..].into();
 
-    if offset == 2 {
+    let pc_offset = target_pc.to_ls_bytes();
+    if pc_offset.len() == 1 {
         new_buffer.push(OpCode::PUSH1.into());
     } else {
         new_buffer.push(OpCode::PUSH2.into());
     }
 
-    let pc_offset = target_pc.to_ls_bytes();
-    tracing::trace!("push bytes: 0x{:x?} at 0x{:x}", pc_offset, original_pc);
+    tracing::trace!(
+        "push bytes: 0x{} at 0x{}",
+        hex::encode(&pc_offset),
+        hex::encode(original_pc.to_ls_bytes())
+    );
     new_buffer.extend_from_slice(&pc_offset);
     new_buffer.extend_from_slice(&rest_buffer);
 
@@ -88,5 +114,5 @@ fn pc(buffer: &mut Buffer, original_pc: u16, target_pc: u16, offset: u16) -> Res
     }
 
     *buffer = new_buffer;
-    Ok(())
+    Ok(1 + pc_offset.len())
 }

codegen/src/jump/table.rs

@@ -81,12 +81,20 @@ impl JumpTable {
     }
 
     /// Get the target of a jump.
-    pub fn target(&mut self, jump: &Jump) -> Result<u16> {
+    pub fn target(&self, jump: &Jump) -> Result<u16> {
         match jump {
             Jump::Offset(offset) => Ok(*offset),
             Jump::Label(label) => Ok(*label),
             Jump::Func(func) => Ok(*self.func.get(func).ok_or(Error::FuncNotFound(*func))?),
             Jump::ExtFunc(ext) => Ok(self.code.offset_of(ext).ok_or(Error::ExtFuncNotFound)?),
         }
     }
+
+    /// Get the target offset
+    pub fn target_offset(&self, jump: &Jump, offset: u16) -> Result<u16> {
+        let target = self.target(jump)?;
+
+        // [PUSH_N, PC]
+        Ok(if (target + offset) > 0xff { 3 } else { 2 })
+    }
 }

codegen/src/masm/ret.rs

@@ -34,6 +34,7 @@ impl MacroAssembler {
         let len = results.len() as u8;
 
         tracing::trace!("cleaning frame stack, target: {}", len + 1);
+        // TODO: clean stacks via the count of nested control stacks.
         while self.sp() > len + 1 {
             self._drop()?;
         }