// Copyright 2019 The SwiftShader Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "LLVMReactorDebugInfo.hpp"

#ifdef ENABLE_RR_DEBUG_INFO

#include "Reactor.hpp"
#include "LLVMReactor.hpp"

#include "boost/stacktrace.hpp"

#include "llvm/Demangle/Demangle.h"
#include "llvm/ExecutionEngine/JITEventListener.h"
#include "llvm/IR/DIBuilder.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/IRBuilder.h"

#include <cctype>
#include <fstream>
#include <mutex>
#include <regex>
#include <sstream>
#include <string>

#if 0
#define LOG(msg, ...) printf(msg "\n", ##__VA_ARGS__)
#else
#define LOG(msg, ...)
#endif

namespace
{
	std::pair<llvm::StringRef, llvm::StringRef> splitPath(const char* path)
	{
		return llvm::StringRef(path).rsplit('/');
	}

	// Note: createGDBRegistrationListener() returns a pointer to a singleton.
	// Nothing is actually created.
	auto jitEventListener = llvm::JITEventListener::createGDBRegistrationListener(); // guarded by jitEventListenerMutex
	std::mutex jitEventListenerMutex;

} // anonymous namespaces

namespace rr
{
	DebugInfo::DebugInfo(
			llvm::IRBuilder<> *builder,
			llvm::LLVMContext *context,
			llvm::Module *module,
			llvm::Function *function)
		: builder(builder), context(context), module(module), function(function)
	{
		using namespace ::llvm;

		auto location = getCallerLocation();

		auto fileAndDir = splitPath(location.function.file.c_str());
		diBuilder.reset(new llvm::DIBuilder(*module));
		diCU = diBuilder->createCompileUnit(
			llvm::dwarf::DW_LANG_C,
			diBuilder->createFile(fileAndDir.first, fileAndDir.second),
			"Reactor",
			0, "", 0);

		registerBasicTypes();

		SmallVector<Metadata *, 8> EltTys;
		auto funcTy = diBuilder->createSubroutineType(diBuilder->getOrCreateTypeArray(EltTys));

		auto file = getOrCreateFile(location.function.file.c_str());
		auto sp = diBuilder->createFunction(
			file,                   // scope
			"ReactorFunction",      // function name
			"ReactorFunction",      // linkage
			file,                   // file
			location.line,          // line
			funcTy,                 // type
			false,                  // internal linkage
			true,                   // definition
			location.line,          // scope line
			DINode::FlagPrototyped, // flags
			false                   // is optimized
		);
		diSubprogram = sp;
		function->setSubprogram(sp);
		diRootLocation = DILocation::get(*context, location.line, 0, sp);
		builder->SetCurrentDebugLocation(diRootLocation);
	}

	DebugInfo::~DebugInfo() = default;

	void DebugInfo::Finalize()
	{
		while (diScope.size() > 0)
		{
			emitPending(diScope.back(), builder);
			diScope.pop_back();
		}
		diBuilder->finalize();
	}

	void DebugInfo::EmitLocation()
	{
		auto const& backtrace = getCallerBacktrace();
		syncScope(backtrace);
		builder->SetCurrentDebugLocation(getLocation(backtrace, backtrace.size() - 1));

#ifdef ENABLE_RR_EMIT_PRINT_LOCATION
		static Location lastLocation;
		if (backtrace.size() == 0)
		{
			return;
		}
		Location currLocation = backtrace[backtrace.size() - 1];
		if (currLocation != lastLocation)
		{
			rr::Print("rr> {0} [{1}:{2}]\n", currLocation.function.name.c_str(), currLocation.function.file.c_str(), currLocation.line);
			lastLocation = std::move(currLocation);
		}
#endif // ENABLE_RR_EMIT_PRINT_LOCATION
	}

	void DebugInfo::Flush()
	{
		emitPending(diScope.back(), builder);
	}

