luapy (5) lua vm

1. 修改lua state

from lua_stack import LuaStack
from lua_type import LuaType
from lua_value import LuaValue
from arith_op import ArithOp
from arithmetic import Arithmetic
from cmp_op import CmpOp
from compare import Compare


class LuaState:
    def __init__(self, proto):
        self.stack = LuaStack()
        self.proto = proto
        self.pc = 0

    # ...

    def get_pc(self):
        return self.pc

    def add_pc(self, n):
        self.pc += n

    def fetch(self):
        code = self.proto.get_code()[self.pc]
        self.pc += 1
        return code

    def get_const(self, idx):
        self.stack.push(self.proto.get_constants()[idx])

    def get_rk(self, rk):
        if rk > 0xff:   # constant
            self.get_const(rk & 0xff)
        else:           # register
            self.push_value(rk + 1)

2. Instruction

from arith_op import ArithOp
from cmp_op import CmpOp
from lua_type import LuaType


# OpMode
IABC = 0
IABx = 1
IAsBx = 2
IAx = 3

# OpArg
OpArgN = 0
OpArgU = 1
OpArgR = 2
OpArgK = 3


class OpCode:
    MOVE = 0
    LOADK = 1
    LOADKX = 2
    LOADBOOL = 3
    LOADNIL = 4
    GETUPVAL = 5
    GETTABUP = 6
    GETTABLE = 7
    SETTABUP = 8
    SETUPVAL = 9
    SETTABLE = 10
    NEWTABLE = 11
    SELF = 12
    ADD = 13
    SUB = 14
    MUL = 15
    MOD = 16
    POW = 17
    DIV = 18
    IDIV = 19
    BAND = 20
    BOR = 21
    BXOR = 22
    SHL = 23
    SHR = 24
    UNM = 25
    BNOT = 26
    NOT = 27
    LEN = 28
    CONCAT = 29
    JMP = 30
    EQ = 31
    LT = 32
    LE = 33
    TEST = 34
    TESTSET = 35
    CALL = 36
    TAILCALL = 37
    RETURN = 38
    FORLOOP = 39
    FORPREP = 40
    TFORCALL = 41
    TFORLOOP = 42
    SETLIST = 43
    CLOSURE = 44
    VARARG = 45
    EXTRAARG = 46

    def __init__(self, test_flag, set_a_flag, arg_b_mode, arg_c_mode, op_mode, name, action):
        self.test_flag = test_flag
        self.set_a_flag = set_a_flag
        self.arg_b_mode = arg_b_mode
        self.arg_c_mode = arg_c_mode
        self.op_mode = op_mode
        self.name = name
        self.action = action


# R(A) := R(B)
def move(inst, vm):
    a, b, _ = inst.a_b_c()
    a += 1
    b += 1
    vm.copy(b, a)


# R(A) := Kst(Bx)
def loadk(inst, vm):
    a, bx = inst.a_bx()
    a += 1
    vm.get_const(bx)
    vm.replace(a)


# R(A) := Kst(extra arg)
def loadkx(inst, vm):
    a, _ = inst.a_bx()
    a += 1
    ax = Instruction(vm.fetch()).ax()
    vm.get_const(ax)
    vm.replace(a)


# R(A) := (bool)B; if (C) pc++
def loadbool(inst, vm):
    a, b, c = inst.a_b_c()
    vm.push_boolean(b != 0)
    vm.replace(a+1)

    if c != 0:
        vm.add_pc(1)


# R(A), R(A+1) ... R(A+B) := nil
def loadnil(inst, vm):
    a, b, _ = inst.a_b_c()
    a += 1

    vm.push_nil()
    for i in range(a, a+b):
        vm.copy(-1, i)
    vm.pop(1)


# arith
def arith_binary(inst, vm, op):
    a, b, c = inst.a_b_c()
    a += 1

    vm.get_rk(b)
    vm.get_rk(c)
    vm.arith(op)
    vm.replace(a)


def arith_unary(inst, vm, op):
    a, b, _ = inst.a_b_c()
    a += 1
    b += 1

    vm.push_value(b)
    vm.arith(op)
    vm.replace(a)


def add(inst, vm):
    arith_binary(inst, vm, ArithOp.ADD)


def sub(inst, vm):
    arith_binary(inst, vm, ArithOp.SUB)


def mul(inst, vm):
    arith_binary(inst, vm, ArithOp.MUL)


def mod(inst, vm):
    arith_binary(inst, vm, ArithOp.MOD)


def luapow(inst, vm):
    arith_binary(inst, vm, ArithOp.POW)


def div(inst, vm):
    arith_binary(inst, vm, ArithOp.DIV)


