// Test of multiple JIT Based on HowToUseJIT example

//printf symbol
#include <stdio.h>

#include "llvm/LLVMContext.h"
#include "llvm/Module.h"
#include "llvm/Constants.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Instructions.h"
#include "llvm/ExecutionEngine/JIT.h"
#include "llvm/ExecutionEngine/Interpreter.h"
#include "llvm/ExecutionEngine/GenericValue.h"
#include "llvm/Target/TargetSelect.h"
#include "llvm/Support/ManagedStatic.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;

void createModule1(LLVMContext &Context1, Module * &M1, Function * &FooF1) {
  M1 = new Module("test1", Context1);
  // function add1
  Function *Add1F1 =
    cast<Function>(M1->getOrInsertFunction("add1", Type::getInt32Ty(Context1),
                                          Type::getInt32Ty(Context1),
                                          (Type *)0));
  BasicBlock *BB1 = BasicBlock::Create(Context1, "EntryBlock", Add1F1);
  Value *One1 = ConstantInt::get(Type::getInt32Ty(Context1), 1);
  Argument *ArgX1 = Add1F1->arg_begin(); 
  ArgX1->setName("ArgX1");
  Instruction *Add1 = BinaryOperator::CreateAdd(One1, ArgX1, "addresult", BB1);
  ReturnInst::Create(Context1, Add1, BB1);
  // `foo'
  FooF1 =
    cast<Function>(M1->getOrInsertFunction("foo1", Type::getInt32Ty(Context1),
                                          (Type *)0));
  BB1 = BasicBlock::Create(Context1, "EntryBlock", FooF1);
  Value *Ten1 = ConstantInt::get(Type::getInt32Ty(Context1), 10);
  CallInst *Add1CallRes1 = CallInst::Create(Add1F1, Ten1, "add1", BB1);
  ReturnInst::Create(Context1, Add1CallRes1, BB1);
}

void createModule2(LLVMContext &Context2, Module * &M2, Function * &FooF2) {
  M2 = new Module("test2", Context2);
  // function add2
  Function *Add2F2 =
    cast<Function>(M2->getOrInsertFunction("add2", Type::getInt32Ty(Context2),
                                          Type::getInt32Ty(Context2),
                                          (Type *)0));
  BasicBlock *BB2 = BasicBlock::Create(Context2, "EntryBlock", Add2F2);
  Value *Two2 = ConstantInt::get(Type::getInt32Ty(Context2), 2);
  Argument *ArgX2 = Add2F2->arg_begin();
  ArgX2->setName("ArgX2");
  Instruction *Add2 = BinaryOperator::CreateAdd(Two2, ArgX2, "addresult", BB2);
  ReturnInst::Create(Context2, Add2, BB2);
  //foo
  FooF2 =
    cast<Function>(M2->getOrInsertFunction("foo2", Type::getInt32Ty(Context2),
                                          (Type *)0));
  BB2 = BasicBlock::Create(Context2, "EntryBlock", FooF2);
  Value *Ten2 = ConstantInt::get(Type::getInt32Ty(Context2), 10);
  CallInst *Add1CallRes2 = CallInst::Create(Add2F2, Ten2, "add2", BB2);
  ReturnInst::Create(Context2, Add1CallRes2, BB2);
}

void testEagerMode() {
  LLVMContext Context1;
  Module *M1 = 0;
  Function *FooF1 = 0;
  createModule1(Context1, M1, FooF1);

  LLVMContext Context2;
  Module *M2 = 0;
  Function *FooF2 = 0;
  createModule2(Context2, M2, FooF2);
  
  // Now we create the JIT in eager mode
  outs() << "JIT in eager mode:\n\n";
  ExecutionEngine* EE1 = EngineBuilder(M1).create();
  EE1->DisableLazyCompilation(true);
  ExecutionEngine* EE2 = EngineBuilder(M2).create();
  EE2->DisableLazyCompilation(true);

  // Call the `foo' function with no arguments:
  std::vector<GenericValue> noargs;
  GenericValue gv1 = EE1->runFunction(FooF1, noargs);
  GenericValue gv2 = EE2->runFunction(FooF2, noargs);

  // Import result of execution:
  if (gv1.IntVal != 11) {
	  errs() << "gv1 is not 11 but " << gv1.IntVal << "\n";
	  return;
  }
  if (gv2.IntVal != 12) {
	  errs() << "gv2 is not 12 but " << gv2.IntVal << "\n";
	  return;
  }
  EE1->freeMachineCodeForFunction(FooF1);
  EE2->freeMachineCodeForFunction(FooF2);
  delete EE1;
  delete EE2;
}

void testLazyMode() {
  LLVMContext Context1;
  Module *M1 = 0;
  Function *FooF1 = 0;
  createModule1(Context1, M1, FooF1);

  LLVMContext Context2;
  Module *M2 = 0;
  Function *FooF2 = 0;
  createModule2(Context2, M2, FooF2);

  // Now we create the JIT in lazy mode
  outs() << "JIT in lazy mode:\n\n";
  ExecutionEngine* EE1 = EngineBuilder(M1).create();
  EE1->DisableLazyCompilation(false);
  ExecutionEngine* EE2 = EngineBuilder(M2).create();
  EE2->DisableLazyCompilation(false);

  // Call the `foo' function with no arguments:
  std::vector<GenericValue> noargs;
  GenericValue gv1 = EE1->runFunction(FooF1, noargs);
  GenericValue gv2 = EE2->runFunction(FooF2, noargs);

  // Import result of execution:
  if (gv1.IntVal != 11) {
	  errs() << "gv1 is not 11 but " << gv1.IntVal << "\n";
	  return;
  }
  if (gv2.IntVal != 12) {
	  errs() << "gv2 is not 12 but " << gv2.IntVal << "\n";
	  return;
  }
  EE1->freeMachineCodeForFunction(FooF1);
  EE2->freeMachineCodeForFunction(FooF2);
  delete EE1;
  delete EE2;
}

extern "C" {
  extern void *getPointerToNamedFunction(const char *Name);
}

void testJitPool() {
  LLVMContext Context1;
  Module *M1 = 0;
  Function *FooF1 = 0;
  createModule1(Context1, M1, FooF1);

  LLVMContext Context2;
  Module *M2 = 0;
  Function *FooF2 = 0;
  createModule2(Context2, M2, FooF2);

  // Now we create the JIT
  outs() << "run test on JitPool\n"; 
  ExecutionEngine* EE1 = EngineBuilder(M1).create();
  ExecutionEngine* EE2 = EngineBuilder(M2).create();

  Function *F1 = EE1->FindFunctionNamed("foo1");
  void *foo1 = EE1->getPointerToFunction(F1);

  Function *F2 = EE2->FindFunctionNamed("foo2");
  void *foo2 = EE2->getPointerToFunction(F2);

  //function in M1
  if (getPointerToNamedFunction("foo1") != foo1) {
	  errs() << "fail to correctly resolve foo1 \n";
    return;
  }
  //function in M2
  if (getPointerToNamedFunction("foo2") != foo2) {
	  errs() << "fail to correctly resolve foo2 \n";
    return;
  }
  //Symbol search
  if (getPointerToNamedFunction("printf") != (void*)&printf) {
	  errs() << "fail to correctly resolve testEagerMode : " << getPointerToNamedFunction("testEagerMode") << "\n";
    return;
  }

  delete EE1;
  delete EE2;
}

int main() {
  
  InitializeNativeTarget();

  testEagerMode();
  testLazyMode();
  testJitPool();

  llvm_shutdown();
  return 0;
}