DeadlyBranches

-- Configurable Stuff starts here
local corrupt = true    -- Actually preform Corruption
local recorrupt = true  -- Unflag accessed areas for potential recorruption.
local range = 50        -- Percentage from 0 to 100 of breaking branches
local maxStack = 128    -- Maximal states on stack
local release = 1       -- Unflag accessed areas every (release)
local clean = false     -- Clean up on resets
local badcodes = true   -- Backtrack on "illegal" opcodes

local debug = false     -- Put debugging messages

-- Comment these out for them not to be checked
local bops = {
	{0x10, 128,   0}, -- BPL
	{0x30, 128, 128}, -- BMI
	{0x50,  64,   0}, -- BVC
	{0x70,  64,  64}, -- BVS
	{0x90,   1,   0}, -- BCC
	{0xB0,   1,   1}, -- BCS
--	{0xD0,   2,   0}, -- BNE
	{0xF0,   2,   2}, -- BEQ
}

local bopcount
if debug then
	bopcount = {0,0,0,0,0,0,0,0}
end

local keys = {
	decrease_range = "leftbracket", -- '['
	increase_range = "rightbracket", -- ']'
	flush_map = 'O', -- 'O'
	restore_good = 'U', -- 'G'
}
-- Configurable Stuff ends here

--[[
Note: possible keys are:

leftclick, rightclick, middleclick, capslock, numlock, scrolllock, 0, 1, 2, 3,
4, 5, 6, 7, 8, 9, A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R, S, T,
U, V, W, X, Y, Z, F1, F2, F3, F4, F5, F6,  F7, F8, F9, F10, F11, F12, F13, F14,
F15, F16, F17, F18, F19, F20, F21, F22, F23, F24, backspace, tab, enter, shift,
control, alt, pause, escape, space, pageup, pagedown, end, home, left, up,
right, down, numpad0, numpad1, numpad2, numpad3, numpad4, numpad5, numpad6,
numpad7, numpad8, numpad9, numpad*, insert, delete, numpad+, numpad-, numpad.,
numpad/, semicolon, plus, minus, comma, period, slash, backslash, tilde, quote,
leftbracket, rightbracket.
--]]

math.randomseed(os.time())

local accessMap = {}
local stack = {}
local stackg = {}
local oldinput = {}
local NewState = false
local GoodState = false
local badTrigger = false
local frameStart = true

local message
local messagetime = -math.huge
local entries = 0
local didkill = 1

local _badopcodes = {0x0B,0x2B,0x87,0x97,0x83,0x8F,0x6B,0x4B,0xAB,0x9F,0x93,0xCB,0xC7,0xD7,0xCF,0xDF,0xDB,0xC3,0xD3,0x04,0x14,0x34,0x44,0x54,0x64,0x74,0x80,0x82,0x89,0xC2,0xD4,0xE2,0xF4,0xE7,0xF7,0xEF,0xFF,0xFB,0xE3,0xF3,0x02,0x12,0x22,0x32,0x42,0x52,0x62,0x72,0x92,0xB2,0xD2,0xF2,0xBB,0xA7,0xB7,0xAF,0xBF,0xA3,0xB3,0x1A,0x3A,0x5A,0x7A,0xDA,0xFA,0x27,0x37,0x2F,0x3F,0x3B,0x23,0x33,0x67,0x77,0x6F,0x7F,0x7B,0x63,0x73,0xEB,0x07,0x17,0x0F,0x1F,0x1B,0x03,0x13,0x47,0x57,0x4F,0x5F,0x5B,0x43,0x53,0x9E,0x9C,0x0C,0x1C,0x3C,0x5C,0x7C,0xDC,0xFC,0x8B,0x9B}
local badopcodes = {}
for i = 1,#_badopcodes do
	badopcodes[_badopcodes[i]] = true
end
_badopcodes = nil
local outbreak = {}

local function gui_print(str)
	message = str
	messagetime = os.time()
end

