Jit feature preparation: Introduce "proxy blocks".

When these are invalidated, the block they point to gets invalidated too.

Will be useful to implement various types of block merging and function inlining
without affecting correctness of cache clears etc.

Also, with this commit we can now fully inline replaced functions. fabsf() boils
down to 1-2 instructions and the block continues, for example.

Core/MIPS/ARM/ArmCompBranch.cpp

@@ -433,6 +433,9 @@ void Jit::Comp_Jump(MIPSOpcode op)
 		break;
 
 	case 3: //jal
+		if (ReplaceJalTo(targetAddr))
+			return;
+
 		gpr.SetImm(MIPS_REG_RA, js.compilerPC + 8);
 		CompileDelaySlot(DELAYSLOT_NICE);
 		if (jo.continueJumps && js.numInstructions < jo.continueMaxInstructions) {

Core/MIPS/ARM/ArmJit.cpp

@@ -339,6 +339,35 @@ void Jit::Comp_RunBlock(MIPSOpcode op)
 	ERROR_LOG(JIT, "Comp_RunBlock should never be reached!");
 }
 
+bool Jit::ReplaceJalTo(u32 dest) {
+	MIPSOpcode op(Memory::Read_Opcode_JIT(dest));
+	if (!MIPS_IS_REPLACEMENT(op.encoding))
+		return false;
+
+	int index = op.encoding & MIPS_EMUHACK_VALUE_MASK;
+	const ReplacementTableEntry *entry = GetReplacementFunc(index);
+	if (!entry) {
+		ERROR_LOG(HLE, "ReplaceJalTo: Invalid replacement op %08x at %08x", op.encoding, dest);
+		return false;
+	}
+
+	// Warning - this might be bad if the code at the destination changes...
+	if (entry->flags & REPFLAG_ALLOWINLINE) {
+		// Jackpot! Just do it, no flushing. The code will be entirely inlined.
+
+		// First, compile the delay slot. It's unconditional so no issues.
+		CompileDelaySlot(DELAYSLOT_NICE);
+		// Technically, we should write the unused return address to RA, but meh.
+		MIPSReplaceFunc repl = entry->jitReplaceFunc;
+		int cycles = (this->*repl)();
+		js.downcountAmount += cycles;
+		// No writing exits, keep going!
+		return true;
+	} else {
+		return false;
+	}
+}
+
 void Jit::Comp_ReplacementFunc(MIPSOpcode op)
 {
 	// We get here if we execute the first instruction of a replaced function. This means

Core/MIPS/ARM/ArmJit.h

@@ -196,6 +196,8 @@ private:
 	void MovFromPC(ARMReg r);
 	void MovToPC(ARMReg r);
 
+	bool ReplaceJalTo(u32 dest);
+
 	void SaveDowncount();
 	void RestoreDowncount();
 

Core/MIPS/JitCommon/JitBlockCache.cpp

@@ -124,6 +124,7 @@ void JitBlockCache::Clear()
 		DestroyBlock(i, false);
 	links_to.clear();
 	block_map.clear();
+	proxyBlockIndices_.clear();
 	num_blocks = 0;
 }
 
@@ -150,6 +151,28 @@ int JitBlockCache::AllocateBlock(u32 em_address)
 		b.linkStatus[i] = false;
 	}
 	b.blockNum = num_blocks;
+	b.isProxy = false;
+	num_blocks++; //commit the current block
+	return num_blocks - 1;
+}
+
+int JitBlockCache::CreateProxyBlock(u32 rootAddress, u32 startAddress, u32 size, const u8 *codePtr) {
+	JitBlock &b = blocks[num_blocks];
+	b.invalid = false;
+	b.originalAddress = startAddress;
+	b.originalSize = size;
+	for (int i = 0; i < MAX_JIT_BLOCK_EXITS; ++i) {
+		b.exitAddress[i] = INVALID_EXIT;
+		b.exitPtrs[i] = 0;
+		b.linkStatus[i] = false;
+	}
+	b.exitAddress[0] = rootAddress;
+	b.blockNum = num_blocks;
+	b.isProxy = true;
+	// Make binary searches and stuff work ok
+	b.normalEntry = codePtr;
+	b.checkedEntry = codePtr;
+	proxyBlockIndices_.push_back(num_blocks);
 	num_blocks++; //commit the current block
 	return num_blocks - 1;
 }
