/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * File:   control.cpp
 * Author: pgess <v.melnychenko@xreate.org>
 *
 * Created on June 26, 2016, 6:00 PM
 */

#include "analysis/typeinference.h"
#include "compilation/control.h"
#include "compilation/containers.h"
#include "compilation/transformersaturation.h"

#include "query/containers.h"
#include "llvmlayer.h"
#include "ast.h"

using namespace std;
using namespace llvm;
using namespace xreate;
using namespace xreate::containers;
using namespace xreate::compilation;

#define NAME(x) (hintRetVar.empty()? x : hintRetVar)
#define UNUSED(x) (void)(x)
#define EXPAND_CONTEXT         \
    LLVMLayer* llvm = context.pass->man->llvm; \
    compilation::IBruteScope* scope = context.scope; \
    compilation::IBruteFunction* function = context.function;

ControlIR::ControlIR(compilation::Context ctx)
: context(ctx), tyNum(static_cast<llvm::IntegerType*> (ctx.pass->man->llvm->toLLVMType(ExpandedType(TypeAnnotation(TypePrimitive::Int))))) {
}

Value*
ControlIR::compileStructIndex(llvm::Value* aggregate, ExpandedType aggrT, const list<string>& indices) {
  EXPAND_CONTEXT UNUSED(scope); UNUSED(function);
  TypesHelper types(llvm);
  llvm::Value* result = aggregate;
  assert(indices.size());

  for (const string& indexField: indices){
    const std::vector<std::string>& tyfields = types.getRecordFields(aggrT);
    unsigned idx = -1; bool flagFound = false;
    for (unsigned i = 0, size = tyfields.size(); i < size; ++i) {
      if (tyfields.at(i) == indexField) {
        idx = i; flagFound = true; break;
      }
    }

    if (flagFound){
      result = llvm->irBuilder.CreateExtractValue(result, llvm::ArrayRef<unsigned>{idx});
      aggrT = typeinference::getSubtype(aggrT, indexField);
    } else {
      assert(false && "not found required struct field");
    }
  }

  return result;

//dereference pointer
//if (types.isPointerT(t)) {
//          llvm::Value* addr = llvm->irBuilder.CreateConstGEP2_32(nullptr, aggregate, 0, i);
//          return llvm->irBuilder.CreateLoad(addr);
//}
}

