/* Resp: A programming language
 * Copyright (C) 2008-2009 David Robillard <dave@drobilla.net>
 *
 * Resp is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 *
 * Resp 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 Affero General
 * Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with Resp.  If not, see <http://www.gnu.org/licenses/>.
 */

/** @file
 * @brief Compile to LLVM IR
 */

#define __STDC_LIMIT_MACROS    1
#define __STDC_CONSTANT_MACROS 1

#include <map>
#include <sstream>
#include <boost/format.hpp>
#include "llvm/Value.h"
#include "llvm/Analysis/Verifier.h"
#include "llvm/Assembly/AsmAnnotationWriter.h"
#include "llvm/DerivedTypes.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/JIT.h"
#include "llvm/Instructions.h"
#include "llvm/LLVMContext.h"
#include "llvm/Module.h"
#include "llvm/PassManager.h"
#include "llvm/Support/raw_os_ostream.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Target/TargetSelect.h"
#include "llvm/Transforms/Scalar.h"
#include "resp.hpp"

using namespace llvm;
using namespace std;
using boost::format;

/***************************************************************************
 * LLVM Engine                                                             *
 ***************************************************************************/

struct LLVMEngine : public Engine {
	LLVMEngine()
		: builder(context)
	{
		InitializeNativeTarget();
		module = new Module("resp", context);
		engine = EngineBuilder(module).create();
		opt    = new FunctionPassManager(module);

		// Set up optimiser pipeline
		const TargetData* target = engine->getTargetData();
		opt->add(new TargetData(*target));          // Register target arch
		opt->add(createInstructionCombiningPass()); // Simple optimizations
		opt->add(createReassociatePass());          // Reassociate expressions
		opt->add(createGVNPass());                  // Eliminate Common Subexpressions
		opt->add(createCFGSimplificationPass());    // Simplify control flow

		// Declare host provided allocation primitive
		std::vector<const Type*> argsT(1, Type::getInt32Ty(context)); // unsigned size
		FunctionType* funcT = FunctionType::get(PointerType::get(Type::getInt8Ty(context), 0), argsT, false);
		alloc = Function::Create(funcT, Function::ExternalLinkage,
				"resp_gc_allocate", module);
	}

	~LLVMEngine()
	{
		delete engine;
		delete opt;
	}

	inline Value*    llVal(CVal v)   { return static_cast<Value*>(v); }
	inline Function* llFunc(CFunc f) { return static_cast<Function*>(f); }

	const Type*
	llType(const AType* t)
	{
		if (t == NULL) {
			return NULL;
		} else if (t->kind == AType::PRIM) {
			if (t->head()->str() == "Nothing") return Type::getVoidTy(context);
			if (t->head()->str() == "Bool")    return Type::getInt1Ty(context);
			if (t->head()->str() == "Int")     return Type::getInt32Ty(context);
			if (t->head()->str() == "Float")   return Type::getFloatTy(context);
			if (t->head()->str() == "String")  return PointerType::get(Type::getInt8Ty(context), NULL);
			if (t->head()->str() == "Quote")   return PointerType::get(Type::getInt8Ty(context), NULL);
			if (t->head()->str() == "Lexeme")  return PointerType::get(Type::getInt8Ty(context), NULL);
			throw Error(t->loc, string("Unknown primitive type `") + t->str() + "'");
		} else if (t->kind == AType::EXPR && t->head()->str() == "Fn") {
			AType::const_iterator i     = t->begin();
			const ATuple*         protT = (*++i)->to<const ATuple*>();
			const AType*          retT  = (*++i)->as<const AType*>();
			if (!llType(retT))
				return NULL;

			vector<const Type*> cprot;
			FOREACHP(ATuple::const_iterator, i, protT) {
				const Type* lt = llType((*i)->to<const AType*>());
				if (!lt)
					return NULL;
				cprot.push_back(lt);
			}

			return PointerType::get(FunctionType::get(llType(retT), cprot, false), 0);
		} else if (t->kind == AType::EXPR && isupper(t->head()->str()[0])) {
			vector<const Type*> ctypes;
			for (AType::const_iterator i = t->begin() + 1; i != t->end(); ++i) {
				const Type* lt = llType((*i)->to<const AType*>());
				if (!lt)
					return NULL;
				ctypes.push_back(lt);
			}

			return PointerType::get(StructType::get(context, ctypes, false), 0);
		}
		return PointerType::get(Type::getInt8Ty(context), NULL);
	}

