(BROKEN!) A simple JITed emulator for my MCPU instruction set(does not work yet! Will update soon)

emulator_jit.t.lua

local debug_print = true

-- index into the state array("optimized" for instruction decoding)
local STATE_PC = 0
local STATE_ADDR = 1
local STATE_ALU_A = 2
local STATE_IMM = 3
local STATE_ALU_B = 4
local STATE_I = 5
local STATE_J = 6
local STATE_K = 7

-- perform the B pre-ALU step to get the effective value use for B in ALU test/data operations
terra alu_b_op(state : &uint32) : uint32
	-- decode ALU instruction
	var alu_op = state[STATE_IMM]
	var b_reg = state[STATE_ALU_B]
	var b_op = alu_op and 0x60
	var inv = alu_op and 0x08
	var b_res = b_reg
	if b_op == 0x20 then
		b_res = state[STATE_IMM] >> 7
	elseif b_op == 0x40 then
		b_res = b_res >> 1
	elseif b_op == 0x60 then
		b_res = b_res << 1
	end
	if inv ~= 0 then
		return b_res ^ 0xffffffff
	else
		return b_res
	end
end

-- perform a single ALU data operation
terra alu_data_op(state : &uint32) : uint32
	-- get A value(from register) and B value(from B pre-operation)
	var a_val : uint32 = state[STATE_ALU_A]
	var b_val : uint32 = alu_b_op(state)

	-- decode ALU instruction
	var alu_op = state[STATE_IMM]
	var cin = alu_op and 0x10
	var d_op = alu_op and 0x7

	-- return correct value
	if d_op == 0 then -- ADD
		if cin == 0 then
			return a_val + b_val
		else
			return a_val + b_val + 1
		end
	elseif d_op == 1 then -- AND
		return a_val and b_val
	elseif d_op == 2 then -- OR
		return a_val or b_val
	elseif d_op == 3 then -- XOR
		return a_val ^ b_val
	elseif d_op == 4 then -- A
		return a_val
	elseif d_op == 5 then -- B
		return b_val
	elseif d_op == 6 then -- X
		return 0
	elseif d_op == 7 then -- Y
		return 0
	end
end

terra alu_test_op(state : &uint32) : bool
	-- get A value(from register) and B value(from B pre-operation)
	var a_val : uint32 = state[STATE_ALU_A]
	var b_val : uint32 = alu_b_op(state)

	-- decode ALU instruction
	var alu_op = state[STATE_IMM]
	var t_op = alu_op and 0x7

	if t_op == 0 then -- A_EQ_Z
		return a_val == 0
	elseif t_op == 1 then -- B_EQ_Z
		return b_val == 0
	elseif t_op == 2 then -- A_GT_B
		return a_val > b_val
	elseif t_op == 3 then -- A_EQ_B
		return a_val == b_val
	elseif t_op == 4 then -- A_LT_B
		return a_val < b_val
	elseif t_op == 5 then -- B[bit 0]
		return (b_val and 1) == 1
	elseif t_op == 6 then -- B[bit 31]
		return (b_val and 0x80000000) == 0x80000000
	elseif t_op == 7 then -- SENSE
		return false
	end
end

