NSMB RNG

-- Author: Suuper

-- Set how many changes the script will check up to.
local maxChanges = 10000
local rngAddress = 0x088A68

local initialValue = mainmemory.read_u32_le(rngAddress)
local initialFrame = emu.framecount()

local lastValue = initialValue
local totalChanges = 0

print("Watching RNG, starting frame: " .. initialFrame)

function hex(v)
    return string.format("%x", v)
end

function RNGStep(iRNG)
    -- This forumla gives correct results:
    -- u54 a = (u64)0x0019660D * iRNG + 0x3C6EF35F
    -- return (a + (a >> 32)) & 0xffffffff
    -- However, Lua does not have 64-bit integers.
    -- We will work around this problem by turning the values to be
    -- multiplied into 2-digit base-0x10000 numbers and multiplying those
    -- digits manually.
    local aVal = 0x3C6EF35F
    local mHigh = 0x0019
    local mLow  = 0x660D
    -- Note that (a >> 32) is the same thing as reading the high 4 bytes
    -- as a u32: *(((u32*)a) + 1). We use this fact to avoid creating large -- numbers that might lose integer precision.

    local rngHigh = math.floor(iRNG / 0x10000)
    local rngLow = iRNG % 0x10000

    local low = rngLow * mLow
    -- lowHigh is guaranteed to fit within 32 bits because mLow + mHigh
    -- is less than 0x10000
    local lowHigh = rngHigh * mLow + mHigh * rngLow -- * 0x10000
    local high = rngHigh * mHigh -- * 0x100000000
    -- Instead of multiplying lowHigh and high, we get the bytes that would
    -- end up as the high 32 bits of a u64 directly.
    local u64High = math.floor(lowHigh / 0x10000) + high
    local u64Low = (lowHigh % 0x10000) * 0x10000 + low + aVal
    if u64Low > 0xffffffff then -- this may happen, because +aVal
        u64High = u64High + 1
        -- We mod 0x100000000 later, so changing u64Low isn't necessary.
        -- u64Low = u64Low - 0x100000000
    end

    return (u64Low + u64High) % 0x100000000
end

function watch()
    local rng = mainmemory.read_u32_le(rngAddress)

    local changesThisFrame = 0
    while rng ~= lastValue and changesThisFrame < maxChanges do
        -- Advance one step
        lastValue = RNGStep(lastValue)
        changesThisFrame = changesThisFrame + 1
    end

    -- Make sure the values match. (We might have hit maxChanges)
    if (rng ~= lastValue) then
        -- Or, we might have loaded a savestate.
        local rewindedChanges = checkRewind(rng)
        if (rewindedChanges == -1) then
            print("RNG did not advance as expected. Restting at frame " .. emu.framecount() .. ".")
            -- Reset counter
            totalChanges = 1
            changesThisFrame = -1
            initialValue = rng
            initialFrame = emu.framecount()
            --print("Watching RNG, starting frame: " .. initialFrame)
        else
            print("RNG set back.")
            changesThisFrame = rewindedChanges - totalChanges
        end
    end

    -- Set up for next frame
    totalChanges = totalChanges + changesThisFrame
    lastValue = rng

     -- Print if there was change
    if changesThisFrame ~= 0 then
        print("Total RNG steps: " .. totalChanges .. ", this frame: " .. changesThisFrame)
    end

end

-- If the RNG didn't advance as expected, check if it went backwards.
-- This is probably due to a savestate. In this case, we don't want
-- to reset the counter to 0, we want to go back to that frame's state.
function checkRewind(newRngValue)
    local rng = initialValue
    local changes = 0
    while rng ~= newRngValue and changes < totalChanges do
        rng = RNGStep(rng)
        changes = changes + 1
    end

    if (rng == newRngValue) then
        return changes
    else
        return -1
    end
end

-- Please don't use this in conjuction with the watch function.
function manualAdvance()
    local inputs = joypad.getimmediate();
    if (inputs.B) then
        local rng = mainmemory.read_u32_le(rngAddress);
        rng = RNGStep(rng);
        mainmemory.write_u32_le(rngAddress, rng);
    end
end

-- What is 5-1
local changeHistory = {}
local previousFrame = 0
function test()
    local rng = mainmemory.read_u32_le(rngAddress)

    local changesThisFrame = 0
    while rng ~= lastValue and changesThisFrame < maxChanges do
        -- Advance one step
        lastValue = RNGStep(lastValue)
        changesThisFrame = changesThisFrame + 1
    end

    -- Change history tracking
    local fnum = emu.framecount()
    if fnum == previousFrame + 1 then
        if changeHistory[fnum] ~= changesThisFrame then
            if changesThisFrame ~= 0 or changeHistory[fnum] ~= nil then
                print("Frame " .. fnum .. ": " .. changesThisFrame)
            end
            changeHistory[fnum] = changesThisFrame
        end
    end

    -- Set up for next frame
    lastValue = rng
    previousFrame = fnum
end
--test()

while true do
    -- manualAdvance();
    watch();
    -- test()

    emu.frameadvance();
end