// Copyright (c) 2012- PPSSPP Project / Dolphin Project.

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 2.0 or later versions.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License 2.0 for more details.

// A copy of the GPL 2.0 should have been included with the program.
// If not, see http://www.gnu.org/licenses/

// Official git repository and contact information can be found at
// https://github.com/hrydgard/ppsspp and http://www.ppsspp.org/.

#include "ppsspp_config.h"
#include <cstddef>
#include <algorithm>

#include "ext/xxhash.h"
#include "Common/CommonTypes.h"
#include "Common/Profiler/Profiler.h"

#ifdef _WIN32
#include "Common/CommonWindows.h"
#endif

#include "Core/Core.h"
#include "Core/MemMap.h"
#include "Core/CoreTiming.h"
#include "Core/Reporting.h"
#include "Core/Config.h"

#include "Core/MIPS/MIPS.h"
#include "Core/MIPS/MIPSTables.h"
#include "Core/MIPS/MIPSAnalyst.h"

#include "Core/MIPS/JitCommon/JitBlockCache.h"
#include "Core/MIPS/JitCommon/JitCommon.h"

constexpr u32 INVALID_EXIT = 0xFFFFFFFF;
constexpr u32 SENTINEL_VAL = 0xc0ffeefe;

static uint64_t HashJitBlock(const JitBlock &b) {
	PROFILE_THIS_SCOPE("jithash");
	if (JIT_USE_COMPILEDHASH) {
		// Includes the emuhack (or emuhacks) in memory.
		if (Memory::IsValidRange(b.originalAddress, b.originalSize * 4)) {
			return XXH3_64bits(Memory::GetPointerUnchecked(b.originalAddress), b.originalSize * 4);
		} else {
			// Hm, this would be bad.
			return 0;
		}
	}
	return 0;
}

bool JitBlock::ContainsAddress(u32 address) const {
	// WARNING - THIS DOES NOT WORK WITH JIT INLINING ENABLED.
	// However, that doesn't exist yet so meh.
	return address >= originalAddress && address < originalAddress + 4 * originalSize;
}

void JitBlock::DoIntegrityCheck(u32 emAddress, int num, const char *reason) const {
	_assert_msg_(emAddress == originalAddress, "%s: Bad orig %08x (expected: %08x, snt: %08x) in block %d (b.num: %d) sz: %d", reason, originalAddress, emAddress, sentinel, num, blockNum, codeSize);
	// Also alert just if the sentinel got corrupted.
	_assert_msg_(sentinel == SENTINEL_VAL, "%s: Block %d(%d) sentinel got corrupted: %08x != %08x (origAddr: %08x)", reason, num, blockNum, sentinel, SENTINEL_VAL, originalAddress);
}

JitBlockCache::JitBlockCache(MIPSState *mipsState, CodeBlockCommon *codeBlock) : codeBlock_(codeBlock) {}

JitBlockCache::~JitBlockCache() {
	Shutdown();
}

bool JitBlockCache::IsFull() const {
	// Subtract some amount to safely leave space for some proxy blocks (now obsolete, but let's still have some extra margin).
	return num_blocks_ >= MAX_NUM_BLOCKS - 64;
}

void JitBlockCache::Init() {
	blocks_ = new JitBlock[MAX_NUM_BLOCKS];
	Clear();
}

void JitBlockCache::Shutdown() {
	Clear(); // Make sure proxy block links are deleted
	delete [] blocks_;
	blocks_ = 0;
	num_blocks_ = 0;
}

// This clears the JIT cache. It's called from JitCache.cpp when the JIT cache
// is full and when saving and loading states.
void JitBlockCache::Clear() {
	// Note: We intentionally clear the block_map_ first to avoid O(N^2) behavior in RemoveBlockMap
	block_map_.clear();
	for (int i = 0; i < num_blocks_; i++) {
		DestroyBlock(i, DestroyType::CLEAR);
	}
	links_to_.clear();
	num_blocks_ = 0;

	blockMemRanges_[JITBLOCK_RANGE_SCRATCH] = std::make_pair(0xFFFFFFFF, 0x00000000);
	blockMemRanges_[JITBLOCK_RANGE_RAMBOTTOM] = std::make_pair(0xFFFFFFFF, 0x00000000);
	blockMemRanges_[JITBLOCK_RANGE_RAMTOP] = std::make_pair(0xFFFFFFFF, 0x00000000);
}

void JitBlockCache::Reset() {
	Shutdown();
	Init();
}