-- assemble instructions into a list of terra statements
-- until the first instruction that could modify the program counter
local function assemble_basic_block(pc_start, irom, dram, state)
	-- list of terra statements that make up the body of this basic block
	local block_list = terralib.newlist()

	-- state for assembling multiple instructions into a single basic block
	local last_imm = false
	local imm_val = 0
	local halts = false
	local pc_offset = 0

	-- add instructions to block_list until encountering a HALT or MOV/CMOV to program counter(end of basic block)
	while true do
		local instr = irom:byte(pc_start+pc_offset+1) or 0
		pc_offset = pc_offset + 1
		if instr >= 128 then
			-- IMM instruction(code to load IMM register is emitted on the next non-IMM instruction)
			if last_imm then
				imm_val = imm_val * 128 + (instr - 128)
			else
				imm_val = (instr - 128)
				last_imm = true
			end
			block_list:insert(quote state[STATE_PC] = state[STATE_PC] + 1 end)
		else
			-- MOV/CMOV instruction
			if last_imm then -- emit completed IMM instruction
				block_list:insert(quote state[STATE_IMM] = imm_val end)
				last_imm = false
			end

			-- decode instruction
			local is_cond = false
			if instr >= 64 then
				is_cond = true
				instr = instr - 64
			end
			local op_src = bit.band(instr, 0x7)
			local op_dst = bit.rshift(bit.band(instr, 0x38), 3)

			-- default statement for directly reading from state(registers on the bus)
			local src_q = `state[op_src]
			-- default statement for directly writing to state(registers on the bus)
			local dst_q = `state[op_dst]

			-- special statements for reading(from functions on the bus)
			if op_src == 2 then -- RAM read
				src_q = `dram[state[1]] 
			elseif op_src == 4 then -- ALU result
				src_q = `alu_data_op(state)
			end

			-- special statements for writing(to functions on the bus)
			if op_dst == 2 then -- RAM write
				dst_q = `dram[state[1]]
			elseif op_dst == 3 then -- ALU A(re-order in state)
				dst_q = `state[STATE_ALU_A]
			elseif op_dst == 4 then -- ALU B(re-order in state)
				dst_q = `state[STATE_ALU_B]
			end

			-- emit code for move instruction
			if is_cond then
				-- conditional move instruction
				block_list:insert(quote if alu_test_op(state) then dst_q = src_q end end)
			else
				-- regular move instruction
				block_list:insert(quote dst_q = src_q end)
			end

			-- emit code to handle the program counter
			if (op_src == 0) and (op_dst == 0) then
				-- HALT instruction, terminate block and halt
				halts = true
				break
			elseif op_dst == 0 then
				-- JUMP/BRANCH instruction (destination is program counter)
				-- terminate current block
				break
			else
				-- regular instruction, increment PC afterwards
				block_list:insert(quote state[STATE_PC] = state[STATE_PC] + 1 end)
			end
		end
	end

	-- return the completed block
	local block_func = terra() [ block_list ] end
	return {
		pc_start = pc_start, -- start address this compiled block represents
		instructions = pc_offset, -- number of instructions in the basic block
		func = block_func, -- reference to (compiled) function
		halts = halts, -- if this block will terminate execution
		block_list = block_list, -- list of instruction statements in the block
	}
end

-- execute a basic block
local function execute_basic_block(irom, code_blocks, dram, state)
	-- check if a compiled version of the basic block already exists
	--local pc_start = state[STATE_PC]
	local pc_start = state:get()[STATE_PC]
	local block = code_blocks[pc_start]
	if not block then
		-- need to generate the basic block
		block = assemble_basic_block(pc_start, irom, dram, state)
		if debug_print then
			print(("Compiled block at: %.8x"):format(pc_start), block.func)
			block.func:disas()
		end
		code_blocks[block.pc_start] = block
	end
	if debug_print then
		print(("Executing block at: %.8x"):format(pc_start))
		print("block.func",block.func)
	end
	-- execute the basic block
	block.func()
	-- terminate execution if needed(last instruction was HALT)
	if block.halts then return; end
	-- continue executing the next basic block
	return execute_basic_block(irom, code_blocks, dram, state)
end

-- read IROM image from file
local irom_str = assert(io.open(arg[1] or "output.bin", "rb")):read("*a")

-- mapping of irom address to compiled code block("code cache")
local code_blocks = {}

-- create dram/state global variables
local mem_size = tonumber(arg[2]) or 0x10000
local dram = global(uint32[mem_size])
local state = global(uint32[8])

-- start compiling and executing basic blocks
execute_basic_block(irom_str, code_blocks, dram, state)

-- print final state
if debug_print then
	print("Halted")
	local state_f = state:get()
	print(("PC     0x%.8x"):format(state_f[STATE_PC]))
	print(("ADDR   0x%.8x"):format(state_f[STATE_ADDR]))
	print(("ALU_A  0x%.8x"):format(state_f[STATE_ALU_A]))
	print(("IMM    0x%.8x"):format(state_f[STATE_IMM]))
	print(("ALU_B  0x%.8x"):format(state_f[STATE_ALU_B]))
	print(("I      0x%.8x"):format(state_f[STATE_I]))
	print(("J      0x%.8x"):format(state_f[STATE_J]))
	print(("K      0x%.8x"):format(state_f[STATE_K]))
end