Add the break.c example

[jwcc.git] / lib / backend / blargcpu.lua

local outf = io.stderr
local maxstack = 0
local minstack = 0
local labelnum = 0

local regstate = {
	r0 = {},
	r1 = {},
	r2 = {},
	r3 = {}
}

local regvars = {}

function regalloc_reset()
	regvars = {}
	for k,v in pairs(regstate) do
		regstate[k] = {}
	end
end

function regalloc_tmp()
	for k,v in pairs(regstate) do		-- first look for an empty one
		if not v.contains then
			regstate[k] = { contains = " tmp", name = k, locked = 0 }
			table.insert(regvars, regstate[k])
			return regstate[k]
		end
	end
	for k,v in pairs(regstate) do		-- now look for an unlocked variable to evict
		if regalloc_unlocked(v) and v.contains ~= " tmp" then
			regalloc_evict(v)
			regstate[k] = { contains = " tmp", name = k, locked = 0 }
			table.insert(regvars, regstate[k])
			return regstate[k]
		end
	end
	for k,v in pairs(regstate) do		-- now look for an unlocked temporary to evict
		if regalloc_unlocked(v) then
			regalloc_evict(v)
			regstate[k] = { contains = " tmp", name = k, locked = 0 }
			table.insert(regvars, regstate[k])
			return regstate[k]
		end
	end
	error("register allocator: tried to allocate a temporary, but couldn't find anything to evict")
end

function regalloc_get(var)
	for k,v in pairs(regvars) do
		if v.contains == var then
			return v
		end
	end
	local reg = { contains = var, locked = 0, evicted = true, stackpos = var.stackpos }
	table.insert(regvars, reg)
	return reg
end

function regalloc_lock(reg)
	if reg.evicted then
		regalloc_unevict(reg)
	end
	reg.locked = reg.locked + 1
end

function regalloc_unlock(reg)
	assert(reg.locked > 0)
	reg.locked = reg.locked - 1
end

function regalloc_unlocked(reg)
	return not reg.locked or reg.locked == 0
end

function regalloc_free(reg)
	if reg.contains == " tmp" then
		for k,v in pairs(regvars) do
			if v == reg then
				table.remove(regvars, k)
			end
		end
		regstate[reg.name] = {}
		reg.name = nil
		reg.contains = nil
	else
		reg.locked = 0
	end
end

function regalloc_savestate(save)
	-- store all NON-evicted variables that MIGHT get evicted in the savestate
	for k,v in pairs(regstate) do
		if v.contains and regalloc_unlocked(v) then
			save[v] = { var = v, reg = v.name, dirty = v.dirty }
		end
	end
end

function regalloc_restorestate(save)
	-- restore any variables that got evicted
	for k,v in pairs(regstate) do
		if v.dirty and ((not save[v]) or (not save[v].dirty)) then
			regalloc_clean(v)
		end
	end
	for k,v in pairs(save) do
		regalloc_unevict(v.var, v.reg)
	end
end

function regalloc_printstate()
	for k,reg in pairs(regstate) do
		if not reg.contains then
			print(k..": EMPTY")
			for k1,v1 in pairs(reg) do
				print(k.."."..tostring(k1)..": "..tostring(v1))
			end
		else
			local name
			if reg.contains == " tmp" then
				name = "compiler internal variable"
			elseif reg.contains.name then
				name = reg.contains.name
			else
				name = "unknown variable"
			end
			print(k..": "..name.." ("..
				(reg.evicted and "evicted" or "not evicted") .. ", "..
				((not regalloc_unlocked(reg)) and "locked" or "not locked") .. ", " ..
				(reg.dirty and "dirty" or "clean") .. ")")
		end
	end
end

function regalloc_evict(reg)
	assert(reg.contains, "can't evict nil")
	assert(reg.contains ~= " tmp", "tmp eviction not implemented")
	
	if reg.dirty then
		regalloc_clean(reg)
	end
	reg.evicted = true
	if reg.name then
		regstate[reg.name] = {}
		reg.name = nil
	end
end