local function branchCheck()
	local pc = memory.getregister("pc")
	local reset = (memory.readbyte(0xFFFD) * 256) + memory.readbyte(0xFFFC)
	local opcode = memory.readbyte(pc)
	if pc == reset and clean then
		emu.print("[" .. os.date() .. "] Detected Reset")
		stack = {}
		stackg = {}
		accessMap = {}
		entries = 0
		NewState = false
	end
	if (pc >= 0xFFFA or (badcodes and badopcodes[opcode])) and not badTrigger then -- Shouldn't be here
		badTrigger = true
		if (pc >= 0xFFFA) then
			emu.print(string.format("[%s] In vectors at $%04X", os.date(), pc))
		else
			emu.print(string.format("[%s] Bad opcode 0x%02X at $%04X", os.date(), opcode, pc))
		end
		if #stackg > 0 then -- Jump to latest state
			emu.print("Restoring to last state ...")
			NewState = true
			GoodState = true
		elseif #stack > 0 then -- Jump to latest bad state
			emu.print("Restoring to last corrupt state ...")
			NewState = true
			GoodState = false
		else -- Reset emulator
			emu.print("Resetting Emulation ...")
			emu.softreset()
		end
	elseif pc < 0xFFFA and ((not badcodes) or (badcodes and not badopcodes[opcode])) and badTrigger and frameStart then
		badTrigger = false
	end
	if corrupt and not accessMap[pc] and math.random(1,100) <= range then
		-- New location
		accessMap[pc] = true
		entries = entries + 1
		if opcode % 0x20 == 0x10 then -- Branch Opcodes
			local jump, realjump
			local offset = memory.readbytesigned(pc + 1)
			local flags = memory.getregister("p")
			for i = 1,#bops do
				if opcode == bops[i][1] then
					if debug then
						bopcount[i] = bopcount[i] + 1
					end
					if AND(flags,bops[i][2]) == bops[i][3] then
						jump = 2
						realjump = offset + 2
					else
						jump = offset + 2
						realjump = 2
					end
					break
				end
			end
			if jump ~= nil then
				if #stack >= maxStack then
					table.remove(stack, 1)
				end
				local state = savestate.object()
				savestate.save(state)
				stack[#stack+1] = {state, pc + jump}
				
				if (type(outbreak) ~= "table" or outbreak[pc + realjump] ~= true) then
					if #stackg >= maxStack then	
						table.remove(stackg, 1)
					end
					state = savestate.object()
					savestate.save(state)
					stackg[#stackg+1] = {state, pc + realjump}
				end
			end
		end
	elseif accessMap[pc] and not NewState then
		NewState = true
		GoodState = false
	end
	frameStart = false
end

local function everyFrame()
	if debug then
		emu.print(bopcount[1], bopcount[2], bopcount[3], bopcount[4], bopcount[5], bopcount[6], bopcount[7], bopcount[8])
	end
	frameStart = true
	-- Randomly clear an accessMap path
	if recorrupt then
		didkill = didkill + 1
		if didkill >= release then
			didkill = 1
			if entries >= 1 then
				local kill = math.random(1,entries)
				local index = 0
				for k,v in pairs(accessMap) do
					index = index + 1
					if index >= kill then
						accessMap[k] = nil
						entries = entries - 1
						break
					end
				end
			end
		end
	end
	
	-- Do keyboard input here
	local inputs = input.read()
	for k,v in pairs(oldinput) do
		if not inputs[k] then
			oldinput[k] = nil
		end
	end
	for k,v in pairs(inputs) do
		if oldinput[k] and inputs[k] then
			inputs[k] = false
		elseif not oldinput[k] and inputs[k] then
			oldinput[k] = true
		end
	end
	if inputs[keys.decrease_range] then
		range = math.max(range - 1,0)
		gui_print("Range set to " .. range)
	end
	if inputs[keys.increase_range] then
		range = math.min(range + 1,100)
		gui_print("Range set to " .. range)
	end
	if inputs[keys.flush_map] then
		accessMap = {}
		entries = 0
		gui_print("Flushed access map")
	end
	if inputs[keys.restore_good] then
		if #stackg > 0 then
			NewState = true
			GoodState = true
			gui_print("Restoring from last good state")
		else
			gui_print("No good states to restore to")
		end
	end

	-- Check if we want to load a state
	if NewState then
		local pc = memory.getregister("pc")
		NewState = false
		if GoodState then
			if #stackg > 0 then
				local state, pc2 = stackg[#stackg][1], stackg[#stackg][2]
				stackg[#stackg] = nil
				savestate.load(state)
				memory.setregister("pc", pc2)
				outbreak[pc2] = true
			end
		else
			if #stack > 0 then
				local state, pc2 = stack[#stack][1], stack[#stack][2]
				stack[#stack] = nil
				savestate.load(state)
				memory.setregister("pc", pc2)
			end
		end
	end
end

local function dispMessage()
	-- Display message if available
	if os.time() - messagetime <= 5 and message ~= nil then
		gui.text(1, 9, message, "white", "#0000FF80")
	end
	if debug then
		gui.text(1, 17, "States: " .. #stack, "white", "#0000FF80")
		gui.text(1, 25, "GState: " .. #stackg, "white", "#0000FF80")
		gui.text(1, 33, "Entries: " .. string.format("%04X",entries), "white", "#0000FF80")
		gui.text(1, 41, "PC: $" .. string.format("%04X",memory.getregister("pc")), "white", "#0000FF80")
	end
end

memory.registerexecute(0, 65536, branchCheck)
emu.registerbefore(everyFrame)
gui.register(dispMessage)
if debug then
	emu.softreset()
end