int JitBlockCache::AllocateBlock(u32 startAddress) {
	_assert_(num_blocks_ < MAX_NUM_BLOCKS);

	const int numBlocks = num_blocks_;

	JitBlock &b = blocks_[numBlocks];

	b.invalid = false;
	b.originalAddress = startAddress;
	for (int i = 0; i < MAX_JIT_BLOCK_EXITS; ++i) {
		b.exitAddress[i] = INVALID_EXIT;
		b.exitPtrs[i] = 0;
		b.linkStatus[i] = false;
	}
	b.blockNum = numBlocks;
	b.sentinel = SENTINEL_VAL;

	num_blocks_ = numBlocks + 1; //commit the current block
	return numBlocks;
}

void JitBlockCache::AddBlockMap(int block_num) {
	const JitBlock &b = blocks_[block_num];
	// Convert the logical address to a physical address for the block map
	u32 pAddr = b.originalAddress & 0x1FFFFFFF;
	block_map_[std::make_pair(pAddr + 4 * b.originalSize, pAddr)] = block_num;
}

void JitBlockCache::RemoveBlockMap(int block_num) {
	const JitBlock &b = blocks_[block_num];
	if (b.invalid) {
		return;
	}

	const u32 pAddr = b.originalAddress & 0x1FFFFFFF;
	auto it = block_map_.find(std::make_pair(pAddr + 4 * b.originalSize, pAddr));
	if (it != block_map_.end() && it->second == (u32)block_num) {
		block_map_.erase(it);
	} else {
		// It wasn't in there, or it has the wrong key.  Let's search...
		// TODO: This is O(n), so O(n^2) when called for every block.
		for (auto it = block_map_.begin(); it != block_map_.end(); ++it) {
			if (it->second == (u32)block_num) {
				_dbg_assert_(false);
				block_map_.erase(it);
				break;
			}
		}
	}
}

static void ExpandRange(std::pair<u32, u32> &range, u32 newStart, u32 newEnd) {
	range.first = std::min(range.first, newStart);
	range.second = std::max(range.second, newEnd);
}

void JitBlockCache::FinalizeBlock(int block_num, bool block_link) {
	JitBlock &b = blocks_[block_num];

	b.originalFirstOpcode = Memory::Read_Opcode_JIT(b.originalAddress);
	MIPSOpcode opcode = GetEmuHackOpForBlock(block_num);
	Memory::Write_Opcode_JIT(b.originalAddress, opcode);

	// Note that this hashes the emuhack too, which is intentional.
	b.compiledHash = HashJitBlock(b);

	AddBlockMap(block_num);

	if (block_link) {
		for (int i = 0; i < MAX_JIT_BLOCK_EXITS; i++) {
			if (b.exitAddress[i] != INVALID_EXIT) {
				links_to_.emplace(b.exitAddress[i], block_num);
			}
		}

		LinkBlock(block_num);
		LinkBlockExits(block_num);
	}

	const u32 blockEnd = b.originalAddress + b.originalSize * 4 - 4;
	if (Memory::IsScratchpadAddress(b.originalAddress)) {
		ExpandRange(blockMemRanges_[JITBLOCK_RANGE_SCRATCH], b.originalAddress, blockEnd);
	}
	const u32 halfUserMemory = (PSP_GetUserMemoryEnd() - PSP_GetUserMemoryBase()) / 2;
	if (b.originalAddress < PSP_GetUserMemoryBase() + halfUserMemory) {
		ExpandRange(blockMemRanges_[JITBLOCK_RANGE_RAMBOTTOM], b.originalAddress, blockEnd);
	}
	if (blockEnd > PSP_GetUserMemoryBase() + halfUserMemory) {
		ExpandRange(blockMemRanges_[JITBLOCK_RANGE_RAMTOP], b.originalAddress, blockEnd);
	}
}

bool JitBlockCache::RangeMayHaveEmuHacks(u32 start, u32 end) const {
	for (int i = 0; i < JITBLOCK_RANGE_COUNT; ++i) {
		if (end >= blockMemRanges_[i].first && start <= blockMemRanges_[i].second) {
			return true;
		}
	}
	return false;
}

static int binary_search(const JitBlock blocks_[], const u8 *baseoff, int imin, int imax) {
	while (imin < imax) {
		int imid = (imin + imax) >> 1;
		if (blocks_[imid].normalEntry < baseoff)
			imin = imid + 1;
		else
			imax = imid;
	}
	if ((imax == imin) && (blocks_[imin].normalEntry == baseoff))
		return imin;
	else
		return -1;
}

