From b82a5b3b42cdee6283dc63fe077ac08857629eca Mon Sep 17 00:00:00 2001
From: clearloop <26088946+clearloop@users.noreply.github.com>
Date: Tue, 24 Oct 2023 17:10:35 +0800
Subject: [PATCH] feat(codegen): introduce offset in jump table

---
 codegen/src/backtrace.rs        |  2 +-
 codegen/src/dispatcher.rs       | 13 +++++++++++--
 codegen/src/jump/mod.rs         | 15 +++++++++++----
 codegen/src/jump/relocate.rs    |  7 ++++++-
 codegen/src/jump/table.rs       |  7 +++++++
 codegen/src/visitor/handlers.rs |  2 ++
 6 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/codegen/src/backtrace.rs b/codegen/src/backtrace.rs
index eab6bb15f..e5a4aae97 100644
--- a/codegen/src/backtrace.rs
+++ b/codegen/src/backtrace.rs
@@ -4,7 +4,7 @@ use std::collections::BTreeMap;
 /// Backtrace implementation for the code generation.
 ///
 /// TODO: full implementation (#21)
-#[derive(Default)]
+#[derive(Debug, Default)]
 pub struct Backtrace {
     /// Compiled instructions.
     ///
diff --git a/codegen/src/dispatcher.rs b/codegen/src/dispatcher.rs
index 61ddf9601..3668e557c 100644
--- a/codegen/src/dispatcher.rs
+++ b/codegen/src/dispatcher.rs
@@ -156,15 +156,24 @@ impl Dispatcher {
 
         let func = self.query_func(&abi.name)?;
 
+        // Jump to the end of the dispatcher.
+        //
+        // TODO: detect the bytes of the position.
+        self.asm.increment_sp(1)?;
+        let pc = self.asm.pc_offset();
+        self.table.offset(pc, if last { 0xb } else { 0xc });
+
         if !last {
-            self.asm._dup1()?;
+            self.asm._dup2()?;
         }
 
         self.asm.push(&selector)?;
         self.asm._eq()?;
         self.asm.increment_sp(1)?;
+
         self.table.call(self.asm.pc_offset(), func);
         self.asm._jumpi()?;
+        self.asm._jumpdest()?;
 
         Ok(())
     }
@@ -182,7 +191,7 @@ impl Dispatcher {
 
         let mut len = selectors.len();
         for (_, func) in selectors.iter() {
-            self.emit_selector(func, len == 0)?;
+            self.emit_selector(func, len == 1)?;
             len -= 1;
         }
 
diff --git a/codegen/src/jump/mod.rs b/codegen/src/jump/mod.rs
index e3a9c7af2..b314ad1e7 100644
--- a/codegen/src/jump/mod.rs
+++ b/codegen/src/jump/mod.rs
@@ -9,10 +9,12 @@ mod relocate;
 mod table;
 
 /// Jump types
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Jump {
-    /// Jump to the given label, the label here is the original
-    /// program counter.
+    /// offset to the program counter.
+    Offset(u16),
+    /// Jump to the given label, the label here
+    /// is the original program counter.
     Label(u16),
     /// Jump to function.
     Func(u32),
@@ -23,6 +25,11 @@ pub enum Jump {
 impl Jump {
     /// If the target is a label.
     pub fn is_label(&self) -> bool {
-        matches!(self, Jump::Label(_))
+        matches!(self, Jump::Label { .. })
+    }
+
+    /// If the target is fixed to offset of the program counter.
+    pub fn is_offset(&self) -> bool {
+        matches!(self, Jump::Offset(_))
     }
 }
diff --git a/codegen/src/jump/relocate.rs b/codegen/src/jump/relocate.rs
index 8a2e99a81..4f7ded928 100644
--- a/codegen/src/jump/relocate.rs
+++ b/codegen/src/jump/relocate.rs
@@ -18,7 +18,12 @@ impl JumpTable {
             tracing::debug!("run relocateion for {jump:?}",);
 
             let offset = relocate::offset(pc)?;
-            relocate::pc(buffer, pc, self.target(&jump)?, offset)?;
+            let mut target = self.target(&jump)?;
+            if jump.is_offset() {
+                target += pc;
+            }
+
+            relocate::pc(buffer, pc, target, offset)?;
             self.shift_label_pc(pc, offset)?;
         }
 
diff --git a/codegen/src/jump/table.rs b/codegen/src/jump/table.rs
index e34db59d9..e8e51e8b7 100644
--- a/codegen/src/jump/table.rs
+++ b/codegen/src/jump/table.rs
@@ -7,6 +7,7 @@ use crate::{
 use std::collections::BTreeMap;
 
 /// Jump table implementation.
+///
 #[derive(Default, Debug)]
 pub struct JumpTable {
     /// Jump table.
@@ -49,6 +50,11 @@ impl JumpTable {
         self.jump.insert(pc, Jump::Label(label));
     }
 
+    /// Register a label.
+    pub fn offset(&mut self, pc: u16, offset: u16) {
+        self.jump.insert(pc, Jump::Offset(offset));
+    }
+
     /// Merge two jump tables.
     ///
     /// Merge other jump table into this one, update the program
@@ -80,6 +86,7 @@ impl JumpTable {
     /// Get the target of a jump.
     pub fn target(&mut 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::ExtNotFound(*ext))?),
diff --git a/codegen/src/visitor/handlers.rs b/codegen/src/visitor/handlers.rs
index 060f09397..17e3f4dbf 100644
--- a/codegen/src/visitor/handlers.rs
+++ b/codegen/src/visitor/handlers.rs
@@ -52,6 +52,8 @@ impl CodeGen {
 
         // TODO: handle the length of results > u8::MAX.
         self.masm.shift_pc(len, false)?;
+
+        // TODO: [PC, PUSHN, BYTES, JUMPDEST]
         self.masm.push(&[0x04])?;
         self.masm._add()?;
         self.masm._jump()?;