llvm::Value*
ControlIR::compileFold(const Expression& loopE, const std::string& hintAlias) {
  EXPAND_CONTEXT
  assert(loopE.op == Operator::FOLD);
  AST* ast = context.pass->man->root;


  //initialization:
  const Expression aggrE = loopE.getOperands().at(0);
  const ExpandedType& aggrT = ast->getType(aggrE);
  llvm::Value* aggrRaw = context.scope->process(aggrE);

  IFwdIteratorIR* aggrItIR = IFwdIteratorIR::create(aggrE, aggrT, context);
  llvm::Value* idxBeginRaw = aggrItIR->begin(aggrRaw);
  llvm::Value* idxEndRaw = aggrItIR->end(aggrRaw);
  ExpandedType loopT = ast->getType(loopE);
  std::string elAlias = loopE.bindings[0];
  std::string accumAlias = loopE.bindings[1];
  const Expression& accumE = loopE.getOperands().at(1);
  ExpandedType accumT = ast->getType(accumE, loopT.get());
  llvm::Type* accumRawT = llvm->toLLVMType(accumT);
  llvm::Value* accumInitRaw = scope->process(accumE, accumAlias, accumT.get());

  llvm::BasicBlock *blockProlog = llvm::BasicBlock::Create(llvm->llvmContext, "fold_prlg", function->raw);
  llvm::BasicBlock *blockHeader = llvm::BasicBlock::Create(llvm->llvmContext, "fold_hdr", function->raw);
  llvm::BasicBlock *blockBody = llvm::BasicBlock::Create(llvm->llvmContext, "fold", function->raw);
  llvm::BasicBlock *blockFooter = llvm::BasicBlock::Create(llvm->llvmContext, "fold_ftr", function->raw);
  llvm::BasicBlock *blockEpilog = llvm::BasicBlock::Create(llvm->llvmContext, "fold_eplg", function->raw);

  std::unique_ptr<TransformerSaturation> transformerSaturation(new TransformerSaturation(blockProlog, context.pass->managerTransformations));

  //Header:
  // * create phi
  llvm->irBuilder.SetInsertPoint(blockHeader);
  llvm::PHINode *accum = llvm->irBuilder.CreatePHI(accumRawT, 2, accumAlias);
  accum->addIncoming(accumInitRaw, blockProlog);
  llvm::PHINode *idxCurrentRaw = llvm->irBuilder.CreatePHI(idxBeginRaw->getType(), 2, "foldIt");
  idxCurrentRaw->addIncoming(idxBeginRaw, blockProlog);

  // * loop checks
  Value* condRange = llvm->irBuilder.CreateICmpNE(idxCurrentRaw, idxEndRaw);
  llvm->irBuilder.CreateCondBr(condRange, blockBody, blockEpilog);

  //Body:
  llvm->irBuilder.SetInsertPoint(blockBody);
  CodeScope* scopeLoop = loopE.blocks.front();
  compilation::IBruteScope* loopUnit = function->getBruteScope(scopeLoop);
  Value* elIn = aggrItIR->get(aggrRaw, idxCurrentRaw);
  loopUnit->bindArg(accum, move(accumAlias));
  loopUnit->bindArg(elIn, move(elAlias));
  Value* accumNext = loopUnit->compile();

  // * Loop saturation checks
  bool flagSaturationTriggered = transformerSaturation->insertSaturationChecks(blockFooter, blockEpilog, context);
  llvm::BasicBlock* blockSaturation = llvm->irBuilder.GetInsertBlock();
  if (!flagSaturationTriggered){
      llvm->irBuilder.CreateBr(blockFooter);
  }

  //Footer:
  // * computing next iteration state
  llvm->irBuilder.SetInsertPoint(blockFooter);
  Value *itLoopNext = aggrItIR->advance(idxCurrentRaw);
  accum->addIncoming(accumNext, llvm->irBuilder.GetInsertBlock());
  idxCurrentRaw->addIncoming(itLoopNext, llvm->irBuilder.GetInsertBlock());
  llvm->irBuilder.CreateBr(blockHeader);

  //Prolog:
  llvm->irBuilder.SetInsertPoint(context.scope->lastBlockRaw);
  llvm->irBuilder.CreateBr(blockProlog);

  llvm->irBuilder.SetInsertPoint(blockProlog);
  llvm->irBuilder.CreateBr(blockHeader);

  // Epilog:
  llvm->irBuilder.SetInsertPoint(blockEpilog);
  if (!flagSaturationTriggered){
      return accum;
  }

  llvm::PHINode* result = llvm->irBuilder.CreatePHI(accumRawT, 2, hintAlias);
  result->addIncoming(accum, blockHeader);
  result->addIncoming(accumNext, blockSaturation);

  return result;
}