	CFunc startFunction(CEnv& cenv,
			const std::string& name, const AType* retT, const ATuple& argsT,
			const vector<string> argNames)
	{
		Function::LinkageTypes linkage = Function::ExternalLinkage;

		vector<const Type*> cprot;
		FOREACH(ATuple::const_iterator, i, argsT) {
			AType* at = (*i)->as<AType*>();
			THROW_IF(!llType(at), Cursor(), string("non-concrete parameter :: ")
				+ at->str())
			cprot.push_back(llType(at));
		}

		THROW_IF(!llType(retT), Cursor(),
				(format("return has non-concrete type `%1%'") % retT->str()).str());

		FunctionType* fT = FunctionType::get(llType(retT), cprot, false);
		Function*     f  = Function::Create(fT, linkage, name, module);

		// Note f->getName() may be different from name
		// however LLVM chooses to mangle is fine, we keep a pointer

		// Set argument names in generated code
		Function::arg_iterator a = f->arg_begin();
		if (!argNames.empty())
			for (vector<string>::const_iterator i = argNames.begin(); i != argNames.end(); ++a, ++i)
				a->setName(*i);

		BasicBlock* bb = BasicBlock::Create(context, "entry", f);
		builder.SetInsertPoint(bb);

		return f;
	}

	void finishFunction(CEnv& cenv, CFunc f, CVal ret) {
		builder.CreateRet(llVal(ret));
		if (verifyFunction(*static_cast<Function*>(f), llvm::PrintMessageAction)) {
			module->dump();
			throw Error(Cursor(), "Broken module");
		}
		if (cenv.args.find("-g") == cenv.args.end())
			opt->run(*static_cast<Function*>(f));
	}

	void eraseFunction(CEnv& cenv, CFunc f) {
		if (f)
			llFunc(f)->eraseFromParent();
	}

	CVal compileCall(CEnv& cenv, CFunc f, const AType* funcT, const vector<CVal>& args) {
		vector<Value*> llArgs(*reinterpret_cast<const vector<Value*>*>(&args));
		Value* closure = builder.CreateBitCast(llArgs[0],
				llType(funcT->prot()->head()->as<const AType*>()),
				cenv.penv.gensymstr("you"));
		llArgs[0] = closure;
		return builder.CreateCall(llFunc(f), llArgs.begin(), llArgs.end());
	}

	CFunc compileFunction(CEnv& cenv, AFn* fn, const AType* type);

	CVal compileTup(CEnv& cenv, const AType* type, const vector<CVal>& fields);
	CVal compileDot(CEnv& cenv, CVal tup, int32_t index);
	CVal compileLiteral(CEnv& cenv, AST* lit);
	CVal compileString(CEnv& cenv, const char* str);
	CVal compilePrimitive(CEnv& cenv, APrimitive* prim);
	CVal compileIf(CEnv& cenv, AIf* aif);
	CVal compileGlobal(CEnv& cenv, const AType* type, const string& sym, CVal val);
	CVal getGlobal(CEnv& cenv, const string& sym, CVal val);

	void writeModule(CEnv& cenv, std::ostream& os) {
		AssemblyAnnotationWriter writer;
		llvm::raw_os_ostream raw_stream(os);
		module->print(raw_stream, &writer);
	}

	const string call(CEnv& cenv, CFunc f, const AType* retT) {
		void*       fp = engine->getPointerToFunction(llFunc(f));
		const Type* t  = llType(retT);
		THROW_IF(!fp, Cursor(), "unable to get function pointer");
		THROW_IF(!t,  Cursor(), "function with non-concrete return type called");

		std::stringstream ss;
		if (t == Type::getInt32Ty(context)) {
			ss << ((int32_t (*)())fp)();
		} else if (t == Type::getFloatTy(context)) {
			ss << showpoint << ((float (*)())fp)();
		} else if (t == Type::getInt1Ty(context)) {
			ss << (((bool (*)())fp)() ? "#t" : "#f");
		} else if (retT->head()->str() == "String") {
			const std::string s(((char* (*)())fp)());
			ss << "\"";
			for (std::string::const_iterator i = s.begin(); i != s.end(); ++i) {
				switch (*i) {
				case '\"':
				case '\\':
					ss << '\\';
				default:
					ss << *i;
					break;
				}
			}
			ss << "\"";
		} else if (retT->head()->str() == "Lexeme") {
			ss << ((char* (*)())fp)();
		} else if (retT->head()->str() == "Quote") {
			ss << "(quote " << ((char* (*)())fp)() << ")";
		} else if (t != Type::getVoidTy(context)) {
			ss << ((void* (*)())fp)();
		} else {
			((void (*)())fp)();
		}
		return ss.str();
	}

	LLVMContext             context;
	Module*                 module;
	ExecutionEngine*        engine;
	IRBuilder<>             builder;
	Function*               alloc;
	FunctionPassManager*    opt;
};

Engine*
resp_new_llvm_engine()
{
	return new LLVMEngine();
}

/***************************************************************************
 * Code Generation                                                         *
 ***************************************************************************/