def idiv(inst, vm):
    arith_binary(inst, vm, ArithOp.IDIV)


def band(inst, vm):
    arith_binary(inst, vm, ArithOp.BAND)


def bor(inst, vm):
    arith_binary(inst, vm, ArithOp.BOR)


def bxor(inst, vm):
    arith_binary(inst, vm, ArithOp.BXOR)


def shl(inst, vm):
    arith_binary(inst, vm, ArithOp.SHL)


def shr(inst, vm):
    arith_binary(inst, vm, ArithOp.SHR)


def unm(inst, vm):
    arith_unary(inst, vm, ArithOp.UNM)


def bnot(inst, vm):
    arith_unary(inst, vm, ArithOp.BNOT)


# R(A) := length of R(B)
def length(inst, vm):
    a, b, _ = inst.a_b_c()
    a += 1
    b += 1

    vm.len(b)
    vm.replace(a)


def concat(inst, vm):
    a, b, c = inst.a_b_c()
    a += 1
    b += 1
    c += 1
    n = c - b + 1

    vm.check_stack()
    for i in range(b, c+1):
        vm.push_value(i)

    vm.concat(n)
    vm.repalce(a)


def jmp(inst, vm):
    a, sbx = inst.a_sbx()
    vm.add_pc(sbx)
    assert(a == 0)


def compare(inst, vm, op):
    a, b, c = inst.a_b_c()
    vm.get_rk(b)
    vm.get_rk(c)
    if vm.compare(-2, op, -1) != (a != 0):
        vm.add_pc(1)
    vm.pop(2)


def eq(inst, vm):
    compare(inst, vm, CmpOp.EQ)


def lt(inst, vm):
    compare(inst, vm, CmpOp.LT)


def le(inst, vm):
    compare(inst, vm, CmpOp.LE)


# R(A) := not R(B)
def luanot(inst, vm):
    a, b, _ = inst.a_b_c()
    a += 1
    b += 1

    vm.push_boolean(not vm.to_boolean(b))
    vm.replace(a)


# if not (R(A) <=> C) then pc++
def test(inst, vm):
    a, _, c = inst.a_b_c()
    if vm.to_boolean(a) != (c != 0):
        vm.add_pc(1)


# if (R(B) <=> C) then R(A) := R(B) else pc++
def testset(inst, vm):
    a, b, c = inst.a_b_c()
    a += 1
    b += 1

    if vm.to_boolean(b) != (c != 0):
        vm.copy(b, a)
    else:
        vm.add_pc(1)


# R(A) += R(A+2)
# if R(A) <?= R(A+1) then {
#   pc += sBx;
#   R(A+3) = R(A)
# }
def forloop(inst, vm):
    a, sbx = inst.a_sbx()
    a += 1

    vm.push_value(a+2)
    vm.push_value(a)
    vm.arith(ArithOp.ADD)
    vm.replace(a)

    positive_step = vm.to_number(a+2) >= 0
    if (positive_step and vm.compare(a, CmpOp.LE, a+1)) \
            or ((not positive_step) and vm.compare(a+1, CmpOp.LE, a)):
        vm.add_pc(sbx)
        vm.copy(a, a+3)


# R(A)-=R(A+2); pc+=sBx
def forprep(inst, vm):
    a, sbx = inst.a_sbx()
    a += 1

    if vm.type(a) == LuaType.STRING:
        vm.push_number(vm.to_number(a))
        vm.replace(a)

    if vm.type(a+1) == LuaType.STRING:
        vm.push_number(vm.to_number(a+1))
        vm.replace(a+1)

    if vm.type(a+2) == LuaType.STRING:
        vm.push_number(vm.to_number(a+2))
        vm.replace(a+2)

    vm.push_value(a)
    vm.push_value(a+2)
    vm.arith(ArithOp.SUB)
    vm.replace(a)
    vm.add_pc(sbx)