llvm::Value*
ControlIR::compileFoldInf(const Expression& fold, const std::string& hintRetVar) {
    EXPAND_CONTEXT
    assert(fold.op == Operator::FOLD_INF);

    std::string accumName = fold.bindings[0];
    llvm::Value* accumInit = scope->process(fold.getOperands()[0]);

    llvm::BasicBlock *blockBeforeLoop = llvm->irBuilder.GetInsertBlock();
    llvm::BasicBlock *blockLoop = llvm::BasicBlock::Create(llvm->llvmContext, "foldinf", function->raw);
    llvm::BasicBlock *blockNext = llvm::BasicBlock::Create(llvm->llvmContext, "foldinf_next", function->raw);
    llvm::BasicBlock *blockAfterLoop = llvm::BasicBlock::Create(llvm->llvmContext, "foldinf_post", function->raw);
    std::unique_ptr<TransformerSaturation> transformerSaturation(new TransformerSaturation(blockBeforeLoop, context.pass->managerTransformations));

    llvm->irBuilder.CreateBr(blockLoop);

    // * create phi
    llvm->irBuilder.SetInsertPoint(blockLoop);
    llvm::PHINode *accum = llvm->irBuilder.CreatePHI(accumInit->getType(), 2, accumName);
    accum->addIncoming(accumInit, blockBeforeLoop);

    // * loop body
    CodeScope* scopeLoop = fold.blocks.front();
    compilation::IBruteScope* unitLoop = function->getBruteScope(scopeLoop);
    unitLoop->bindArg(accum, move(accumName));
    Value* accumNext = unitLoop->compile();

    // * Loop saturation checks
    bool flagSaturationTriggered = transformerSaturation->insertSaturationChecks(blockNext, blockAfterLoop, context);
    assert(flagSaturationTriggered);

    // * computing next iteration state
    llvm->irBuilder.SetInsertPoint(blockNext);
    accum->addIncoming(accumNext, llvm->irBuilder.GetInsertBlock());
    llvm->irBuilder.CreateBr(blockLoop);

    // finalization:
    llvm->irBuilder.SetInsertPoint(blockAfterLoop);
    return accumNext;
}

llvm::Value*
ControlIR::compileIf(const Expression& exprIf, const std::string& hintRetVar) {
    EXPAND_CONTEXT

    const Expression& condExpr = exprIf.getOperands()[0];
    llvm::IRBuilder<>& builder = llvm->irBuilder;
    assert(builder.GetInsertBlock() == scope->lastBlockRaw);

            //initialization:
    llvm::BasicBlock *blockEpilog = llvm::BasicBlock::Create(llvm->llvmContext, "ifAfter", function->raw);
    llvm::BasicBlock *blockTrue = llvm::BasicBlock::Create(llvm->llvmContext, "ifTrue", function->raw);
    llvm::BasicBlock *blockFalse = llvm::BasicBlock::Create(llvm->llvmContext, "ifFalse", function->raw);

    llvm::Value* cond = scope->process(condExpr);

    builder.SetInsertPoint(blockTrue);
    CodeScope* scopeTrue = exprIf.blocks.front();
    llvm::Value* resultTrue = function->getBruteScope(scopeTrue)->compile();
    llvm::BasicBlock * blockTrueEnd = builder.GetInsertBlock();
    builder.CreateBr(blockEpilog);

    builder.SetInsertPoint(blockFalse);
    CodeScope* scopeFalse = exprIf.blocks.back();
    llvm::Value* resultFalse = function->getBruteScope(scopeFalse)->compile();
    llvm::BasicBlock *  blockFalseEnd = builder.GetInsertBlock();
    builder.CreateBr(blockEpilog);

    builder.SetInsertPoint(scope->lastBlockRaw);
    llvm->irBuilder.CreateCondBr(cond, blockTrue, blockFalse);

    builder.SetInsertPoint(blockEpilog);
    llvm::PHINode *ret = builder.CreatePHI(resultTrue->getType(), 2, hintRetVar);

    ret->addIncoming(resultTrue, blockTrueEnd);
    ret->addIncoming(resultFalse, blockFalseEnd);

    return ret;
}