int JitBlockCache::GetBlockNumberFromEmuHackOp(MIPSOpcode inst, bool ignoreBad) const {
	if (!num_blocks_ || !MIPS_IS_EMUHACK(inst)) // definitely not a JIT block
		return -1;
	int off = (inst & MIPS_EMUHACK_VALUE_MASK);

	const u8 *baseoff = codeBlock_->GetBasePtr() + off;
	if (baseoff < codeBlock_->GetBasePtr() || baseoff >= codeBlock_->GetCodePtr()) {
		if (!ignoreBad) {
			ERROR_LOG(Log::JIT, "JitBlockCache: Invalid Emuhack Op %08x", inst.encoding);
		}
		return -1;
	}

	int bl = binary_search(blocks_, baseoff, 0, num_blocks_ - 1);
	if (bl >= 0 && blocks_[bl].invalid) {
		return -1;
	} else {
		return bl;
	}
}

MIPSOpcode JitBlockCache::GetEmuHackOpForBlock(int blockNum) const {
	int off = (int)(blocks_[blockNum].normalEntry - codeBlock_->GetBasePtr());
	return MIPSOpcode(MIPS_EMUHACK_OPCODE | off);
}

int JitBlockCache::GetBlockNumberFromStartAddress(u32 addr) const {
	if (!blocks_ || !Memory::IsValidAddress(addr))
		return -1;

	MIPSOpcode inst = MIPSOpcode(Memory::Read_U32(addr));
	int bl = GetBlockNumberFromEmuHackOp(inst);
	if (bl < 0) {
		return -1;
	}

	if (blocks_[bl].originalAddress != addr)
		return -1;

	return bl;
}

void JitBlockCache::GetBlockNumbersFromAddress(u32 em_address, std::vector<int> *block_numbers) const {
	for (int i = 0; i < num_blocks_; i++)
		if (blocks_[i].ContainsAddress(em_address))
			block_numbers->push_back(i);
}

int JitBlockCache::GetBlockNumberFromAddress(u32 em_address) const {
	for (int i = 0; i < num_blocks_; i++) {
		if (blocks_[i].ContainsAddress(em_address))
			return i;
	}
	return -1;
}

u32 JitBlockCache::GetAddressFromBlockPtr(const u8 *ptr) const {
	if (!codeBlock_->IsInSpace(ptr))
		return (u32)-1;

	for (int i = 0; i < num_blocks_; ++i) {
		const auto &b = blocks_[i];
		if (!b.invalid && ptr >= b.checkedEntry && ptr < b.normalEntry + b.codeSize) {
			return b.originalAddress;
		}
	}

	// It's in jit somewhere, but we must have deleted it.
	return 0;
}

MIPSOpcode JitBlockCache::GetOriginalFirstOp(int block_num) const {
	if (block_num >= num_blocks_ || block_num < 0) {
		return MIPSOpcode(block_num);
	}
	return blocks_[block_num].originalFirstOpcode;
}

void JitBlockCache::LinkBlockExits(int i) {
	JitBlock &b = blocks_[i];
	if (b.invalid) {
		// 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]);
			if (destinationBlock == -1) {
				continue;
			}

			JitBlock &eb = blocks_[destinationBlock];
			// Make sure the destination is not invalid.
			if (!eb.invalid) {
				MIPSComp::jit->LinkBlock(b.exitPtrs[e], eb.checkedEntry);
				b.linkStatus[e] = true;
			}
		}
	}
}

void JitBlockCache::LinkBlock(int i) {
	LinkBlockExits(i);
	JitBlock &b = blocks_[i];
	// equal_range(b) returns pair<iterator,iterator> representing the range
	// of element with key b
	auto ppp = links_to_.equal_range(b.originalAddress);
	if (ppp.first == ppp.second)
		return;
	for (auto iter = ppp.first; iter != ppp.second; ++iter) {
		// INFO_LOG(Log::JIT, "Linking block %i to block %i", iter->second, i);
		LinkBlockExits(iter->second);
	}
}