/** Convert a size in bits to bytes, rounding up as necessary */
static inline size_t
bitsToBytes(size_t bits)
{
	return ((bits % 8 == 0) ? bits : (((bits / 8) + 1) * 8)) / 8;
}

CVal
LLVMEngine::compileTup(CEnv& cenv, const AType* type, const vector<CVal>& fields)
{
	// Find size of memory required
	size_t s = 0;
	assert(type->begin() != type->end());
	for (AType::const_iterator i = type->begin() + 1; i != type->end(); ++i)
		s += engine->getTargetData()->getTypeSizeInBits(llType((*i)->as<AType*>()));

	// Allocate struct
	Value* structSize = ConstantInt::get(Type::getInt32Ty(context), bitsToBytes(s));
	Value* mem        = builder.CreateCall(alloc, structSize, "tupMem");
	Value* structPtr  = builder.CreateBitCast(mem, llType(type), "tup");

	// Set struct fields
	size_t i = 0;
	for (vector<CVal>::const_iterator f = fields.begin(); f != fields.end(); ++f, ++i) {
		builder.CreateStore(llVal(*f),
				builder.CreateStructGEP(structPtr, i, (format("tup%1%") % i).str().c_str()));
	}

	return structPtr;
}

CVal
LLVMEngine::compileDot(CEnv& cenv, CVal tup, int32_t index)
{
	Value* ptr = builder.CreateStructGEP(llVal(tup), index, "dotPtr");
	return builder.CreateLoad(ptr, 0, "dotVal");
}

CVal
LLVMEngine::compileLiteral(CEnv& cenv, AST* lit)
{
	ALiteral<int32_t>* ilit = dynamic_cast<ALiteral<int32_t>*>(lit);
	if (ilit)
		return ConstantInt::get(Type::getInt32Ty(context), ilit->val, true);

	ALiteral<float>* flit = dynamic_cast<ALiteral<float>*>(lit);
	if (flit)
		return ConstantFP::get(Type::getFloatTy(context), flit->val);

	ALiteral<bool>* blit = dynamic_cast<ALiteral<bool>*>(lit);
	if (blit)
		return ConstantFP::get(Type::getFloatTy(context), blit->val);

	throw Error(lit->loc, "Unknown literal type");
}

CVal
LLVMEngine::compileString(CEnv& cenv, const char* str)
{
	return builder.CreateGlobalStringPtr(str);
}

CFunc
LLVMEngine::compileFunction(CEnv& cenv, AFn* fn, const AType* type)
{
	assert(type->concrete());

	const AType* argsT = type->prot()->as<const AType*>();
	const AType* retT  = type->last()->as<const AType*>();

	vector<string> argNames;
	for (ATuple::const_iterator i = fn->prot()->begin(); i != fn->prot()->end(); ++i)
		argNames.push_back((*i)->str());

	// Write function declaration
	const string name = (fn->name == "") ? cenv.penv.gensymstr("_fn") : fn->name;
	Function* f = llFunc(cenv.engine()->startFunction(cenv, name, retT, *argsT, argNames));

	cenv.push();

	// Bind argument values in CEnv
	vector<Value*> args;
	AFn::const_iterator    p  = fn->prot()->begin();
	ATuple::const_iterator pT = argsT->begin();
	assert(fn->prot()->size() == argsT->size());
	assert(fn->prot()->size() == f->num_args());
	for (Function::arg_iterator a = f->arg_begin(); a != f->arg_end(); ++a, ++p, ++pT) {
		const AType* t  = (*pT)->as<const AType*>();
		const Type*  lt = llType(t);
		THROW_IF(!lt, fn->loc, "untyped parameter\n");
		cenv.def((*p)->as<ASymbol*>(), *p, t, &*a);
	}

	assert(!cenv.currentFn);

	// Write function body
	try {
		cenv.currentFn = f;
		CVal retVal = NULL;
		for (AFn::iterator i = fn->begin() + 2; i != fn->end(); ++i)
			retVal = (*i)->compile(cenv);
		cenv.engine()->finishFunction(cenv, f, retVal);
	} catch (Error& e) {
		f->eraseFromParent(); // Error reading body, remove function
		cenv.pop();
		cenv.currentFn = NULL;
		throw e;
	}
	cenv.pop();
	cenv.currentFn = NULL;
	return f;
}