	void DebugInfo::syncScope(Backtrace const& backtrace)
	{
		auto shrink = [this](size_t newsize)
		{
			while (diScope.size() > newsize)
			{
				auto &scope = diScope.back();
				LOG("- STACK(%d): di: %p, location: %s:%d",
					int(diScope.size() - 1), scope.di,
					scope.location.function.file.c_str(),
					int(scope.location.line));
				emitPending(scope, builder);
				diScope.pop_back();
			}
		};

		if (backtrace.size() < diScope.size())
		{
			shrink(backtrace.size());
		}

		for (size_t i = 0; i < diScope.size(); i++)
		{
			auto &scope = diScope[i];
			auto const &oldLocation = scope.location;
			auto const &newLocation = backtrace[i];

			if (oldLocation.function != newLocation.function)
			{
				LOG("  STACK(%d): Changed function %s -> %s", int(i),
					oldLocation.function.name.c_str(), newLocation.function.name.c_str());
				shrink(i);
				break;
			}

			if (oldLocation.line > newLocation.line)
			{
				// Create a new di block to shadow all the variables in the loop.
				auto file = getOrCreateFile(newLocation.function.file.c_str());
				auto di = diBuilder->createLexicalBlock(scope.di, file, newLocation.line, 0);
				LOG("  STACK(%d): Jumped backwards %d -> %d. di: %p -> %p", int(i),
					oldLocation.line, newLocation.line, scope.di, di);
				emitPending(scope, builder);
				scope = {newLocation, di};
				shrink(i+1);
				break;
			}

			scope.location = newLocation;
		}

		while (backtrace.size() > diScope.size())
		{
			auto i = diScope.size();
			auto location = backtrace[i];
			auto file = getOrCreateFile(location.function.file.c_str());
			auto funcTy = diBuilder->createSubroutineType(diBuilder->getOrCreateTypeArray({}));

			char buf[1024];
			size_t size = sizeof(buf);
			int status = 0;
			llvm::itaniumDemangle(location.function.name.c_str(), buf, &size, &status);
			auto name = "jit!" + (status == 0 ? std::string(buf) : location.function.name);

			auto func = diBuilder->createFunction(
				file,                           // scope
				name,                           // function name
				"",                             // linkage
				file,                           // file
				location.line,                  // line
				funcTy,                         // type
				false,                          // internal linkage
				true,                           // definition
				location.line,                  // scope line
				llvm::DINode::FlagPrototyped,   // flags
				false                           // is optimized
			);
			diScope.push_back({location, func});
			LOG("+ STACK(%d): di: %p, location: %s:%d", int(i), di,
				location.function.file.c_str(), int(location.line));
		}
	}

	llvm::DILocation* DebugInfo::getLocation(const Backtrace &backtrace, size_t i)
	{
		if (backtrace.size() == 0) { return nullptr; }
		assert(backtrace.size() == diScope.size());
		return llvm::DILocation::get(
			*context,
			backtrace[i].line,
			0,
			diScope[i].di,
			i > 0 ? getLocation(backtrace, i - 1) : diRootLocation
		);
	}