void JitBlockCache::UnlinkBlock(int i) {
	JitBlock &b = blocks_[i];
	auto ppp = links_to_.equal_range(b.originalAddress);
	if (ppp.first == ppp.second)
		return;
	for (auto iter = ppp.first; iter != ppp.second; ++iter) {
		if ((size_t)iter->second >= num_blocks_) {
			// Something probably went very wrong. Try to stumble along nevertheless.
			ERROR_LOG(Log::JIT, "UnlinkBlock: Invalid block number %d", iter->second);
			continue;
		}
		JitBlock &sourceBlock = blocks_[iter->second];
		for (int e = 0; e < MAX_JIT_BLOCK_EXITS; e++) {
			if (sourceBlock.exitAddress[e] == b.originalAddress)
				sourceBlock.linkStatus[e] = false;
		}
	}
}

std::vector<u32> JitBlockCache::SaveAndClearEmuHackOps() {
	std::vector<u32> result;
	result.resize(num_blocks_);

	for (int block_num = 0; block_num < num_blocks_; ++block_num) {
		JitBlock &b = blocks_[block_num];
		if (b.invalid)
			continue;

		const u32 emuhack = GetEmuHackOpForBlock(block_num).encoding;
		if (Memory::ReadUnchecked_U32(b.originalAddress) == emuhack)
		{
			result[block_num] = emuhack;
			Memory::Write_Opcode_JIT(b.originalAddress, b.originalFirstOpcode);
		}
		else
			result[block_num] = 0;
	}

	return result;
}

void JitBlockCache::RestoreSavedEmuHackOps(const std::vector<u32> &saved) {
	if (num_blocks_ != (int)saved.size()) {
		ERROR_LOG(Log::JIT, "RestoreSavedEmuHackOps: Wrong saved block size.");
		return;
	}

	for (int block_num = 0; block_num < num_blocks_; ++block_num) {
		const JitBlock &b = blocks_[block_num];
		if (b.invalid || saved[block_num] == 0)
			continue;

		// Only if we restored it, write it back.
		if (Memory::ReadUnchecked_U32(b.originalAddress) == b.originalFirstOpcode.encoding)
			Memory::Write_Opcode_JIT(b.originalAddress, MIPSOpcode(saved[block_num]));
	}
}

void JitBlockCache::DestroyBlock(int block_num, DestroyType type) {
	if (block_num < 0 || block_num >= num_blocks_) {
		ERROR_LOG_REPORT(Log::JIT, "DestroyBlock: Invalid block number %d", block_num);
		return;
	}
	JitBlock *b = &blocks_[block_num];
	// No point it being in there anymore.
	RemoveBlockMap(block_num);

	// 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 (type == DestroyType::INVALIDATE)
			ERROR_LOG(Log::JIT, "Invalidating invalid block %d", block_num);
		return;
	}

	b->invalid = true;

	if (!Memory::IsValid4AlignedAddress(b->originalAddress)) {
		_dbg_assert_msg_(false, "Destroying block with invalid original address: %08x (block num: %d)", b->originalAddress, block_num);
		return;
	}

	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 (b->checkedEntry) {
		// We can skip this if we're clearing anyway, which cuts down on protect back and forth on WX exclusive.
		if (type != DestroyType::CLEAR) {
			u8 *writableEntry = codeBlock_->GetWritablePtrFromCodePtr(b->checkedEntry);
			MIPSComp::jit->UnlinkBlock(writableEntry, b->originalAddress);
		}
	} else {
		ERROR_LOG(Log::JIT, "Unlinking block with no entry: %08x (%d)", b->originalAddress, block_num);
	}
}

void JitBlockCache::InvalidateICache(u32 address, const u32 length) {
	// Convert the logical address to a physical address for the block map
	const u32 pAddr = address & 0x1FFFFFFF;
	const u32 pEnd = pAddr + length;

	if (pEnd < pAddr) {
		ERROR_LOG(Log::JIT, "Bad InvalidateICache: %08x with len=%d", address, length);
		return;
	}

	if (pAddr == 0 && pEnd >= 0x1FFFFFFF) {
		InvalidateChangedBlocks();
		return;
	}

	// Blocks may start and end in overlapping ways, and destroying one invalidates iterators.
	// So after destroying one, we start over.
	do {
	restart:
		auto next = block_map_.lower_bound(std::make_pair(pAddr, 0));
		auto last = block_map_.upper_bound(std::make_pair(pEnd + MAX_BLOCK_INSTRUCTIONS, 0));
		// Note that if next is end(), last will be end() too (equal.)
		for (; next != last; ++next) {
			const u32 blockStart = next->first.second;
			const u32 blockEnd = next->first.first;
			if (blockStart < pEnd && blockEnd > pAddr) {
				DestroyBlock(next->second, DestroyType::INVALIDATE);
				// Our iterator is now invalid.  Break and search again.
				// Most of the time there shouldn't be a bunch of matching blocks.
				goto restart;
			}
		}
		// We got here - it wasn't in the map at all (or anymore.)
	} while (false);
}