CVal
LLVMEngine::compileIf(CEnv& cenv, AIf* aif)
{
	typedef vector< pair<Value*, BasicBlock*> > Branches;
	LLVMEngine* engine  = reinterpret_cast<LLVMEngine*>(cenv.engine());
	Function*   parent  = engine->builder.GetInsertBlock()->getParent();
	BasicBlock* mergeBB = BasicBlock::Create(context, "endif");
	BasicBlock* nextBB  = NULL;
	Branches    branches;
	size_t idx = 1;
	for (AIf::iterator i = aif->begin() + 1; ; ++i, idx += 2) {
		AIf::iterator next = i;
		if (++next == aif->end())
			break;

		Value*      condV  = llVal((*i)->compile(cenv));
		BasicBlock* thenBB = BasicBlock::Create(context, (format("then%1%") % ((idx+1)/2)).str());

		nextBB = BasicBlock::Create(context, (format("else%1%") % ((idx+1)/2)).str());

		engine->builder.CreateCondBr(condV, thenBB, nextBB);

		// Emit then block for this condition
		parent->getBasicBlockList().push_back(thenBB);
		engine->builder.SetInsertPoint(thenBB);
		Value* thenV = llVal((*next)->compile(cenv));
		engine->builder.CreateBr(mergeBB);
		branches.push_back(make_pair(thenV, engine->builder.GetInsertBlock()));

		parent->getBasicBlockList().push_back(nextBB);
		engine->builder.SetInsertPoint(nextBB);

		i = next; // jump 2 each iteration (to the next predicate)
	}

	// Emit final else block
	engine->builder.SetInsertPoint(nextBB);
	Value* elseV = llVal(aif->last()->compile(cenv));
	engine->builder.CreateBr(mergeBB);
	branches.push_back(make_pair(elseV, engine->builder.GetInsertBlock()));

	// Emit merge block (Phi node)
	parent->getBasicBlockList().push_back(mergeBB);
	engine->builder.SetInsertPoint(mergeBB);
	PHINode* pn = engine->builder.CreatePHI(llType(cenv.type(aif)), "ifval");

	FOREACH(Branches::iterator, i, branches)
		pn->addIncoming(i->first, i->second);

	return pn;
}

CVal
LLVMEngine::compilePrimitive(CEnv& cenv, APrimitive* prim)
{
	APrimitive::iterator i = prim->begin();

	LLVMEngine*  engine  = reinterpret_cast<LLVMEngine*>(cenv.engine());
	bool         isFloat = cenv.type(*++i)->str() == "Float";
	Value*       a       = llVal((*i++)->compile(cenv));
	Value*       b       = llVal((*i++)->compile(cenv));
	const string n       = prim->head()->to<ASymbol*>()->str();

	// Binary arithmetic operations
	Instruction::BinaryOps op = (Instruction::BinaryOps)0;
	if (n == "+")   op = Instruction::Add;
	if (n == "-")   op = Instruction::Sub;
	if (n == "*")   op = Instruction::Mul;
	if (n == "and") op = Instruction::And;
	if (n == "or")  op = Instruction::Or;
	if (n == "xor") op = Instruction::Xor;
	if (n == "/")   op = isFloat ? Instruction::FDiv : Instruction::SDiv;
	if (n == "%")   op = isFloat ? Instruction::FRem : Instruction::SRem;
	if (op != 0) {
		Value* val = engine->builder.CreateBinOp(op, a, b);
		while (i != prim->end())
			val = engine->builder.CreateBinOp(op, val, llVal((*i++)->compile(cenv)));
		return val;
	}

	// Comparison operations
	CmpInst::Predicate pred = (CmpInst::Predicate)0;
	if (n == "=")  pred = isFloat ? CmpInst::FCMP_OEQ : CmpInst::ICMP_EQ ;
	if (n == "!=") pred = isFloat ? CmpInst::FCMP_ONE : CmpInst::ICMP_NE ;
	if (n == ">")  pred = isFloat ? CmpInst::FCMP_OGT : CmpInst::ICMP_SGT;
	if (n == ">=") pred = isFloat ? CmpInst::FCMP_OGE : CmpInst::ICMP_SGE;
	if (n == "<")  pred = isFloat ? CmpInst::FCMP_OLT : CmpInst::ICMP_SLT;
	if (n == "<=") pred = isFloat ? CmpInst::FCMP_OLE : CmpInst::ICMP_SLE;
	if (pred != 0) {
		if (isFloat)
			return engine->builder.CreateFCmp(pred, a, b);
		else
			return engine->builder.CreateICmp(pred, a, b);
	}

	throw Error(prim->loc, "unknown primitive");
}

CVal
LLVMEngine::compileGlobal(CEnv& cenv, const AType* type, const string& sym, CVal val)
{
	LLVMEngine*     engine = reinterpret_cast<LLVMEngine*>(cenv.engine());
	GlobalVariable* global = new GlobalVariable(*module, llType(type), false,
			GlobalValue::ExternalLinkage, Constant::getNullValue(llType(type)), sym);

	engine->builder.CreateStore(llVal(val), global);
	return global;
}

CVal
LLVMEngine::getGlobal(CEnv& cenv, const string& sym, CVal val)
{
	LLVMEngine* engine = reinterpret_cast<LLVMEngine*>(cenv.engine());
	return engine->builder.CreateLoad(llVal(val), sym + "Ptr");
}