feat: Support lua bytecode encoding

lib/src/backend/lua/bytecode.rs

@@ -1,9 +1,135 @@
-use crate::backend::lua::encoding::LuaOpCode;
-
 use indexmap::IndexSet;
 use std::fmt::{Display, Formatter, Write};
 use std::hash::{Hash, Hasher};
 
+macro_rules! excess_k {
+    ($v: expr, $k: expr) => {
+        ($v as u32).wrapping_add(((2u32.pow($k) - 1) / 2)) & (2u32.pow($k) - 1)
+    };
+}
+
+macro_rules! excess_sBx {
+    ($v: expr) => { excess_k!($v, 17) };
+}
+
+#[repr(u8)]
+#[allow(non_camel_case_types)]
+enum LuaOpCode {
+    // A B, R[A] := R[B]
+    OP_MOVE = 0,
+    // A sBx, R[A] := sBx
+    OP_LOADI = 1,
+    // A sBx, R[A] := (lua_Number)xBx
+    OP_LOADF = 2,
+    // A Bx, R[A] := K[Bx]
+    OP_LOADK = 3,
+    // A, R[A] := K[extra arg]
+    OP_LOADKX = 4,
+    OP_LOADFALSE = 5,
+    OP_LFALSESKIP = 6,
+    OP_LOADTRUE = 7,
+    OP_LOADNIL = 8,
+    OP_GETUPVAL = 9,
+    OP_SETUPVAL = 10,
+
+    OP_GETTABUP = 11,
+    OP_GETTABLE = 12,
+    OP_GETI = 13,
+    OP_GETFIELD = 14,
+
+    OP_SETTABUP = 15,
+    OP_SETTABLE = 16,
+    OP_SETI = 17,
+    OP_SETFIELD = 18,
+
+    OP_NEWTABLE = 19,
+
+    OP_SELF = 20,
+
+    OP_ADDI = 21,
+
+    OP_ADDK = 22,
+    OP_SUBK = 23,
+    OP_MULK = 24,
+    OP_MODK = 25,
+    OP_POWK = 26,
+    OP_DIVK = 27,
+    OP_IDIVK = 28,
+
+    OP_BANDK = 29,
+    OP_BORK = 30,
+    OP_BXORK = 31,
+
+    OP_SHRI = 32,
+    OP_SHLI = 33,
+
+    OP_ADD = 34,
+    OP_SUB = 35,
+    OP_MUL = 36,
+    OP_MOD = 37,
+    OP_POW = 38,
+    OP_DIV = 39,
+    OP_IDIV = 40,
+
+    OP_BAND = 41,
+    OP_BOR = 42,
+    OP_BXOR = 43,
+    OP_SHL = 44,
+    OP_SHR = 45,
+
+    OP_MMBIN = 46,
+    OP_MMBINI = 47,
+    OP_MMBINK = 48,
+
+    OP_UNM = 49,
+    OP_BNOT = 50,
+    OP_NOT = 51,
+    OP_LEN = 52,
+
+    OP_CONCAT = 53,
+
+    OP_CLOSE = 54,
+    OP_TBC = 55,
+    OP_JMP = 56,
+    OP_EQ = 57,
+    OP_LT = 58,
+    OP_LE = 59,
+
+    OP_EQK = 60,
+    OP_EQI = 61,
+    OP_LTI = 62,
+    OP_LEI = 63,
+    OP_GTI = 64,
+    OP_GEI = 65,
+
+    OP_TEST = 66,
+    OP_TESTSET = 67,
+
+    OP_CALL = 68,
+    OP_TAILCALL = 69,
+
+    OP_RETURN = 70,
+    OP_RETURN0 = 71,
+    OP_RETURN1 = 72,
+
+    OP_FORLOOP = 73,
+    OP_FORPREP = 74,
+
+    OP_TFORPREP = 75,
+    OP_TFORCALL = 76,
+    OP_TFORLOOP = 77,
+
+    OP_SETLIST = 78,
+
+    OP_CLOSURE = 79,
+
+    OP_VARARG = 80,
+
+    OP_VARARGPREP = 81,
+
+    OP_EXTRAARG = 82,
+}
+
 pub struct LuaExecState {}
 
 #[derive(PartialEq, Debug, Clone)]
