BizHawk Lua Script to automatically find cheat codes

find_cheat.lua

console.clear()

st1 = 'save_stage1'
st2 = 'save_stage2'
st3 = 'save_stage3'

-- initial memory lookup method [increase, decrease, diff]
lookup_method = 'increase'

-- joystick definitions
frames_per_input = 1
joystick_number = 1
input = joypad.get(joystick_number)

-- input = {}
-- input['A'] = false
-- input['B'] = false
-- input['Down'] = false
-- input['Left'] = false
-- input['Right'] = false
-- input['Select'] = false
-- input['Start'] = false
-- input['Up'] = false

-- initializing vars
candidates = {}
code_list = {}
final_input_list = {}
min_code_length = 5
max_code_length = 20

-- ROM and memory Info
rom_name = gameinfo.getromname()
mem_size = mainmemory.getcurrentmemorydomainsize()
mem_domain = mainmemory.getname()
console.log(string.format('ROM Name: %s', rom_name))
console.log(string.format('Mem Domain: %s', mem_domain))
console.log(string.format('Mem Size: 0x%X', mem_size))


function frame_advance(frames_per_input)
    for i=1, frames_per_input do
        emu.frameadvance()
    end
end

function press_button ( input, button, joystick_number, frames_per_input )
    frame_advance(frames_per_input)
    input[button] = true
    joypad.set(input, joystick_number)
    frame_advance(frames_per_input)
    input[button] = false
end

function lookup_offset( st, mem )
    for button, _ in pairs(input) do
        savestate.load(st)        
        press_button ( input, button, joystick_number, frames_per_input )
        mem_new = mainmemory.read_bytes_as_dict(0x0, mem_size)
        
        for offset, _ in pairs(mem) do
            if lookup_method == 'increase' then
                if mem_new[offset] == math.fmod(256 + (mem[offset] + 1), 256) then
                    candidates[offset] = mem_new[offset]
                end
            elseif lookup_method == 'decrease' then
                if mem_new[offset] == mem[offset] - 1 then
                    candidates[offset] = mem_new[offset]
                end
            elseif lookup_method == 'diff' then
                if mem_new[offset] ~= mem[offset] then
                    candidates[offset] = mem_new[offset]
                end
            end
        end
    end
end

function lookup_final_input( st, offset, input )
    final_input_list[offset] = {}
    savestate.save(st)
    frame_advance(2 * frames_per_input)
    mem_final = mainmemory.read_bytes_as_dict(0x0, mem_size)

    for button, _ in pairs(input) do
        savestate.load(st)
        press_button ( input, button, joystick_number, frames_per_input )
        mem_final_after_input = mainmemory.read_bytes_as_dict(0x0, mem_size)

        mem_changes = 0
        for new_offset, mem_value in pairs(mem_final_after_input) do          
            if mem_final_after_input[new_offset] ~= mem_final[new_offset] then
                mem_changes = mem_changes + 1
                --console.log(string.format("%s\tchanges [0x%X] from %s to %s [framecount %s]", button, new_offset, mem_final[new_offset], mem_final_after_input[new_offset], emu.framecount()))
            end
        end        
        final_input_list[offset][button] = mem_changes
    end
end

function update_results(offset, button, mem_old, mem, mem_new, input_count)
    code_list[offset][input_count] = button
    -- debug printing
    if input_count == 1 then console.log(string.format('=== Offset: 0x%X ===',offset)) end
    console.log(string.format("%s\tchanges [0x%X] from %s to %s [framecount %s, mem_old %s, input_count %s]", button, offset, mem, mem_new, emu.framecount(), mem_old, input_count))
end

function lookup_recursion( lookup_method, offset, mem_old, mem, mem_new, button, input, input_count )
    savestate.save(st2)
    frame_advance(2 * frames_per_input)

    mem_old = mem
    mem = mem_new

    for button, _ in pairs(input) do
        savestate.load(st2)
        press_button(input, button, joystick_number, frames_per_input)
        mem_new = mainmemory.readbyte(offset)    

        lookup_input(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)                
    end
end

function start_recursion(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)
    if code_list[offset][input_count] == nil and input_count <= max_code_length then        
        update_results(offset, button, mem_old, mem, mem_new, input_count)
        input_count = input_count + 1
        lookup_recursion( lookup_method, offset, mem_old, mem, mem_new, button, input, input_count )        
    else
        if final_input_list[offset] == nil then
            lookup_final_input( st3, offset, input )
        end
    end
end 

function lookup_input(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)
    if lookup_method == 'increase' then
        if mem_new == math.fmod(256 + (mem + 1), 256) then
            start_recursion(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)
        end
    elseif lookup_method == 'decrease' then
        if mem_new == mem_old - 1 then
            start_recursion(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)
        end
    elseif lookup_method == 'diff' then
        if mem_new ~= mem_old then
            start_recursion(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)
        end
    end
    savestate.load(st1)
end

function main()
    
    -- take base savestate and advance (2 * frames_per_input)
    savestate.save(st1)
    frame_advance(2 * frames_per_input)
    mem_original = mainmemory.read_bytes_as_dict(0x0,mem_size)
    client.unpause()
    
    -- find all memory offsets that change after button press
    lookup_offset(st1, mem_original)
    
    for offset, v in pairs(candidates) do
        input_count = 1
        code_list[offset] = {}

        mem = mainmemory.readbyte(offset)
        mem_old = mem        
        for button, mem_value in pairs(input) do             
            savestate.load(st1)
            press_button(input, button, joystick_number, frames_per_input)
            mem_new = mainmemory.readbyte(offset)
            lookup_input(lookup_method, offset, mem_old, mem, mem_new, button, input, input_count)            
        end
        
    end

    
    -- reload base savestate
    savestate.load(st1)

    -- print results
    console.log('===================')
    console.log('===   Results  ====')
    console.log('===================')
    for offset, result in pairs(code_list) do 
        if #result >= min_code_length then
            console.log(string.format('[0x%X] %s', offset, table.concat(result,' + ')))
            if final_input_list[offset] ~= nil then
                for final_button, mem_changes in pairs(final_input_list[offset]) do
                    console.log(string.format('> Input %s changes %s memory addresses', final_button, final_input_list[offset][final_button]))
                end
            end
        end
    end
    console.log('===================')
    console.log('Done!')
end

main()