op_codes = [
    #      T  A  B       C       mode   name
    OpCode(0, 1, OpArgR, OpArgN, IABC,  "MOVE    ", move),      # R(A) := R(B)
    OpCode(0, 1, OpArgK, OpArgN, IABx,  "LOADK   ", loadk),     # R(A) := Kst(Bx)
    OpCode(0, 1, OpArgN, OpArgN, IABx,  "LOADKX  ", loadkx),    # R(A) := Kst(extra arg)
    OpCode(0, 1, OpArgU, OpArgU, IABC,  "LOADBOOL", loadbool),  # R(A) := (bool)B; if (C) pc++
    OpCode(0, 1, OpArgU, OpArgN, IABC,  "LOADNIL ", loadnil),   # R(A), R(A+1), ..., R(A+B) := nil
    OpCode(0, 1, OpArgU, OpArgN, IABC,  "GETUPVAL", None),      # R(A) := UpValue[B]
    OpCode(0, 1, OpArgU, OpArgK, IABC,  "GETTABUP", None),      # R(A) := UpValue[B][RK(C)]
    OpCode(0, 1, OpArgR, OpArgK, IABC,  "GETTABLE", None),      # R(A) := R(B)[RK(C)]
    OpCode(0, 0, OpArgK, OpArgK, IABC,  "SETTABUP", None),      # UpValue[A][RK(B)] := RK(C)
    OpCode(0, 0, OpArgU, OpArgN, IABC,  "SETUPVAL", None),      # UpValue[B] := R(A)
    OpCode(0, 0, OpArgK, OpArgK, IABC,  "SETTABLE", None),      # R(A)[RK(B)] := RK(C)
    OpCode(0, 1, OpArgU, OpArgU, IABC,  "NEWTABLE", None),      # R(A) := {} (size = B,C)
    OpCode(0, 1, OpArgR, OpArgK, IABC,  "SELF    ", None),      # R(A+1) := R(B); R(A) := R(B)[RK(C)]
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "ADD     ", add),       # R(A) := RK(B) + RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "SUB     ", sub),       # R(A) := RK(B) - RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "MUL     ", mul),       # R(A) := RK(B) * RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "MOD     ", mod),       # R(A) := RK(B) % RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "POW     ", luapow),    # R(A) := RK(B) ^ RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "DIV     ", div),       # R(A) := RK(B) / RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "IDIV    ", idiv),      # R(A) := RK(B) // RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "BAND    ", band),      # R(A) := RK(B) & RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "BOR     ", bor),       # R(A) := RK(B) | RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "BXOR    ", bxor),      # R(A) := RK(B) ~ RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "SHL     ", shl),       # R(A) := RK(B) << RK(C)
    OpCode(0, 1, OpArgK, OpArgK, IABC,  "SHR     ", shr),       # R(A) := RK(B) >> RK(C)
    OpCode(0, 1, OpArgR, OpArgN, IABC,  "UNM     ", unm),       # R(A) := -R(B)
    OpCode(0, 1, OpArgR, OpArgN, IABC,  "BNOT    ", bnot),      # R(A) := ~R(B)
    OpCode(0, 1, OpArgR, OpArgN, IABC,  "NOT     ", luanot),    # R(A) := not R(B)
    OpCode(0, 1, OpArgR, OpArgN, IABC,  "LEN     ", length),    # R(A) := length of R(B)
    OpCode(0, 1, OpArgR, OpArgR, IABC,  "CONCAT  ", concat),    # R(A) := R(B).. ... ..R(C)
    OpCode(0, 0, OpArgR, OpArgN, IAsBx, "JMP     ", jmp),       # pc+=sBx; if (A) close all upvalues >= R(A - 1)
    OpCode(1, 0, OpArgK, OpArgK, IABC,  "EQ      ", eq),        # if ((RK(B) == RK(C)) ~= A) then pc++
    OpCode(1, 0, OpArgK, OpArgK, IABC,  "LT      ", lt),        # if ((RK(B) <  RK(C)) ~= A) then pc++
    OpCode(1, 0, OpArgK, OpArgK, IABC,  "LE      ", le),        # if ((RK(B) <= RK(C)) ~= A) then pc++
    OpCode(1, 0, OpArgN, OpArgU, IABC,  "TEST    ", test),      # if not (R(A) <=> C) then pc++
    OpCode(1, 1, OpArgR, OpArgU, IABC,  "TESTSET ", testset),   # if (R(B) <=> C) then R(A) := R(B) else pc++
    OpCode(0, 1, OpArgU, OpArgU, IABC,  "CALL    ", None),      # R(A), ...,R(A+C-2) := R(A)(R(A+1), ...,R(A+B-1))
    OpCode(0, 1, OpArgU, OpArgU, IABC,  "TAILCALL", None),      # return R(A)(R(A+1), ... ,R(A+B-1))
    OpCode(0, 0, OpArgU, OpArgN, IABC,  "RETURN  ", None),      # return R(A), ... ,R(A+B-2)
    OpCode(0, 1, OpArgR, OpArgN, IAsBx, "FORLOOP ", forloop),   # R(A)+=R(A+2); if R(A) <?= R(A+1) then { pc+=sBx; R(A+3)=R(A) }
    OpCode(0, 1, OpArgR, OpArgN, IAsBx, "FORPREP ", forprep),   # R(A)-=R(A+2); pc+=sBx
    OpCode(0, 0, OpArgN, OpArgU, IABC,  "TFORCALL", None),      # R(A+3), ... ,R(A+2+C) := R(A)(R(A+1), R(A+2));
    OpCode(0, 1, OpArgR, OpArgN, IAsBx, "TFORLOOP", None),      # if R(A+1) ~= nil then { R(A)=R(A+1); pc += sBx }
    OpCode(0, 0, OpArgU, OpArgU, IABC,  "SETLIST ", None),      # R(A)[(C-1)*FPF+i] := R(A+i), 1 <= i <= B
    OpCode(0, 1, OpArgU, OpArgN, IABx,  "CLOSURE ", None),      # R(A) := closure(KPROTO[Bx])
    OpCode(0, 1, OpArgU, OpArgN, IABC,  "VARARG  ", None),      # R(A), R(A+1), ..., R(A+B-2) = vararg
    OpCode(0, 0, OpArgU, OpArgU, IAx,   "EXTRAARG", None),      # extra (larger) argument for previous opcode
]


