#include <iostream>
#include <string>
#include <vector>
#include <map>
using namespace std;

#include "llvm/Constants.h"
#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/Instruction.h"
#include "llvm/Instructions.h"
#include "llvm/Module.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Value.h"
#include "llvm/ValueSymbolTable.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Type.h"
#include <llvm/LLVMContext.h>
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Bitcode/ReaderWriter.h>
#include "llvm/Support/system_error.h"
#include <llvm/ADT/ArrayRef.h>
#include "llvm/Support/TypeBuilder.h"
using namespace llvm;

#include <assert.h>

struct Hello: public ModulePass{	
public:
	static char ID;

	Hello() : ModulePass(ID) {
		my_section = new StringRef("my_section");	
	}

	virtual bool runOnModule(Module &M) {
		/* Run "FunctionPass" */
		bool retval = false;
		for(Module::iterator f = M.begin(); f != M.end(); ++f) 
			retval |= runOnFunction(*f, M);

		//M.dump();
		return retval;
	}

private:
	StringRef *my_section;
	map<Value *, GlobalVariable *> my_shadowMap;
	map<Function *, pair<Value *, Value *> > my_ptrAssigns;

	Function *my_createAssignHook(Module &M, Type *ptrType) {
		/* Create arguments for new function */
		vector<Type *> v_params;
		v_params.push_back(ptrType);					/* Shadow */
		v_params.push_back(ptrType);					/* New Value */
		v_params.push_back(ptrType);					/* Pointer */
		//v_params.push_back(Type::getInt32PtrTy(M.getContext()));	/* Function */
		ArrayRef<Type *> a_params(v_params);

		/* Create new function */
		FunctionType *my_funcType = FunctionType::get(Type::getVoidTy(M.getContext()), a_params, false);
		Function *newFunc = Function::Create(my_funcType, GlobalValue::PrivateLinkage, "", &M);
		newFunc->setSection(*my_section);

		/* Start Code */
		BasicBlock *block = BasicBlock::Create(M.getContext(), "", newFunc);
		
		Value *val = newFunc->arg_begin();
		Value *ptr = ++(newFunc->arg_begin());
		new StoreInst(val, ptr, block);

		/* End Code */
		ReturnInst::Create(M.getContext(), block);
		return newFunc;
	}

	bool runOnFunction(Function &F, Module &M) {
		bool retval = false; 
		for(inst_iterator i = inst_begin(F); i != inst_end(F); ++i) {
			StoreInst *store = dyn_cast<StoreInst>(&*i);
			if(store != NULL) {
				if(store->getValueOperand()->getType()->isPointerTy()) {
					PointerType *pointerType = cast<PointerType>(store->getValueOperand()->getType()); 
					if(pointerType->getTypeAtIndex(0u)->isFunctionTy() == true) {
						/* Store Assignment */
						pair<Value *, Value *> assign = make_pair(store->getPointerOperand(), store->getValueOperand());
						my_ptrAssigns.insert(make_pair(&F, assign));  //TODO: mem leak?

						/* Create shadow copy of variable, if needed */
						GlobalVariable *shadow;

						map<Value *, GlobalVariable *>::iterator it = my_shadowMap.find(store->getPointerOperand());
						if(it == my_shadowMap.end()) {
							shadow = new GlobalVariable(M, pointerType, false, GlobalValue::PrivateLinkage, 
										    Constant::getNullValue(pointerType), "");
							shadow->setSection(*my_section);
				
							pair<map<Value *, GlobalVariable *>::iterator, bool> result;
							my_shadowMap.insert(make_pair(store->getPointerOperand(), shadow));
						} else {
							shadow = it->second;
						}

						/* Make arguments */
						vector<Value *> v_args;
						v_args.push_back(shadow);
						v_args.push_back(store->getValueOperand());
						v_args.push_back(store->getPointerOperand());
						//v_args.push_back(&F);
						ArrayRef<Value *> a_args(v_args);
						
						/* Insert Hook */
						++i;
						Function *hook = my_createAssignHook(M, store->getValueOperand()->getType());
						M.dump();
						CallInst::Create(hook, a_args, "", &*i);
						//M.dump();
						--i;
						retval = true;
					}
				}
			}
		}
		//M.dump();
		return retval; 
	}
};

char Hello::ID = 0;
static RegisterPass<Hello> X("hello", "Hello World Pass");