@@ -53,7 +179,7 @@ pub enum LuaByteCode {
     /// A B: R[A] := R[B]
     Move(u8, u8),
     /// A sBx: R[A] := sBx
-    LoadI(u8, u32),
+    LoadI(u8, i32),
     /// A B C: R[A] := R[B] + R[C]
     Add(u8, u8, u8),
     /// A sB k: if ((R[A] >= sB) ~= k) then pc++
@@ -94,6 +220,28 @@ impl LuaByteCode {
             LuaByteCode::Eq(..) => LuaOpCode::OP_EQ,
         }
     }
+
+    pub fn encode(&self) -> u32 {
+        let payload = match *self {
+            // ABC
+            LuaByteCode::Call(a, b, c)
+            | LuaByteCode::GetTabUp(a, b, c)
+            | LuaByteCode::SetTabUp(a, b, c)
+            | LuaByteCode::Gei(a, b, c)
+            | LuaByteCode::Gti(a, b, c)
+            | LuaByteCode::Eq(a, b, c)
+            | LuaByteCode::Add(a, b, c) => (c as u32) << 24 | (b as u32) << 15 | (a as u32) << 7,
+            // ABx
+            LuaByteCode::LoadK(a, bx) => bx << 8 | a as u32,
+            // AsBx
+            LuaByteCode::LoadI(a, sbx) => excess_sBx!(sbx) << 8 | a as u32,
+            // A B
+            LuaByteCode::Move(a, b) => todo!(),
+        };
+
+        let op = self.opcode() as u32;
+        payload << 7 | op
+    }
 }
 
 #[derive(Debug)]
@@ -103,9 +251,13 @@ pub struct LuaCompiledCode {
 }
 
 impl LuaCompiledCode {
+    pub fn byte_codes(&self) -> &[LuaByteCode] {
+        &self.byte_codes
+    }
+
     pub fn disassembly_code(&self, code: &LuaByteCode) -> String {
         let mut s = String::with_capacity(16);
-        write!(s, "{:<10} ", code.mnemonic()).unwrap();
+        write!(s, "{:<08X} {:<10} ", code.encode(), code.mnemonic()).unwrap();
 
         match code {
             // ABC
@@ -168,3 +320,17 @@ impl Display for LuaCompiledCode {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::LuaByteCode;
+
+    #[test]
+    fn test_encoding() {
+        let code = LuaByteCode::LoadI(0, 2);
+        assert_eq!(code.encode(), 0x80008001);
+
+        let code = LuaByteCode::LoadI(3, -65535);
+        assert_eq!(code.encode(), 0x00000181);
+    }
+}

lib/src/backend/lua/dump.rs

@@ -77,6 +77,12 @@ fn lua_dump_function(
     lua_dump_byte(w, max_stack_size(p))?;
 
     // Dump Code
+    let f = f.read();
+    let code = f.compiled_code().as_ref().unwrap();
+    let lua_code = code.as_any().downcast_ref::<LuaCompiledCode>().unwrap();
+    for c in lua_code.byte_codes() {
+        println!("{:?} {}", c, c.encode())
+    }
 
     // Dump Constants
 
@@ -164,6 +170,10 @@ impl CompiledCode for LuaCompiledCode {
             "Can't get code from single Lua function",
         ))
     }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
 }
 
 #[cfg(test)]

lib/src/backend/lua/mod.rs

@@ -3,8 +3,6 @@ mod bytecode;
 use bytecode::*;
 
 pub mod dump;
-/// 32-bits Lua instruction bytecode encoding/decoding
-mod encoding;
 mod register;
 mod utils;
 mod vm;
@@ -214,7 +212,7 @@ impl AstVisitorMut for LuaBackend {
                 .top_attribute()
                 .register
                 .unwrap_or_else(|| self.reg_mgr.alloc());
-            self.byte_codes.push(LuaByteCode::LoadI(r as u8, v));
+            self.byte_codes.push(LuaByteCode::LoadI(r as u8, v as i32));
             self.top_attribute().register = Some(r);
             return;
         }