// Copyright (c) 2012- PPSSPP 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 <set>
#include <algorithm>

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

#include "Common/Log.h"
#include "Common/Serialize/Serializer.h"
#include "Common/StringUtils.h"

#include "Core/Config.h"
#include "Core/Core.h"
#include "Core/CoreTiming.h"
#include "Core/HLE/ReplaceTables.h"
#include "Core/HLE/sceKernelMemory.h"
#include "Core/MemMap.h"
#include "Core/MIPS/MIPS.h"
#include "Core/MIPS/MIPSCodeUtils.h"
#include "Core/MIPS/Interpreter.h"
#include "Core/MIPS/MIPSTables.h"
#include "Core/MIPS/IR/IRRegCache.h"
#include "Core/MIPS/IR/IRInterpreter.h"
#include "Core/MIPS/IR/IRJit.h"
#include "Core/MIPS/IR/IRNativeCommon.h"
#include "Core/MIPS/JitCommon/JitCommon.h"
#include "Core/Reporting.h"
#include "Common/TimeUtil.h"
#include "Core/MIPS/MIPSTracer.h"


namespace MIPSComp {

IRJit::IRJit(MIPSState *mipsState, bool actualJit) : frontend_(mipsState->HasDefaultPrefix()), mips_(mipsState), blocks_(actualJit) {
	// u32 size = 128 * 1024;
	InitIR();

	compileToNative_ = actualJit;

	// If this IRJit instance will be used to drive a "JIT using IR", don't optimize for interpretation.
	jo.optimizeForInterpreter = !actualJit;

	IROptions opts{};
	opts.disableFlags = g_Config.uJitDisableFlags;
#if PPSSPP_ARCH(RISCV64)
	// Assume RISC-V always has very slow unaligned memory accesses.
	opts.unalignedLoadStore = false;
	opts.unalignedLoadStoreVec4 = true;
	opts.preferVec4 = cpu_info.RiscV_V;
#elif PPSSPP_ARCH(ARM) || PPSSPP_ARCH(ARM64)
	opts.unalignedLoadStore = (opts.disableFlags & (uint32_t)JitDisable::LSU_UNALIGNED) == 0;
	opts.unalignedLoadStoreVec4 = true;
	opts.preferVec4 = true;
#else
	opts.unalignedLoadStore = (opts.disableFlags & (uint32_t)JitDisable::LSU_UNALIGNED) == 0;
	// TODO: Could allow on x86 pretty easily...
	opts.unalignedLoadStoreVec4 = false;
	opts.preferVec4 = true;
#endif
	opts.optimizeForInterpreter = jo.optimizeForInterpreter;
	frontend_.SetOptions(opts);
}

IRJit::~IRJit() {
}

void IRJit::DoState(PointerWrap &p) {
	frontend_.DoState(p);
}

void IRJit::UpdateFCR31() {
}

void IRJit::ClearCache() {
	INFO_LOG(Log::JIT, "IRJit: Clearing the block cache!");
	blocks_.Clear();
}

void IRJit::InvalidateCacheAt(u32 em_address, int length) {
	std::vector<int> numbers = blocks_.FindInvalidatedBlockNumbers(em_address, length);
	if (numbers.empty()) {
		return;
	}

	DEBUG_LOG(Log::JIT, "Invalidating IR block cache at %08x (%d bytes): %d blocks", em_address, length, (int)numbers.size());

	for (int block_num : numbers) {
		auto block = blocks_.GetBlock(block_num);
		// TODO: We are invalidating a lot of blocks that are already invalid (yu gi oh).
		// INFO_LOG(Log::JIT, "Block at %08x invalidated: valid: %d", block->GetOriginalStart(), block->IsValid());
		// If we're a native JIT (IR->JIT, not just IR interpreter), we write native offsets into the blocks.
		int cookie = compileToNative_ ? block->GetNativeOffset() : block->GetIRArenaOffset();
		// Before Destroy, which forgets the block's address. Blocks linked to it jump straight in.
		InvalidateNativeBlock(&blocks_, block_num);
		blocks_.RemoveBlockFromPageLookup(block_num);
		block->Destroy(cookie);
	}
}

void IRJit::Compile(u32 em_address) {
	_dbg_assert_(compilerEnabled_);

	PROFILE_THIS_SCOPE("jitc");

	std::vector<IRInst> instructions;
	u32 mipsBytes;
	if (!CompileBlock(em_address, instructions, mipsBytes)) {
		// Ran out of block numbers - need to reset.
		ERROR_LOG(Log::JIT, "Ran out of block numbers, clearing cache");
		ClearCache();
		CompileBlock(em_address, instructions, mipsBytes);
	}

	if (frontend_.CheckRounding(em_address)) {
		// Our assumptions are all wrong so it's clean-slate time.
		ClearCache();
		CompileBlock(em_address, instructions, mipsBytes);
	}
}

// WARNING! This can be called from IRInterpret / the JIT, through the function preload stuff!
bool IRJit::CompileBlock(u32 em_address, std::vector<IRInst> &instructions, u32 &mipsBytes) {
	_dbg_assert_(compilerEnabled_);

	frontend_.DoJit(em_address, instructions, mipsBytes);
	_dbg_assert_(!instructions.empty());

	int block_num = blocks_.AllocateBlock(em_address, mipsBytes, instructions);
	if (block_num < 0) {
		WARN_LOG(Log::JIT, "Failed to allocate block for %08x (%d instructions)", em_address, (int)instructions.size());
		// Out of block numbers.  Caller will handle.
		return false;
	}

	IRBlock *b = blocks_.GetBlock(block_num);
	if (mipsTracer.tracing_enabled) {
		// Hash, then only update page stats, don't link yet.
		// TODO: Should we always hash?  Then we can reuse blocks.
		b->UpdateHash();
	}

	if (!CompileNativeBlock(&blocks_, block_num))
		return false;

	if (mipsTracer.tracing_enabled) {
		mipsTracer.prepare_block(b, blocks_);
	}

	// Updates stats, also enters the block in the block shadow.
	blocks_.FinalizeBlock(block_num);
	FinalizeNativeBlock(&blocks_, block_num);
	return true;
}

void IRJit::RunLoopUntil(u64 globalticks) {
	PROFILE_THIS_SCOPE("jit");

	// ApplyRoundingMode(true);
	// IR Dispatcher
	
	MIPSState *mips = mips_;
	while (true) {
		// RestoreRoundingMode(true);
		CoreTiming::Advance(currentMIPS);
		// ApplyRoundingMode(true);
		if (coreState != 0) {
			break;
		}

#ifdef _DEBUG
		compilerEnabled_ = false;
#endif
		while (mips->downcount >= 0) {
			// The block's offset into the IR arena, or 0 for none, or a replacement hook (handled by
			// compiling a block).
			const u32 offset = *Memory::GetBlockShadowEntry(mips->pc);
			static_assert(Memory::BLOCK_SHADOW_HOOK_NO_JIT == 1, "IR arena reserves offsets 0 and 1");
			if (offset > Memory::BLOCK_SHADOW_HOOK_NO_JIT) {
				const IRInst *instPtr = blocks_.GetArenaPtr() + offset;
				// First op is always, except when using breakpoints, downcount, to save one dispatch inside IRInterpret.
				// This branch is very cpu-branch-predictor-friendly so this still beats the dispatch.
				if (instPtr->op == IROp::Downcount) {
					mips->downcount -= instPtr->constant;
					instPtr++;
				}
#ifdef IR_PROFILING
				IRBlock *block = blocks_.GetBlock(blocks_.GetBlockNumFromIRArenaOffset(offset));
				Instant start = Instant::Now();
				mips->pc = IRInterpret(mips, instPtr);
				int64_t elapsedNanos = start.ElapsedNanos();
				block->profileStats_.executions += 1;
				block->profileStats_.totalNanos += elapsedNanos;
#else
				mips->pc = IRInterpret(mips, instPtr);
#endif
				// Note: this will "jump to zero" on a badly constructed block missing exits.
				if (!Memory::IsValid4AlignedAddress(mips->pc)) {
					int blockNum = blocks_.GetBlockNumFromIRArenaOffset(offset);
					IRBlock *block = blocks_.GetBlockUnchecked(blockNum);
					Core_ExecException(mips->pc, block->GetOriginalStart(), ExecExceptionType::JUMP);
					break;
				}
			} else {
				// RestoreRoundingMode(true);
#ifdef _DEBUG
				compilerEnabled_ = true;
#endif
				Compile(mips->pc);
#ifdef _DEBUG
				compilerEnabled_ = false;
#endif
				// ApplyRoundingMode(true);
			}
		}
#ifdef _DEBUG
		compilerEnabled_ = true;
#endif
	}

	// RestoreRoundingMode(true);
}

bool IRJit::DescribeCodePtr(const u8 *ptr, std::string &name) {
	// Used in native disassembly viewer.
	return false;
}

void IRJit::LinkBlock(u8 *exitPoint, const u8 *checkedEntry) {
	Crash();
}

void IRJit::UnlinkBlock(u8 *checkedEntry, u32 originalAddress) {
	Crash();
}

void IRBlockCache::Clear() {
	for (int i = 0; i < (int)blocks_.size(); ++i) {
		int cookie = compileToNative_ ? blocks_[i].GetNativeOffset() : blocks_[i].GetIRArenaOffset();
		blocks_[i].Destroy(cookie);
	}
	blocks_.clear();
	byPage_.clear();
	arena_.clear();
	arena_.shrink_to_fit();
}

IRBlockCache::IRBlockCache(bool compileToNative) : compileToNative_(compileToNative) {}

int IRBlockCache::AllocateBlock(int emAddr, u32 origSize, const std::vector<IRInst> &insts) {
	const u32 MAX_ARENA_SIZE = 0x1000000 - 1;
	if (arena_.empty()) {
		// Offset 0 would mean "no block" in the block shadow, and 1 a replacement hook.
		arena_.push_back(IRInst{ IROp::Bad });
		arena_.push_back(IRInst{ IROp::Bad });
	}
	int offset = (int)arena_.size();
	if (offset >= MAX_ARENA_SIZE) {
		WARN_LOG(Log::JIT, "Filled JIT arena, restarting");
		return -1;
	}
	// TODO: Use memcpy.
	for (size_t i = 0; i < insts.size(); i++) {
		arena_.push_back(insts[i]);
	}
	int newBlockIndex = (int)blocks_.size();
	blocks_.push_back(IRBlock(emAddr, origSize, offset, (u32)insts.size()));
	return newBlockIndex;
}

int IRBlockCache::GetBlockNumFromIRArenaOffset(int offset) const {
	// Block offsets are always in rising order (we don't go back and replace them when invalidated). So we can binary search.
	int low = 0;
	int high = (int)blocks_.size() - 1;
	int found = -1;
	while (low <= high) {
		int mid = low + (high - low) / 2;
		const int blockOffset = blocks_[mid].GetIRArenaOffset();
		if (blockOffset == offset) {
			found = mid;
			break;
		}
		if (blockOffset < offset) {
			low = mid + 1;
		} else {
			high = mid - 1;
		}
	}
	return found;
#if 1
	return found;
#else
	// Cross check the result. This is not fast so normally not enabled. Called a lot when IR_PROFILING is on.
	for (int i = 0; i < (int)blocks_.size(); i++) {
		if (blocks_[i].GetIRArenaOffset() == offset) {
			_dbg_assert_(i == found);
			return i;
		}
	}
	_dbg_assert_(found == -1);
	return -1;
#endif
}

std::vector<int> IRBlockCache::FindInvalidatedBlockNumbers(u32 address, u32 lengthInBytes) {
	u32 startPage = AddressToPage(address);
	u32 endPage = AddressToPage(address + lengthInBytes);

	std::vector<int> found;
	for (u32 page = startPage; page <= endPage; ++page) {
		const auto iter = byPage_.find(page);
		if (iter == byPage_.end())
			continue;

		const std::vector<int> &blocksInPage = iter->second;
		for (int i : blocksInPage) {
			if (blocks_[i].OverlapsRange(address, lengthInBytes)) {
				// We now try to remove these during invalidation.
				found.push_back(i);
			}
		}
	}

	return found;
}

void IRBlockCache::FinalizeBlock(int blockIndex) {
	// TODO: What's different about preload blocks?
	IRBlock &block = blocks_[blockIndex];
	int cookie = compileToNative_ ? block.GetNativeOffset() : block.GetIRArenaOffset();
	block.Finalize(cookie);

	u32 startAddr, size;
	block.GetRange(&startAddr, &size);

	u32 startPage = AddressToPage(startAddr);
	u32 endPage = AddressToPage(startAddr + size);

	for (u32 page = startPage; page <= endPage; ++page) {
		byPage_[page].push_back(blockIndex);
	}
}

// Call after Destroy-ing it.
void IRBlockCache::RemoveBlockFromPageLookup(int blockIndex) {
	// We need to remove the block from the byPage lookup.
	IRBlock &block = blocks_[blockIndex];

	u32 startAddr, size;
	block.GetRange(&startAddr, &size);

	u32 startPage = AddressToPage(startAddr);
	u32 endPage = AddressToPage(startAddr + size);

	for (u32 page = startPage; page <= endPage; ++page) {
		auto iter = std::find(byPage_[page].begin(), byPage_[page].end(), blockIndex);
		if (iter != byPage_[page].end()) {
			byPage_[page].erase(iter);
		} else if (block.IsValid()) {
			// If it was previously invalidated, we don't care, hence the above check.
			WARN_LOG(Log::JIT, "RemoveBlock: Block at %08x was not found where expected in byPage table.", startAddr);
		}
	}

	// Additionally, we'd like to zap the block in the IR arena.
	// However, this breaks if calling sceKernelIcacheClearAll(), since as soon as we return, we'll be executing garbage.
	/*
	IRInst bad{ IROp::Bad };
	for (int off = block.GetIRArenaOffset(); off < (int)(block.GetIRArenaOffset() + block.GetNumIRInstructions()); off++) {
		arena_[off] = bad;
	}
	*/
}

u32 IRBlockCache::AddressToPage(u32 addr) const {
	// Use relatively small pages since basic blocks are typically small.
	return (addr & 0x3FFFFFFF) >> 10;
}

int IRBlockCache::FindPreloadBlock(u32 em_address) {
	u32 page = AddressToPage(em_address);
	auto iter = byPage_.find(page);
	if (iter == byPage_.end())
		return -1;

	const std::vector<int> &blocksInPage = iter->second;
	for (int i : blocksInPage) {
		if (blocks_[i].GetOriginalStart() == em_address) {
			if (blocks_[i].HashMatches()) {
				return i;
			}
		}
	}

	return -1;
}

int IRBlockCache::FindByCookie(int cookie) {
	if (blocks_.empty())
		return -1;

	// TODO: Maybe a flag to determine native offset mode?
	if (!compileToNative_) {
		return GetBlockNumFromIRArenaOffset(cookie);
	}

	// TODO: This could also use a binary search.
	for (int i = 0; i < GetNumBlocks(); ++i) {
		int offset = blocks_[i].GetNativeOffset();
		if (offset == cookie)
			return i;
	}
	return -1;
}

JitBlockDebugInfo IRBlockCache::GetBlockDebugInfo(int blockNum) const {
	const IRBlock &ir = blocks_[blockNum];
	JitBlockDebugInfo debugInfo{};
	uint32_t start, size;
	ir.GetRange(&start, &size);
	debugInfo.originalAddress = start;  // TODO
	if (!Memory::IsValid4AlignedAddress(start)) {
		return debugInfo;
	}

	debugInfo.origDisasm.reserve(((start + size) - start) / 4);
	for (u32 addr = start; addr < start + size; addr += 4) {
		char temp[256];
		MIPSDisAsm(Replacement_GetDisplayOp(addr), addr, temp, sizeof(temp), true);
		std::string mipsDis = temp;
		debugInfo.origDisasm.push_back(mipsDis);
	}

	debugInfo.irDisasm.reserve(ir.GetNumIRInstructions());
	const IRInst *instructions = GetBlockInstructionPtr(ir);
	for (int i = 0; i < ir.GetNumIRInstructions(); i++) {
		IRInst inst = instructions[i];
		char buffer[256];
		DisassembleIR(buffer, sizeof(buffer), inst);
		debugInfo.irDisasm.push_back(buffer);
	}
	return debugInfo;
}

void IRBlockCache::ComputeStats(BlockCacheStats &bcStats) const {
	double totalBloat = 0.0;
	double maxBloat = 0.0;
	double minBloat = 1000000000.0;
	for (const auto &b : blocks_) {
		double codeSize = (double)b.GetNumIRInstructions() * 4;  // We count bloat in instructions, not bytes. sizeof(IRInst);
		if (codeSize == 0)
			continue;
		u32 origAddr, mipsBytes;
		b.GetRange(&origAddr, &mipsBytes);
		double origSize = (double)mipsBytes;
		double bloat = codeSize / origSize;
		if (bloat < minBloat) {
			minBloat = bloat;
			bcStats.minBloatBlock = origAddr;
		}
		if (bloat > maxBloat) {
			maxBloat = bloat;
			bcStats.maxBloatBlock = origAddr;
		}
		totalBloat += bloat;
	}
	bcStats.numBlocks = (int)blocks_.size();
	bcStats.minBloat = minBloat;
	bcStats.maxBloat = maxBloat;
	bcStats.avgBloat = totalBloat / (double)blocks_.size();
}

int IRBlockCache::GetBlockNumberFromStartAddress(u32 em_address) const {
	u32 page = AddressToPage(em_address);

	const auto iter = byPage_.find(page);
	if (iter == byPage_.end())
		return -1;

	const std::vector<int> &blocksInPage = iter->second;
	int best = -1;
	for (int i : blocksInPage) {
		if (blocks_[i].GetOriginalStart() == em_address) {
			best = i;
			if (blocks_[i].IsValid()) {
				return i;
			}
		}
	}
	return best;
}

void IRBlock::Finalize(int cookie) {
	// Check it wasn't invalidated, in case this is after preload.
	// TODO: Allow reusing blocks when the code matches hash_ again, instead.
	if (origAddr_ && Memory::IsValid4AlignedAddress(origAddr_)) {
		origFirstOpcode_ = Memory::ReadUnchecked_Instruction(origAddr_);
		_dbg_assert_(cookie != 0);
		Memory::WriteBlockShadow(origAddr_, cookie);
	} else {
		WARN_LOG(Log::JIT, "Finalizing invalid block (cookie: %d)", cookie);
	}
}

void IRBlock::Destroy(int cookie) {
	if (origAddr_ && Memory::IsValid4AlignedAddress(origAddr_)) {
		// A block at a mirror may have taken over the entry, leave that alone.
		Memory::ClearBlockShadow(origAddr_, cookie);
		// TODO: Also wipe the block in the IR opcode arena.
		// Let's mark this invalid so we don't try to clear it again.
		origAddr_ = 0;
	}
}

u64 IRBlock::CalculateHash() const {
	if (origAddr_ && Memory::IsValidRange(origAddr_, origSize_)) {
		return XXH3_64bits(Memory::GetPointerUnchecked(origAddr_), origSize_);
	}
	return 0;
}

bool IRBlock::OverlapsRange(u32 addr, u32 size) const {
	addr &= 0x3FFFFFFF;
	u32 origAddr = origAddr_ & 0x3FFFFFFF;
	return addr + size > origAddr && addr < origAddr + origSize_;
}

}  // namespace MIPSComp