Fix 7bag and other jsr-using tests

7bag-test.lua

@@ -4,56 +4,47 @@ require("testing")
 
 local labels = asm.loadlabels("build/7bag.lbl")
 
+local spawnTable = {0x02, 0x07, 0x08, 0x0A, 0x0B, 0x0E, 0x12}
+
 function test_fullbag ()
-	asm.waitbefore()
-	asm.waitbefore()
-	local spawnTable = {0x02,0x07,0x08,0x0A,0x0B,0x0E,0x12}
+	asm.waitexecute(0x813B) -- cmp gameModeState in @mainLoop
+	asm.jsr(labels.disableNmi)
+	local results = {}
+	for _, orientID in pairs(spawnTable) do
+		results[orientID] = 0
+	end
 	math.randomseed(1)
 	memory.writebyte(labels.spawnID, 0)
-	local results = {}
 	for i=1,14 do
-		asm.waitbefore()
 		memory.writebyte(labels.rng_seed, math.random(256))
 		asm.jsr(labels.pickRandomTetrimino_7bag)
 		local a = memory.getregister("a")
-		if results[a] == nil then
-			results[a] = 0
-		end
 		results[a] = results[a] + 1
 	end
-	for _, orientID in ipairs(spawnTable) do
-		local cnt = results[orientID]
+	print(results)
+	for orientID, cnt in pairs(results) do
 		assert(cnt == 2, "orientation: " .. orientID .. " count: " .. cnt)
 	end
 end
 
 function test_distribution ()
-	local spawnTable = {0x02,0x07,0x08,0x0A,0x0B,0x0E,0x12}
-	math.randomseed(1)
+	asm.waitexecute(0x813B) -- cmp gameModeState in @mainLoop
+	asm.jsr(labels.disableNmi)
 	local results = {}
-	for i=1,100*7 do
-		asm.waitbefore()
-		asm.waitbefore()
-		memory.writebyte(labels.spawnID, 0)
-		memory.writebyte(labels.rng_seed, math.random(256))
+	for _, orientID in pairs(spawnTable) do
+		results[orientID] = 0
+	end
+	for i=0,255 do
+		memory.writebyte(labels.spawnID, 0x7F)
+		memory.writebyte(labels.rng_seed, i)
 		asm.jsr(labels.pickRandomTetrimino_7bag)
 		local a = memory.getregister("a")
-		if results[a] == nil then
-			results[a] = 0
-		end
 		results[a] = results[a] + 1
-		asm.waitbefore()
-		emu.poweron()
 	end
-	local imbalanced = false
-	for _, orientID in ipairs(spawnTable) do
-		local cnt = results[orientID]
-		print("orientation: " .. orientID .. " count: " .. cnt)
-		if cnt < 80 or cnt > 120 then
-			imbalanced = true
-		end
+	print(results)
+	for orientID, cnt in pairs(results) do
+		assert(cnt == 36 or cnt == 37, "Imbalanced")
 	end
-	assert(not imbalanced, "Imbalanced")
 end
 
 testing.run()

asm.lua

@@ -86,15 +86,44 @@ function push (byte)
 	memory.setregister("s", s-1)
 end
 
-local jsrretaddr = 1 -- arbitrary address that is not executed elsewhere
--- Jumps to the named address and waits for it to return
+-- Pop byte from stack
+function pop ()
+	local s = memory.getregister("s") + 1
+	memory.setregister("s", s)
+	return memory.readbyte(0x100 + s, byte)
+end
+
+-- Jumps to the named address and waits for it to return. The PC is not
+-- advanced when returning (same PC before and after the call). The current
+-- instruction will be executed, but in a corrupted manner, before running the
+-- subroutine. For example, a lda can put garbage into reg a. A sta could write
+-- to a different memory location. A control instruction will overwrite the
+-- chosen address.
 function jsr (addr)
 	assert(addr ~= nil, "addr must not be nil")
 	local curaddr = memory.getregister("pc")
-	push(((jsrretaddr-1)/256) % 256)
-	push(((jsrretaddr-1)    ) % 256)
-	memory.setregister("pc", addr)
 
-	waitexecute(jsrretaddr)
-	memory.setregister("pc", curaddr)
+	-- Wait until FCEUX has finished the current instruction.
+	-- NOP sled because FCEUX hasn't incremented the PC for the current
+	-- instruction yet, and we don't know how long the instruction is.
+	local lastnopaddr = 0x100 + memory.getregister("s")
+	push(0xEA) -- NOP
+	push(0xEA) -- NOP
+	push(0xEA) -- NOP
+	push(0xEA) -- NOP
+	-- This may corrupt the current instruction, especially for multi-byte
+	-- insturctions. The later bytes will be relative to this new pc.
+	memory.setregister("pc", lastnopaddr-2)
+
+	waitexecute(lastnopaddr)
+	pop()
+	pop()
+	pop()
+	pop()
+
+	push(((curaddr-1)/256) % 256)
+	push(((curaddr-1)    ) % 256)
+
+	memory.setregister("pc", addr-1) -- -1 for the NOP's size
+	waitexecute(curaddr)
 end

chart-test.lua

@@ -5,6 +5,8 @@ require("testing")
 local labels = asm.loadlabels("build/taus.lbl")
 
 function test_chartEffConvert ()
+	asm.waitexecute(0x813B) -- cmp gameModeState in @mainLoop
+	asm.jsr(labels.disableNmi)
 	local chartEffConvertDivisor
 	if labels.chartEffConvert == labels.div3 then
 		chartEffConvertDivisor = 3
@@ -14,8 +16,6 @@ function test_chartEffConvert ()
 		error("unknown chartEffConvert")
 	end
 	for i=0,255 do
-		asm.waitbefore()
-		asm.waitbefore()
 		local raw = i / chartEffConvertDivisor
 		local expected = math.floor(i / chartEffConvertDivisor)
 		memory.setregister("a", i)
@@ -24,14 +24,12 @@ function test_chartEffConvert ()
 		if result ~= expected then
 			error("i " .. i .. " expected: " .. expected .. " actual: " .. result .. " raw: " .. raw)
 		end
-		asm.waitbefore()
-		emu.poweron()
 	end
 end
 
 function test_drawChartBackground ()
-	asm.waitbefore()
-	asm.waitbefore()
+	asm.waitexecute(0x813B) -- cmp gameModeState in @mainLoop
+	asm.jsr(labels.disableNmi)
 	local levelEffs = {8, 8, 9, 9, 8, 9, 16, 17, 48, 1}
 	for i, levelEffs in ipairs(levelEffs) do
 		memory.writebyte(labels.levelEffs+(i-1), levelEffs)
@@ -72,8 +70,8 @@ function test_drawChartBackground ()
 end
 
 function test_drawChartSprites ()
-	asm.waitbefore()
-	asm.waitbefore()
+	asm.waitexecute(0x813B) -- cmp gameModeState in @mainLoop
+	asm.jsr(labels.disableNmi)
 	local levelEffs = {
 		48, 14, 14, 48,
 		16, 14, 14, 16,