	void DebugInfo::EmitVariable(Value *variable)
	{
		auto const& backtrace = getCallerBacktrace();
		syncScope(backtrace);

		for (int i = backtrace.size() - 1; i >= 0; i--)
		{
			auto const &location = backtrace[i];
			auto tokens = getOrParseFileTokens(location.function.file.c_str());
			auto tokIt = tokens->find(location.line);
			if (tokIt == tokens->end())
			{
				break;
			}
			auto token = tokIt->second;
			auto name = token.identifier;
			if (token.kind == Token::Return)
			{
				// This is a:
				//
				//   return <expr>;
				//
				// Emit this expression as two variables -
				// Once as a synthetic 'return_value' variable at this scope.
				// Again by bubbling the expression value up the callstack as
				// Return Value Optimizations (RVOs) are likely to carry across
				// the value to a local without calling a constructor in
				// statements like:
				//
				//   auto val = foo();
				//
				name = "return_value";
			}

			auto &scope = diScope[i];
			if (scope.pending.location != location)
			{
				emitPending(scope, builder);
			}

			auto value = V(variable);
			auto block = builder->GetInsertBlock();

			auto insertAfter = block->size() > 0 ? &block->back() : nullptr;
			while (insertAfter != nullptr && insertAfter->isTerminator())
			{
				insertAfter = insertAfter->getPrevNode();
			}

			scope.pending = Pending{};
			scope.pending.name = name;
			scope.pending.location = location;
			scope.pending.diLocation = getLocation(backtrace, i);
			scope.pending.value = value;
			scope.pending.block = block;
			scope.pending.insertAfter = insertAfter;
			scope.pending.scope = scope.di;

			if (token.kind == Token::Return)
			{
				// Insert a noop instruction so the debugger can inspect the
				// return value before the function scope closes.
				scope.pending.addNopOnNextLine = true;
			}
			else
			{
				break;
			}
		}
	}

	void DebugInfo::emitPending(Scope &scope, IRBuilder *builder)
	{
		auto const &pending = scope.pending;
		if (pending.value == nullptr)
		{
			return;
		}

		if (!scope.symbols.emplace(pending.name).second)
		{
			return;
		}

		bool isAlloca = llvm::isa<llvm::AllocaInst>(pending.value);

		LOG("  EMIT(%s): di: %p, location: %s:%d, isAlloca: %s", pending.name.c_str(), scope.di,
			pending.location.function.file.c_str(), pending.location.line, isAlloca ? "true" : "false");

		auto value = pending.value;

		IRBuilder::InsertPointGuard guard(*builder);
		if (pending.insertAfter != nullptr)
		{
			builder->SetInsertPoint(pending.block, ++pending.insertAfter->getIterator());
		}
		else
		{
			builder->SetInsertPoint(pending.block);
		}
		builder->SetCurrentDebugLocation(pending.diLocation);

		if (!isAlloca)
		{
			// While insertDbgValueIntrinsic should be enough to declare a
			// variable with no storage, variables of RValues can share the same
			// llvm::Value, and only one can be named. Take for example:
			//
			//   Int a = 42;
			//   RValue<Int> b = a;
			//   RValue<Int> c = b;
			//
			// To handle this, always promote named RValues to an alloca.

			llvm::BasicBlock &entryBlock = function->getEntryBlock();
			auto alloca = new llvm::AllocaInst(value->getType(), 0, pending.name);
			entryBlock.getInstList().push_front(alloca);
			builder->CreateStore(value, alloca);
			value = alloca;
		}

		value->setName(pending.name);

		auto diFile = getOrCreateFile(pending.location.function.file.c_str());
		auto diType = getOrCreateType(value->getType()->getPointerElementType());
		auto diVar = diBuilder->createAutoVariable(scope.di, pending.name, diFile, pending.location.line, diType);

		auto di = diBuilder->insertDeclare(value, diVar, diBuilder->createExpression(), pending.diLocation, pending.block);
		if (pending.insertAfter != nullptr) { di->moveAfter(pending.insertAfter); }

		if (pending.addNopOnNextLine)
		{
			builder->SetCurrentDebugLocation(llvm::DILocation::get(
				*context,
				pending.diLocation->getLine() + 1,
				0,
				pending.diLocation->getScope(),
				pending.diLocation->getInlinedAt()
			));
			Nop();
		}

		scope.pending = Pending{};
	}

	void DebugInfo::NotifyObjectEmitted(const llvm::object::ObjectFile &Obj, const llvm::LoadedObjectInfo &L)
	{
		std::unique_lock<std::mutex> lock(jitEventListenerMutex);
		jitEventListener->NotifyObjectEmitted(Obj, static_cast<const llvm::RuntimeDyld::LoadedObjectInfo&>(L));
	}