function regalloc_clean(reg)
	assert(reg.contains, "can't clean nil")
	assert(not reg.evicted, "register already evicted?")
	assert(regalloc_unlocked(reg), "must be unlocked to clean")
	
	local name = "unknown variable"
	if type(reg.contains) == "string" then
		name = reg.contains
	elseif type(reg.contains) == "table" and reg.contains.name then
		name = reg.contains.name
	end
	name = "`"..name.."'"
	
	reg.dirty = false
	
	-- first off, look for an easy eviction -- is there an empty register?
	for k,v in pairs(regstate) do
		if not v.contains then
			local tmp = regalloc_tmp()
			regalloc_lock(tmp)
			emit_opcode("mov", tmp.name..", #"..(reg.stackpos-maxstack-1), "easy evict "..name.." to stack")
			emit_opcode("add", tmp.name..", sp")
			emit_opcode("sto", "["..tmp.name.."], "..reg.name)
			regalloc_free(tmp)
			return
		end
	end
	-- damnit, we have to do an in-place evict
	
	-- Step 1 -- push rN
	emit_opcode("push", "sp, "..reg.name, "in-place evict "..name)
	-- Step 2 -- Move offset from new sp to desired place into rN
	emit_opcode("mov", reg.name..", #"..(reg.stackpos-maxstack-2))
	-- Step 3 -- Add sp to result to get an absolute pointer
	emit_opcode("add", reg.name..", sp")
	-- Step 4 -- Load sp with original contents of rN
	emit_opcode("pop", "sp, sp")
	-- Step 5 -- Store original contents in address now stored in rN
	emit_opcode("sto", "["..reg.name.."], sp")
	-- Step 6 -- Move offset to get back to correct sp into sp
	emit_opcode("mov", "sp, #"..(maxstack-reg.stackpos +1))
	-- Step 7 -- Add rN to sp; sp is now restored
	emit_opcode("add", "sp, "..reg.name)
	-- In this case, the variable is toast. Sorry.
	regstate[reg.name] = {}
	reg.evicted = true
	reg.name = nil
end

function regalloc_unevict(reg, wantreg)
	assert(reg.contains, "can't unevict a nil pointer")
	local name = "unknown variable"
	if type(reg.contains) == "string" then
		name = reg.contains
	elseif type(reg.contains) == "table" and reg.contains.name then
		name = reg.contains.name
	end
	name = "`"..name.."'"
	
	if not reg.evicted and wantreg and wantreg ~= reg.name then
		assert(regalloc_unlocked(regstate[wantreg]), "tried to unevict with a locked wantreg")
		if regstate[wantreg].contains then
			emit_opcode("push", "sp, "..wantreg, "unevict: swap "..name.." and "..wantreg)
			emit_opcode("mov", wantreg..", "..reg.name)
			emit_opcode("pop", "sp, "..reg.name)
			regstate[reg.name] = regstate[wantreg]
			regstate[reg.name].name = reg.name
			regstate[wantreg] = reg
			regstate[wantreg].name = wantreg
			return		-- wee!
		else
			emit_opcode("mov", wantreg..", "..reg.name, "unevict: move "..name.." to "..wantreg)
			regstate[reg.name] = {}
			regstate[wantreg] = reg
			regstate[wantreg].name = wantreg
			return		-- double wee!
		end
	elseif not reg.evicted then
		return	-- well, if we don't actually have to do anything, ...!
	else
		assert(reg.contains ~= " tmp", "tmp uneviction not implemented")
		assert(reg.evicted)
		-- find a register
		for k,v in pairs(regstate) do	-- look for an empty
			if not v.contains and not wantreg then
				wantreg = k
			end
		end
		for k,v in pairs(regstate) do	-- look for a nontmp that's not dirty
			if regalloc_unlocked(v) and v.contains ~= " tmp" and not v.dirty and not wantreg then
				regalloc_evict(v)
				wantreg = k
			end
		end
		for k,v in pairs(regstate) do	-- look for any nontmp
			if regalloc_unlocked(v) and v.contains ~= " tmp" and not wantreg then
				regalloc_evict(v)
				wantreg = k
			end
		end
		for k,v in pairs(regstate) do	-- look for anything to evict
			if regalloc_unlocked(v) and not wantreg then
				regalloc_evict(v)
				wantreg = k
			end
		end
		assert(wantreg and regalloc_unlocked(regstate[wantreg]), "could not find an unlocked register to evict to")
		emit_opcode("mov", wantreg..", #"..(reg.stackpos-maxstack-1), "unevict: pull "..name.." from stack to "..wantreg)
		emit_opcode("add", wantreg..", sp")
		emit_opcode("ldr", wantreg..", ["..wantreg.."]")
		reg.name = wantreg
		reg.evicted = false
		regstate[wantreg] = reg
		return
	end
	error("something got fucked up")