//TODO Switch: default variant no needed when all possible conditions are considered
llvm::Value*
ControlIR::compileSwitch(const Expression& exprSwitch, const std::string& hintRetVar) {
    EXPAND_CONTEXT    UNUSED(function);
    AST* root = context.pass->man->root;
    llvm::IRBuilder<>& builder = llvm->irBuilder;
    assert(exprSwitch.operands.size() >= 2);
    assert(exprSwitch.operands[1].op == Operator::CASE_DEFAULT && "No default case in Switch Statement");

    int countCases = exprSwitch.operands.size() - 1;
    llvm::BasicBlock* blockProlog = builder.GetInsertBlock();
    llvm::BasicBlock *blockEpilog = llvm::BasicBlock::Create(llvm->llvmContext, "switchAfter", function->raw);
    builder.SetInsertPoint(blockEpilog);
    llvm::Type* exprSwitchType = llvm->toLLVMType(root->getType(exprSwitch));
    llvm::PHINode *ret = builder.CreatePHI(exprSwitchType, countCases, hintRetVar);
    llvm::Type* typI8 = llvm::Type::getInt8Ty(llvm->llvmContext);

    builder.SetInsertPoint(blockProlog);
    llvm::Value * conditionSwitch = scope->process(exprSwitch.operands[0]);
    llvm::BasicBlock *blockDefault = llvm::BasicBlock::Create(llvm->llvmContext, "caseDefault", function->raw);
    llvm::SwitchInst * instructionSwitch = builder.CreateSwitch(
        typeinference::doAutomaticTypeConversion(conditionSwitch, typI8, builder),
        blockDefault,
        countCases);

    for (int size = exprSwitch.operands.size(), i = 2; i < size; ++i) {
        llvm::BasicBlock *blockCase = llvm::BasicBlock::Create(llvm->llvmContext, "case" + std::to_string(i), function->raw);

        llvm::Value* condCase = function->getBruteScope(exprSwitch.operands[i].blocks.front())->compile();
        builder.SetInsertPoint(blockCase);
        llvm::Value* resultCase = function->getBruteScope(exprSwitch.operands[i].blocks.back())->compile();
        builder.CreateBr(blockEpilog);

        ret->addIncoming(resultCase, builder.GetInsertBlock());
        builder.SetInsertPoint(blockProlog);
        instructionSwitch->addCase(
            dyn_cast<llvm::ConstantInt>(
                typeinference::doAutomaticTypeConversion(condCase, typI8, builder)),
            blockCase);
    }

    //compile default block:
    builder.SetInsertPoint(blockDefault);
    CodeScope* scopeDefault = exprSwitch.operands[1].blocks.front();
    llvm::Value* resultDefault = function->getBruteScope(scopeDefault)->compile();
    builder.CreateBr(blockEpilog);
    ret->addIncoming(resultDefault, builder.GetInsertBlock());
    builder.SetInsertPoint(blockEpilog);

    return ret;
}