	void DebugInfo::NotifyFreeingObject(const llvm::object::ObjectFile &Obj)
	{
		std::unique_lock<std::mutex> lock(jitEventListenerMutex);
		jitEventListener->NotifyFreeingObject(Obj);
	}

	void DebugInfo::registerBasicTypes()
	{
		using namespace rr;
		using namespace llvm;

		auto vec4 = diBuilder->getOrCreateArray(diBuilder->getOrCreateSubrange(0, 4));
		auto vec8 = diBuilder->getOrCreateArray(diBuilder->getOrCreateSubrange(0, 8));
		auto vec16 = diBuilder->getOrCreateArray(diBuilder->getOrCreateSubrange(0, 16));

		diTypes.emplace(T(Bool::getType()), diBuilder->createBasicType("Bool", sizeof(bool), dwarf::DW_ATE_boolean));
		diTypes.emplace(T(Byte::getType()), diBuilder->createBasicType("Byte", 8, dwarf::DW_ATE_unsigned_char));
		diTypes.emplace(T(SByte::getType()), diBuilder->createBasicType("SByte", 8, dwarf::DW_ATE_signed_char));
		diTypes.emplace(T(Short::getType()), diBuilder->createBasicType("Short", 16, dwarf::DW_ATE_signed));
		diTypes.emplace(T(UShort::getType()), diBuilder->createBasicType("UShort", 16, dwarf::DW_ATE_unsigned));
		diTypes.emplace(T(Int::getType()), diBuilder->createBasicType("Int", 32, dwarf::DW_ATE_signed));
		diTypes.emplace(T(UInt::getType()), diBuilder->createBasicType("UInt", 32, dwarf::DW_ATE_unsigned));
		diTypes.emplace(T(Long::getType()), diBuilder->createBasicType("Long", 64, dwarf::DW_ATE_signed));
		diTypes.emplace(T(Half::getType()), diBuilder->createBasicType("Half", 16, dwarf::DW_ATE_float));
		diTypes.emplace(T(Float::getType()), diBuilder->createBasicType("Float", 32, dwarf::DW_ATE_float));

		diTypes.emplace(T(Byte4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Byte::getType())], {vec16}));
		diTypes.emplace(T(SByte4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(SByte::getType())], {vec16}));
		diTypes.emplace(T(Byte8::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Byte::getType())], {vec16}));
		diTypes.emplace(T(SByte8::getType()), diBuilder->createVectorType(128, 128, diTypes[T(SByte::getType())], {vec16}));
		diTypes.emplace(T(Byte16::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Byte::getType())], {vec16}));
		diTypes.emplace(T(SByte16::getType()), diBuilder->createVectorType(128, 128, diTypes[T(SByte::getType())], {vec16}));
		diTypes.emplace(T(Short2::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Short::getType())], {vec8}));
		diTypes.emplace(T(UShort2::getType()), diBuilder->createVectorType(128, 128, diTypes[T(UShort::getType())], {vec8}));
		diTypes.emplace(T(Short4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Short::getType())], {vec8}));
		diTypes.emplace(T(UShort4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(UShort::getType())], {vec8}));
		diTypes.emplace(T(Short8::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Short::getType())], {vec8}));
		diTypes.emplace(T(UShort8::getType()), diBuilder->createVectorType(128, 128, diTypes[T(UShort::getType())], {vec8}));
		diTypes.emplace(T(Int2::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Int::getType())], {vec4}));
		diTypes.emplace(T(UInt2::getType()), diBuilder->createVectorType(128, 128, diTypes[T(UInt::getType())], {vec4}));
		diTypes.emplace(T(Int4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Int::getType())], {vec4}));
		diTypes.emplace(T(UInt4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(UInt::getType())], {vec4}));
		diTypes.emplace(T(Float2::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Float::getType())], {vec4}));
		diTypes.emplace(T(Float4::getType()), diBuilder->createVectorType(128, 128, diTypes[T(Float::getType())], {vec4}));
	}

	DebugInfo::Location DebugInfo::getCallerLocation() const
	{
		return getCallerBacktrace(1)[0];
	}

	DebugInfo::Backtrace DebugInfo::getCallerBacktrace(size_t limit /* = 0 */) const
	{
		auto shouldSkipFile = [](llvm::StringRef fileSR) {
				return fileSR.empty() ||
					fileSR.endswith_lower("ReactorDebugInfo.cpp") ||
					fileSR.endswith_lower("Reactor.cpp") ||
					fileSR.endswith_lower("Reactor.hpp") ||
					fileSR.endswith_lower("stacktrace.hpp");
		};

		std::vector<DebugInfo::Location> locations;

		// Note that bs::stacktrace() effectively returns a vector of addresses; bs::frame construction is where
		// the heavy lifting is done: resolving the function name, file and line number.
		namespace bs = boost::stacktrace;
		for (bs::frame frame : bs::stacktrace())
		{
			if (shouldSkipFile(frame.source_file()))
			{
				continue;
			}

			DebugInfo::Location location;
			location.function.file = frame.source_file();
			location.function.name = frame.name();
			location.line = frame.source_line();
			locations.push_back(location);

			if (limit > 0 && locations.size() >= limit)
			{
				break;
			}
		}

		std::reverse(locations.begin(), locations.end());

		return locations;
	}

	llvm::DIType *DebugInfo::getOrCreateType(llvm::Type* type)
	{
		auto it = diTypes.find(type);
		if (it != diTypes.end()) { return it->second; }

		if(type->isPointerTy())
		{
			auto dbgTy = diBuilder->createPointerType(
				getOrCreateType(type->getPointerElementType()),
				sizeof(void*)*8, alignof(void*)*8);
			diTypes.emplace(type, dbgTy);
			return dbgTy;
		}
		llvm::errs() << "Unimplemented debug type: " << type << "\n";
		assert(false);
		return nullptr;
	}

	llvm::DIFile *DebugInfo::getOrCreateFile(const char* path)
	{
		auto it = diFiles.find(path);
		if (it != diFiles.end()) { return it->second; }
		auto dirAndName = splitPath(path);
		auto file = diBuilder->createFile(dirAndName.second, dirAndName.first);
		diFiles.emplace(path, file);
		return file;
	}

	DebugInfo::LineTokens const *DebugInfo::getOrParseFileTokens(const char* path)
	{
		static std::regex reLocalDecl(
			"^" // line start
			"\\s*" // initial whitespace
			"(?:For\\s*\\(\\s*)?" // optional 'For ('
			"((?:\\w+(?:<[^>]+>)?)(?:::\\w+(?:<[^>]+>)?)*)" // type (match group 1)
			"\\s+" // whitespace between type and name
			"(\\w+)" // identifier (match group 2)
			"\\s*" // whitespace after identifier
			"(\\[.*\\])?"); // optional array suffix (match group 3)

		auto it = fileTokens.find(path);
		if (it != fileTokens.end())
		{
			return it->second.get();
		}

		auto tokens = std::unique_ptr<LineTokens>(new LineTokens());

		std::ifstream file(path);
		std::string line;
		int lineCount = 0;
		while (std::getline(file, line))
		{
			lineCount++;
			std::smatch match;
			if (std::regex_search(line, match, reLocalDecl) && match.size() > 3)
			{
				bool isArray = match.str(3) != "";
				if (!isArray) // Cannot deal with C-arrays of values.
				{
					if (match.str(1) == "return")
					{
						(*tokens)[lineCount] = Token{Token::Return};
					}
					else
					{
						(*tokens)[lineCount] = Token{Token::Identifier, match.str(2)};
					}
				}
			}
		}

		auto out = tokens.get();
		fileTokens.emplace(path, std::move(tokens));
		return out;
	}

} // namespace rr

#endif // ENABLE_RR_DEBUG_INFO