end

function regalloc_unevict_all_tmp()
	for k,v in pairs(regvars) do
		if v.evicted and v.contains == " tmp" then
			regalloc_unevict(v)
		end
	end
end

function regalloc_destroy(variable)
	for k,v in pairs(regvars) do
		if v.contains == variable and not v.evicted then
			regstate[v.name] = {}
			v.name = nil
		end
	end
end

local cursection = nil
function emit(text)
	outf:write(text.."\n")
end

function emit_section(section)
	if cursection == section then
		return
	end
	emit("\t.section "..section)
	cursection = section
end

function emit_label(label)
	emit(label..":")
end

function emit_opcode(opcode, params, comment)
	if comment then
		emit("\t"..opcode.."\t"..params.."\t\t\t@ "..comment)
	else
		emit("\t"..opcode.."\t"..params)
	end
end

local globalvariables = {}
function do_globalvariable(var)
	globalvariables[var.name] = { type = var.vtype }
	if var.vtype == "int" then
		local initializer = var.initializer or 0
		emit_section(".bss")
		emit_label(var.name)
		emit_opcode(".word",tostring(initializer))
	else
		error("unknown variable emit in do_globalvariable")
	end
end

function do_function(func)
	local stack = {}
	
		
	-- First, read in the arguments and register them as being on the stack.
	for k,v in pairs(func.args) do
		stack[-k] = { name = v.name, type = v.type, stackpos = -k }
		assert(v.type == "int", "can only handle ints on the stack right now")
		minstack = -k
	end
	stack[0] = {name = " return address", type = "int" }	-- XXX not actually an int
	
	regalloc_reset()
	emit_section(".text")
	emit("")
	emit_opcode(".global", func.name)
	emit_label(func.name)
	local emitter = function(expr, emitter)
		if expr.type == "stmtlist" then
			local origstack = maxstack
			local initialized = false
			local savestate = {}
			regalloc_savestate(savestate)
			for k,v in pairs(expr.body) do
				if v.type == "variable" then
					assert(not initialized, "cannot have more variables after stack has been emitted")
					maxstack = maxstack + 1
					stack[maxstack] = { name = v.name, type = v.vtype, initializer = v.initializer, stackpos = maxstack }
				else
					if not initialized and origstack ~= maxstack then
						local reg = regalloc_tmp()
						regalloc_lock(reg)
						emit_opcode("mov", reg.name..", #"..tostring(maxstack-origstack), "increase stack size for variables")
						emit_opcode("add", "sp, "..reg.name)
						regalloc_free(reg)
						for k2,v2 in pairs(stack) do
							if v2.initializer then
								local data = regalloc_tmp()	-- one for the data
								local addr = regalloc_tmp()	-- one for the stack address
								regalloc_lock(data)
								emit_opcode("mov", data.name..", #"..v2.initializer, "initialize variable "..v2.name)
								regalloc_lock(addr)
								emit_opcode("mov", addr.name..", #"..tostring(k2-maxstack-1))
								emit_opcode("add", addr.name..", sp")
								emit_opcode("sto", "["..addr.name.."], "..data.name)
								regalloc_free(data)
								regalloc_free(addr)
								v2.initializer = nil
							end
						end
						initialized = true
					end
					local reg = emitter(v, emitter)
					if reg then
						regalloc_free(reg)
					end
				end
			end
			if origstack ~= maxstack then
				local reg = regalloc_tmp()
				emit_opcode("mov", reg.name..", #"..tostring(origstack-maxstack), "decrease stack size for variables")
				emit_opcode("add", "sp, "..reg.name)
				regalloc_free(reg)
				maxstack = origstack
				for i=origstack+1,maxstack do
					regalloc_destroy(stack[i])
					stack[i] = nil
				end
			end
			regalloc_restorestate(savestate)
		elseif expr.type == "expression" then
			return emitter(expr.value, emitter)
		elseif expr.type == "funccall" then
			-- push args on the stack from last to first
			for arg=table.getn(expr.args),1,-1 do
				local reg = emitter(expr.args[arg], emitter)
				assert(reg, "argument "..arg.." did not return a register")
				regalloc_lock(reg)
				regalloc_unevict_all_tmp()
				emit_opcode("push","sp, "..reg.name, "push argument "..arg.." onto the stack for function call "..expr.name)
				regalloc_free(reg)
				maxstack = maxstack + 1
			end
			-- meh, I hate fucking around with the allocator's internal state here, XXX
			for k,v in pairs(regstate) do
				assert(not(v.contains == " tmp" and not v.evicted), "it's too damn late now to evict tmp")
				if not v.evicted and v.contains then
					regalloc_evict(v)
				end
			end
			local reg = regalloc_tmp()
			regalloc_lock(reg)
			emit_opcode("mov",reg.name..", #"..expr.name, "function call "..expr.name)
			emit_opcode("call","sp, "..reg.name)
			-- first thing: save the return value!
			local retv = regalloc_tmp()
			regalloc_lock(retv)
			if reg.name == "r0" then	-- try hard to get r0 for the return value
				local tmp = reg
				reg = retv
				retv = tmp
			end
			if retv.name ~= "r0" then	-- wee o_O
				emit_opcode("mov", retv.name..", r0", "grab return value")
			end
			-- decrement the stack by the number of args
			emit_opcode("mov",reg.name..", #-"..table.getn(expr.args), "clean up stack from function call")
			emit_opcode("add","sp, "..reg.name)
			maxstack = maxstack - table.getn(expr.args)
			regalloc_free(reg)
			return retv
		elseif expr.type == "while" then
			local headlabel, taillabel = labelnum, labelnum + 1
			local savestate = {}
			labelnum = labelnum + 2
			emit_label(".L"..headlabel)
			-- do expression
			local resreg = emitter(expr.expression, emitter)
			assert(resreg, "while expression did not return a register")
			regalloc_lock(resreg)
			local tmp = regalloc_tmp()
			emit_opcode("mov",tmp.name..", #0")
			emit_opcode("tst",tmp.name..", "..resreg.name)
			emit_opcode("moveq","pc, #.L"..taillabel, "conditional break from while")
			regalloc_free(resreg)
			regalloc_free(tmp)
			regalloc_savestate(savestate)
			local tmp = emitter(expr.chunk, emitter)
			if tmp then
				regalloc_free(tmp)
			end
			regalloc_restorestate(savestate)
			emit_opcode("mov","pc, #.L"..headlabel)
			emit_label(".L"..taillabel)
		elseif expr.type == "return" then
			local reg = emitter(expr.expression, emitter)
			assert(reg, "return expression did not return a register")
			regalloc_unevict(reg, "r0")
			regalloc_lock(reg)
			assert(reg.name == "r0", "failed to unevict to r0?")
			local tmp = regalloc_tmp()
			regalloc_lock(tmp)
			emit_opcode("mov", tmp.name..", #-"..maxstack, "clean up stack before return")
			emit_opcode("add","sp, "..tmp.name)
			emit_opcode("pop","sp, pc", "return")
			regalloc_free(tmp)
			regalloc_free(reg)
		elseif expr.type == "if" then
			local taillabel = labelnum
			local savestate = {}
			labelnum = labelnum + 1
			local resreg = emitter(expr.expression, emitter)
			assert(resreg, "if expression did not return a register")
			regalloc_lock(resreg)
			local tmp = regalloc_tmp()
			emit_opcode("mov",tmp.name..", #0")
			emit_opcode("tst",tmp.name..", "..resreg.name)
			emit_opcode("moveq","pc, #.L"..taillabel, "conditional exit from if")
			regalloc_free(resreg)
			regalloc_free(tmp)
			regalloc_savestate(savestate)
			local tmp = emitter(expr.chunk, emitter)
			if tmp then
				regalloc_free(tmp)
			end
			regalloc_restorestate(savestate)
			emit_label(".L"..taillabel)
		elseif expr.type == "operator" and expr.value ~= "=" then
			-- an arith oper
			local lhs = emitter(expr.left, emitter)
			local rhs = emitter(expr.right, emitter)
			local tmp
			assert(lhs, "lhs did not return a register")
			assert(rhs, "rhs did not return a register")
			regalloc_lock(lhs)
			regalloc_lock(rhs)
			if lhs.contains == " tmp" then
				tmp = lhs
			elseif rhs.contains == " tmp" then
				tmp = rhs
			end
			
			if expr.value == "==" or expr.value == "<" or expr.value == ">" then
				local predicates = { ["=="] = "eq", ["<"] = "lt", [">"] = "gt" }
				if not tmp then
					tmp = regalloc_tmp()
					regalloc_lock(tmp)
				end
				-- tst is fucking backwards
				emit_opcode("tst",rhs.name..", "..lhs.name, "comparison operator "..expr.value)
				emit_opcode("mov",tmp.name..", #0")
				emit_opcode("mov"..predicates[expr.value], tmp.name..", #1")
			elseif expr.value == "+" then
				if not tmp then
					tmp = regalloc_tmp()
					regalloc_lock(tmp)
				end
				if tmp == lhs then
					emit_opcode("add",tmp.name..", "..rhs.name)
				elseif tmp == rhs then
					emit_opcode("add",tmp.name..", "..lhs.name)
				else
					emit_opcode("mov",tmp.name..", "..lhs.name)
					emit_opcode("add",tmp.name..", "..rhs.name)
				end
			elseif expr.value == "-" then
				if not tmp then
					tmp = regalloc_tmp()
					regalloc_lock(tmp)
				end
				emit_opcode("not", rhs.name)
				local tmp2 = regalloc_tmp()
				regalloc_lock(tmp2)
				emit_opcode("mov",tmp2.name..", #1")
				emit_opcode("add",rhs.name..", "..tmp2.name)
				regalloc_free(tmp2)
				if tmp == lhs then
					emit_opcode("add",tmp.name..", "..rhs.name)
				elseif tmp == rhs then
					emit_opcode("add",tmp.name..", "..lhs.name)
				else
					emit_opcode("mov",tmp.name..", "..lhs.name)
					emit_opcode("add",tmp.name..", "..rhs.name)
				end
			else
				print("Dunno how to handle operator "..expr.value)
				return nil
			end
			
			if lhs ~= tmp then
				regalloc_free(lhs)
			end
			if rhs ~= tmp then
				regalloc_free(rhs)
			end
			regalloc_unlock(tmp)
			return tmp
		elseif expr.type == "operator" and expr.value == "=" then
			assert(expr.left.type == "identifier", "lhs has to be an identifier")
			local lhs = emitter(expr.left, emitter)
			local rhs = emitter(expr.right, emitter)
			assert(rhs, "rhs did not return a register")
			regalloc_lock(lhs)
			regalloc_lock(rhs)
			lhs.dirty = true
			emit_opcode("mov",lhs.name..", "..rhs.name)
			regalloc_free(rhs)
			return lhs
		elseif expr.type == "identifier" then
			-- find the variable on the stack
			local var
			for k,v in pairs(stack) do
				if v.name == expr.value and not var then
					var = v
				end
			end
			for k,v in pairs(globalvariables) do
				if v.name == expr.value and not var then
					var = v
				end
			end
			assert(var, "I have no idea what identifier "..expr.value.." is")
			return regalloc_get(var)
		elseif expr.type == "number" then
			local reg = regalloc_tmp()
			regalloc_lock(reg)
			emit_opcode("mov", reg.name..", #"..expr.value)
			regalloc_unlock(reg)
			return reg
		else
			error("unknown expression type "..expr.type)
		end
	end
	emitter(func.body, emitter)	-- call into it, and light it off
	emit_opcode("pop", "sp, pc")
end

function backend(parsetree, outfn)
	if outfn then
		outf = io.open(outfn, "w")
	else
		outf = io.stdout
	end
	emit("@ blargCPU output generated by jwcc")
	for k,v in pairs(parsetree) do
		if v.type == "variable" then
			do_globalvariable(v)
		elseif v.type == "function" then
			do_function(v)
		else
			error("invalid type in global deparse")
		end
	end
	if outfn then
		outf:close()
	end
	outf = io.stderr
end