void JitBlockCache::InvalidateChangedBlocks() {
	// The primary goal of this is to make sure block linking is cleared up.
	for (int block_num = 0; block_num < num_blocks_; ++block_num) {
		JitBlock &b = blocks_[block_num];
		if (b.invalid)
			continue;

		bool changed = false;
		if (JIT_USE_COMPILEDHASH) {
			changed = b.compiledHash != HashJitBlock(b);
		} else {
			const u32 emuhack = GetEmuHackOpForBlock(block_num).encoding;
			changed = Memory::ReadUnchecked_U32(b.originalAddress) != emuhack;
		}

		if (changed) {
			DEBUG_LOG(Log::JIT, "Invalidating changed block at %08x", b.originalAddress);
			DestroyBlock(block_num, DestroyType::INVALIDATE);
		}
	}
}

int JitBlockCache::GetBlockExitSize() {
#if PPSSPP_ARCH(ARM)
	// Will depend on the sequence found to encode the destination address.
	return 0;
#elif PPSSPP_ARCH(X86) || PPSSPP_ARCH(AMD64)
	return 15;
#elif PPSSPP_ARCH(ARM64)
	// Will depend on the sequence found to encode the destination address.
	return 0;
#elif PPSSPP_ARCH(RISCV64)
	// Will depend on the sequence found to encode the destination address.
	return 0;
#elif PPSSPP_ARCH(LOONGARCH64)
	// Will depend on the sequence found to encode the destination address.
	return 0;
#else
#warning GetBlockExitSize unimplemented
	return 0;
#endif
}

void JitBlockCache::ComputeStats(BlockCacheStats &bcStats) const {
	double totalBloat = 0.0;
	double maxBloat = 0.0;
	double minBloat = 1000000000.0;
	for (int i = 0; i < num_blocks_; i++) {
		const JitBlock *b = GetBlock(i);
		double codeSize = (double)b->codeSize;
		if (codeSize == 0)
			continue;
		double origSize = (double)(4 * b->originalSize);
		double bloat = codeSize / origSize;
		if (bloat < minBloat) {
			minBloat = bloat;
			bcStats.minBloatBlock = b->originalAddress;
		}
		if (bloat > maxBloat) {
			maxBloat = bloat;
			bcStats.maxBloatBlock = b->originalAddress;
		}
		totalBloat += bloat;
	}
	bcStats.numBlocks = num_blocks_;
	bcStats.minBloat = (float)minBloat;
	bcStats.maxBloat = (float)maxBloat;
	bcStats.avgBloat = (float)(totalBloat / (double)num_blocks_);
}

JitBlockDebugInfo JitBlockCache::GetBlockDebugInfo(int blockNum) const {
	JitBlockDebugInfo debugInfo{};
	const JitBlock *block = GetBlock(blockNum);
	debugInfo.originalAddress = block->originalAddress;
	debugInfo.origDisasm.reserve(((block->originalAddress + block->originalSize * 4) - block->originalAddress) / 4);
	for (u32 addr = block->originalAddress; addr <= block->originalAddress + block->originalSize * 4; addr += 4) {
		char temp[256];
		MIPSDisAsm(Memory::Read_Instruction(addr), addr, temp, sizeof(temp), true);
		std::string mipsDis = temp;
		debugInfo.origDisasm.push_back(mipsDis);
	}

#if PPSSPP_ARCH(ARM)
	debugInfo.targetDisasm = DisassembleArm2(block->normalEntry, block->codeSize);
#elif PPSSPP_ARCH(ARM64)
	debugInfo.targetDisasm = DisassembleArm64(block->normalEntry, block->codeSize);
#elif PPSSPP_ARCH(X86) || PPSSPP_ARCH(AMD64)
	debugInfo.targetDisasm = DisassembleX86(block->normalEntry, block->codeSize);
#elif PPSSPP_ARCH(RISCV64)
	debugInfo.targetDisasm = DisassembleRV64(block->normalEntry, block->codeSize);
#elif PPSSPP_ARCH(LOONGARCH64)
	debugInfo.targetDisasm = DisassembleLA64(block->normalEntry, block->codeSize);
#endif
	return debugInfo;
}