class Instruction:
    MAXARG_Bx = (1 << 18) - 1
    MAXARG_sBx = MAXARG_Bx >> 1

    def __init__(self, code):
        self.code = code

    # ...

    def execute(self, vm):
        op_codes[self.op_code()].action(self, vm)

    # ...

3. test

from binary_chunk import BinaryChunk
from lua_state import LuaState
from opcode import Instruction
from opcode import OpCode


def lua_main(proto):
    vm = LuaState(proto)
    print('max stack size: ', proto.get_max_stack_size())
    vm.set_top(proto.get_max_stack_size())
    while True:
        pc = vm.get_pc()
        i = vm.fetch()
        inst = Instruction(i)
        if inst.op_code() != OpCode.RETURN:
            inst.execute(vm)
            print('[%02d] %-8s ' % (pc+1, inst.op_name()), end='')
            vm.print_stack()
        else:
            break


def main():
    bc = BinaryChunk('./test/sum.luac')
    bc.print_header()
    bc.check_header()
    bc.print_main_func()

    proto = bc.get_main_func()
    lua_main(proto)


if __name__ == '__main__':
    main()

local sum = 0
for i = 1, 100 do
    if i % 2 == 0 then
        sum = sum + i
    end
end

4. result

signature:        b'\x1bLua'
version:          83
format:           0
luac_data:        b'\x19\x93\r\n\x1a\n'
cint_size:        4
csizet_size:      8
inst_size:        4
lua_int_size:     8
lua_number_size:  8
luac_int:         0x5678
luac_num:         370.5

main <@sum.lua:0,0> (11 instructions)
0+ params, 6 slots, 1 upvalues, 5 locals, 4 constants, 0 functions
	1	[1]	LOADK           0       -1
	2	[2]	LOADK           1       -2
	3	[2]	LOADK           2       -3
	4	[2]	LOADK           3       -2
	5	[2]	FORPREP         1        4
	6	[3]	MOD             5        4       -4
	7	[3]	EQ              0        5       -1
	8	[3]	JMP             0        1
	9	[4]	ADD             0        0        4
	10	[2]	FORLOOP         1       -5
	11	[6]	RETURN          0        1
constants (4):
	1	0
	2	1
	3	100
	4	2
locals (5):
	1	sum	2	12
	2	(for index)	5	11
	3	(for limit)	5	11
	4	(for step)	5	11
	5	i	6	10
upvalues (1):
	1	_ENV	1	0
max stack size:  6
[01] LOADK    [0][nil][nil][nil][nil][nil]
[02] LOADK    [0][1][nil][nil][nil][nil]
[03] LOADK    [0][1][100][nil][nil][nil]
[04] LOADK    [0][1][100][1][nil][nil]
[05] FORPREP  [0][0][100][1][nil][nil]
[10] FORLOOP  [0][1][100][1][1][nil]
[06] MOD      [0][1][100][1][1][1]
[07] EQ       [0][1][100][1][1][1]
[08] JMP      [0][1][100][1][1][1]
[10] FORLOOP  [0][2][100][1][2][1]
[06] MOD      [0][2][100][1][2][0]
[07] EQ       [0][2][100][1][2][0]
[09] ADD      [2][2][100][1][2][0]
[10] FORLOOP  [2][3][100][1][3][0]
# ...
[06] MOD      [2450][99][100][1][99][1]
[07] EQ       [2450][99][100][1][99][1]
[08] JMP      [2450][99][100][1][99][1]
[10] FORLOOP  [2450][100][100][1][100][1]
[06] MOD      [2450][100][100][1][100][0]
[07] EQ       [2450][100][100][1][100][0]
[09] ADD      [2550][100][100][1][100][0]
[10] FORLOOP  [2550][101][100][1][100][0]