@@ -201,10 +224,8 @@ void JitBlockCache::FinalizeBlock(int block_num, bool block_link)
 #endif
 }
 
-int binary_search(JitBlock blocks[], const u8 *baseoff, int imin, int imax)
-{
-	while (imin < imax)
-	{
+static int binary_search(JitBlock blocks[], const u8 *baseoff, int imin, int imax) {
+	while (imin < imax) {
 		int imid = (imin + imax) / 2;
 		if (blocks[imid].normalEntry < baseoff)
 			imin = imid + 1;
@@ -230,7 +251,7 @@ int JitBlockCache::GetBlockNumberFromEmuHackOp(MIPSOpcode inst, bool ignoreBad)
 		return -1;
 	}
 
-	int bl = binary_search(blocks, baseoff, 0, num_blocks-1);
+	int bl = binary_search(blocks, baseoff, 0, num_blocks - 1);
 	if (blocks[bl].invalid)
 		return -1;
 	return bl;
@@ -241,16 +262,25 @@ MIPSOpcode JitBlockCache::GetEmuHackOpForBlock(int blockNum) const {
 	return MIPSOpcode(MIPS_EMUHACK_OPCODE | off);
 }
 
-int JitBlockCache::GetBlockNumberFromStartAddress(u32 addr)
-{
+int JitBlockCache::GetBlockNumberFromStartAddress(u32 addr, bool realBlocksOnly) {
 	if (!blocks)
 		return -1;
+
 	MIPSOpcode inst = MIPSOpcode(Memory::Read_U32(addr));
 	int bl = GetBlockNumberFromEmuHackOp(inst);
-	if (bl < 0)
+	if (bl < 0) {
+		// Wasn't an emu hack op, look through proxyBlockIndices_. 
+		for (size_t i = 0; i < proxyBlockIndices_.size(); i++) {
+			int blockIndex = proxyBlockIndices_[i];
+			if (blocks[blockIndex].originalAddress == addr && blocks[blockIndex].isProxy && !blocks[blockIndex].invalid)
+				return blockIndex;
+		}
 		return -1;
+	}
+
 	if (blocks[bl].originalAddress != addr)
-		return -1;		
+		return -1;	
+
 	return bl;
 }
 
@@ -272,15 +302,13 @@ u32 JitBlockCache::GetAddressFromBlockPtr(const u8 *ptr) const {
 		}
 	}
 
-	// It's in jit somewhere, but we must've deleted it.
+	// It's in jit somewhere, but we must have deleted it.
 	return 0;
 }
 
 MIPSOpcode JitBlockCache::GetOriginalFirstOp(int block_num)
 {
-	if (block_num >= num_blocks || block_num < 0)
-	{
-		//PanicAlert("JitBlockCache::GetOriginalFirstOp - block_num = %u is out of range", block_num);
+	if (block_num >= num_blocks || block_num < 0) {
 		return MIPSOpcode(block_num);
 	}
 	return blocks[block_num].originalFirstOpcode;
@@ -293,6 +321,7 @@ void JitBlockCache::LinkBlockExits(int i)
 		// This block is dead. Don't relink it.
 		return;
 	}
+
 	for (int e = 0; e < MAX_JIT_BLOCK_EXITS; e++) {
 		if (b.exitAddress[e] != INVALID_EXIT && !b.linkStatus[e]) {
 			int destinationBlock = GetBlockNumberFromStartAddress(b.exitAddress[e]);
@@ -400,28 +429,48 @@ void JitBlockCache::DestroyBlock(int block_num, bool invalidate)
 		ERROR_LOG_REPORT(JIT, "DestroyBlock: Invalid block number %d", block_num);
 		return;
 	}
-	JitBlock &b = blocks[block_num];
-	if (b.invalid) {
+	JitBlock *b = &blocks[block_num];
+
+	// Follow a block proxy chain.
+	// Destroy the block that transitively has this as a proxy. Likely the root block once inlined
+	// this block or its 'parent', so now that this block has changed, the root block must be destroyed.
+	while (b->isProxy) {
+		// false to allow the slow search for proxy blocks.
+		block_num = GetBlockNumberFromStartAddress(b->exitAddress[0], false);  // exitAddress is repurposed for this.
+		JitBlock *proxy = b;
+		proxy->invalid = true;  // The proxy block can safely be invalidated immediately. No need to change anything.
+		b = &blocks[block_num];
+		if (b->invalid) {
+			// Block has already been invalidated, possibly by some other proxy block.
+			return;
+		}
+	}
+
+	// TODO: Handle the case when there's a proxy block and a regular JIT block at the same location.
+	// In this case we probably "leak" the proxy block currently (no memory leak but it'll stay enabled).
+
+	if (b->invalid) {
 		if (invalidate)
 			ERROR_LOG(JIT, "Invalidating invalid block %d", block_num);
 		return;
 	}
-	b.invalid = true;
-	if (Memory::ReadUnchecked_U32(b.originalAddress) == GetEmuHackOpForBlock(block_num).encoding)
-		Memory::Write_Opcode_JIT(b.originalAddress, b.originalFirstOpcode);
-	// It's not safe to set normalEntry to 0 here, since we use a binary search.
+	b->invalid = true;
+	if (Memory::ReadUnchecked_U32(b->originalAddress) == GetEmuHackOpForBlock(block_num).encoding)
+		Memory::Write_Opcode_JIT(b->originalAddress, b->originalFirstOpcode);
+
+	// It's not safe to set normalEntry to 0 here, since we use a binary search
+	// that looks at that later to find blocks. Marking it invalid is enough.
 
 	UnlinkBlock(block_num);
 
-
 #if defined(ARM)
 
 	// Send anyone who tries to run this block back to the dispatcher.
 	// Not entirely ideal, but .. pretty good.
 	// I hope there's enough space...
 	// checkedEntry is the only "linked" entrance so it's enough to overwrite that.
-	ARMXEmitter emit((u8 *)b.checkedEntry);
-	emit.MOVI2R(R0, b.originalAddress);
+	ARMXEmitter emit((u8 *)b->checkedEntry);
+	emit.MOVI2R(R0, b->originalAddress);
 	emit.STR(R0, CTXREG, offsetof(MIPSState, pc));
 	emit.B(MIPSComp::jit->dispatcher);
 	emit.FlushIcache();
@@ -431,12 +480,12 @@ void JitBlockCache::DestroyBlock(int block_num, bool invalidate)
 	// Send anyone who tries to run this block back to the dispatcher.
 	// Not entirely ideal, but .. pretty good.
 	// Spurious entrances from previously linked blocks can only come through checkedEntry
-	XEmitter emit((u8 *)b.checkedEntry);
-	emit.MOV(32, M(&mips->pc), Imm32(b.originalAddress));
+	XEmitter emit((u8 *)b->checkedEntry);
+	emit.MOV(32, M(&mips->pc), Imm32(b->originalAddress));
 	emit.JMP(MIPSComp::jit->Asm().dispatcher, true);
 #elif defined(PPC)
-	PPCXEmitter emit((u8 *)b.checkedEntry);
-	emit.MOVI2R(R3, b.originalAddress);
+	PPCXEmitter emit((u8 *)b->checkedEntry);
+	emit.MOVI2R(R3, b->originalAddress);
 	emit.STW(R0, CTXREG, offsetof(MIPSState, pc));
 	emit.B(MIPSComp::jit->dispatcher);
 	emit.FlushIcache();

Core/MIPS/JitCommon/JitBlockCache.h

@@ -55,6 +55,8 @@ const int MAX_JIT_BLOCK_EXITS = 8;
 // Add the VTune include/lib directories to the project directories to get this to build.
 // #define USE_VTUNE
 
+// We should be careful not to access these block structures during runtime as they are large.
+// Fine to mess with them at block compile time though.
 struct JitBlock {
 	bool ContainsAddress(u32 em_address);
 	
@@ -71,6 +73,7 @@ struct JitBlock {
 	u16 blockNum;
 
 	bool invalid;
+	bool isProxy;  // If set, exitAddress[0] points to the block we're proxying for.
 	bool linkStatus[MAX_JIT_BLOCK_EXITS];
 
 #ifdef USE_VTUNE
@@ -80,13 +83,15 @@ struct JitBlock {
 
 typedef void (*CompiledCode)();
 
-class JitBlockCache
-{
+class JitBlockCache {
 public:
 	JitBlockCache(MIPSState *mips_, CodeBlock *codeBlock);
 	~JitBlockCache();
 
 	int AllocateBlock(u32 em_address);
+	// When a proxy block is invalidated, the block located at the rootAddress
+	// is invalidated too.
+	int CreateProxyBlock(u32 rootAddress, u32 startAddress, u32 size, const u8 *codePtr);
 	void FinalizeBlock(int block_num, bool block_link);
 
 	void Clear();
@@ -100,7 +105,7 @@ public:
 	JitBlock *GetBlock(int block_num);
 
 	// Fast way to get a block. Only works on the first source-cpu instruction of a block.
-	int GetBlockNumberFromStartAddress(u32 em_address);
+	int GetBlockNumberFromStartAddress(u32 em_address, bool realBlocksOnly = true);
 
 	// slower, but can get numbers from within blocks, not just the first instruction.
 	// WARNING! WILL NOT WORK WITH JIT INLINING ENABLED (not yet a feature but will be soon)
@@ -134,6 +139,7 @@ private:
 	MIPSState *mips;
 	CodeBlock *codeBlock_;
 	JitBlock *blocks;
+	std::vector<int> proxyBlockIndices_;
 
 	int num_blocks;
 	std::multimap<u32, int> links_to;

Core/MIPS/x86/CompBranch.cpp

@@ -553,8 +553,8 @@ void Jit::Comp_Jump(MIPSOpcode op)
 
 	case 3: //jal
 		// Special case for branches to "replace functions":
-		// if (ReplaceJalTo(targetAddr))
-		// 	return;
+		if (ReplaceJalTo(targetAddr))
+			return;
 
 		// Save return address - might be overwritten by delay slot.
 		gpr.SetImm(MIPS_REG_RA, js.compilerPC + 8);

Core/MIPS/x86/Jit.cpp

@@ -373,8 +373,8 @@ void Jit::Comp_RunBlock(MIPSOpcode op)
 }
 
 bool Jit::ReplaceJalTo(u32 dest) {
-	MIPSOpcode op(Memory::Read_U32(dest));
-	if (!MIPS_IS_REPLACEMENT(dest))
+	MIPSOpcode op(Memory::Read_Opcode_JIT(dest));
+	if (!MIPS_IS_REPLACEMENT(op.encoding))
 		return false;
 
 	int index = op.encoding & MIPS_EMUHACK_VALUE_MASK;
@@ -395,6 +395,10 @@ bool Jit::ReplaceJalTo(u32 dest) {
 		int cycles = (this->*repl)();
 		js.downcountAmount += cycles;
 		// No writing exits, keep going!
+
+		// Add a trigger so that if the inlined code changes, we invalidate this block.
+		// TODO: Correctly determine the size of this block.
+		blocks.CreateProxyBlock(js.blockStart, dest, 4, GetCodePtr());
 		return true;
 	} else {
 		return false;