llvm::Value*
ControlIR::compileSwitchVariant(const Expression& exprSwitch, const std::string& hintRetVar) {
    EXPAND_CONTEXT     UNUSED(function);
    AST* root = context.pass->man->root;
    llvm::IRBuilder<>& builder = llvm->irBuilder;
    llvm::Type* typI8= llvm::Type::getInt8Ty(llvm->llvmContext);
    const ExpandedType& typVariant = root->getType(exprSwitch.operands.at(0));
    llvm::Type* typVariantRaw = llvm->toLLVMType(typVariant);

    assert(typVariant->__operands.size() == exprSwitch.operands.size() - 1 && "Ill-formed Switch Variant");

    int casesCount = exprSwitch.operands.size();
    llvm::BasicBlock* blockProlog = builder.GetInsertBlock();
    llvm::BasicBlock *blockEpilog = llvm::BasicBlock::Create(llvm->llvmContext, "switchAfter", function->raw);
    builder.SetInsertPoint(blockEpilog);
    llvm::Type* resultType = llvm->toLLVMType(root->getType(exprSwitch));
    llvm::PHINode *ret = builder.CreatePHI(resultType, casesCount, hintRetVar);

    builder.SetInsertPoint(blockProlog);
    llvm::Value * conditionSwitchRaw = scope->process(exprSwitch.operands.at(0));
    llvm::Value* idRaw = builder.CreateExtractValue(conditionSwitchRaw, llvm::ArrayRef<unsigned>({0}));

    //Dereference preparation
    const bool flagPrepareDerefence = std::any_of(typVariant->__operands.begin(), typVariant->__operands.end(), [](const TypeAnnotation& op){
        return op.isValid();
    });

    llvm::Value* addrAsStorage = nullptr;
    if (flagPrepareDerefence){
        assert(exprSwitch.bindings.size() && "Switch condition alias not found");
        llvm::Type* typStorageRaw = llvm::cast<llvm::StructType>(typVariantRaw)->getElementType(1);
        llvm::Value* storageRaw = builder.CreateExtractValue(conditionSwitchRaw, llvm::ArrayRef<unsigned>({1}));
        addrAsStorage = llvm->irBuilder.CreateAlloca(typStorageRaw);
        llvm->irBuilder.CreateStore(storageRaw, addrAsStorage);
    }

    llvm::SwitchInst * instructionSwitch = builder.CreateSwitch(idRaw, nullptr, casesCount);
    llvm::BasicBlock* blockDefaultUndefined;
    std::list<CodeScope*>::const_iterator scopeCaseIt = exprSwitch.blocks.begin();
    for (int instancesSize = exprSwitch.operands.size()-1, instId = 0; instId < instancesSize; ++instId) {
        llvm::BasicBlock *blockCase = llvm::BasicBlock::Create(llvm->llvmContext, "case" + std::to_string(instId), function->raw);
        builder.SetInsertPoint(blockCase);
        IBruteScope* unitCase = function->getBruteScope(*scopeCaseIt);
        const ExpandedType& instType = ExpandedType(typVariant->__operands.at(instId));

        //Actual variant derefence
        if (instType->isValid()) {
            string identCondition = exprSwitch.bindings.front();
            llvm::Type* instTypeRaw = llvm->toLLVMType(instType);
            llvm::Value* addrAsInst =  llvm->irBuilder.CreateBitOrPointerCast(addrAsStorage, instTypeRaw->getPointerTo());
            llvm::Value* instRaw = llvm->irBuilder.CreateLoad(instTypeRaw, addrAsInst);
            const Symbol& identSymb =  unitCase->bindArg(instRaw, move(identCondition));
            Attachments::put<TypeInferred>(identSymb, instType);
        }

        llvm::Value* resultCase = function->getBruteScope(*scopeCaseIt)->compile();
        builder.CreateBr(blockEpilog);

        ret->addIncoming(resultCase, blockDefaultUndefined = builder.GetInsertBlock());
        builder.SetInsertPoint(blockProlog);
        instructionSwitch->addCase(dyn_cast<llvm::ConstantInt>(llvm::ConstantInt::get(typI8, exprSwitch.operands.at(instId+1).getValueDouble())), blockCase);
        ++scopeCaseIt;
    }

    instructionSwitch->setDefaultDest(blockDefaultUndefined);
    builder.SetInsertPoint(blockEpilog);
    return ret;
}

llvm::Value*
ControlIR::compileConstantStringAsPChar(const string& data, const std::string& hintRetVar) {
    EXPAND_CONTEXT UNUSED(function); UNUSED(scope);

    Type* typPchar = PointerType::getUnqual(Type::getInt8Ty(llvm->llvmContext));
    //ArrayType* typStr = (ArrayType*) (llvm->toLLVMType(ExpandedType(TypeAnnotation(tag_array, TypePrimitive::I8, size+1))));

    /*
    std::vector<Constant *> chars;
    chars.reserve(size+1);

    for (size_t i=0; i< size; ++i){
        chars[i] = ConstantInt::get(typI8, (unsigned char) data[i]);
    }
    chars[size] = ConstantInt::get(typI8, 0);
     */

    Value* rawData = ConstantDataArray::getString(llvm->llvmContext, data);
    Value* rawPtrData = llvm->irBuilder.CreateAlloca(rawData->getType(), ConstantInt::get(Type::getInt32Ty(llvm->llvmContext), 1, false));
    llvm->irBuilder.CreateStore(rawData, rawPtrData);
    return llvm->irBuilder.CreateCast(llvm::Instruction::BitCast, rawPtrData, typPchar, hintRetVar);
}

llvm::Value*
ControlIR::compileSequence(const Expression &expr){
    EXPAND_CONTEXT     UNUSED(scope); UNUSED(llvm);

    llvm::Value* result;
    for(CodeScope* scope: expr.blocks){
        result = function->getBruteScope(scope)->compile();
    }

    return result;
}