/*
 * Copyright (c) 2020 Trail of Bits, Inc.
 *
 * 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 <remill/BC/TraceLifter.h>

#include "InstructionLifter.h"

#include <set>
#include <sstream>

namespace remill {
namespace {

}  // namespace

TraceManager::~TraceManager(void) {}

// Return an already lifted trace starting with the code at address
// `addr`.
llvm::Function *TraceManager::GetLiftedTraceDeclaration(uint64_t) {
  return nullptr;
}

// Return an already lifted trace starting with the code at address
// `addr`.
llvm::Function *TraceManager::GetLiftedTraceDefinition(uint64_t) {
  return nullptr;
}

// Apply a callback that gives the decoder access to multiple virtual
// targets of this instruction (indirect call or jump).
void TraceManager::ForEachDevirtualizedTarget(
    const Instruction &,
    std::function<void(uint64_t, DevirtualizedTargetKind)>) {

  // Must be extended.
}

// Figure out the name for the trace starting at address `addr`.
std::string TraceManager::TraceName(uint64_t addr) {
  std::stringstream ss;
  ss << "sub_" << std::hex << addr;
  return ss.str();
}

namespace {

using DecoderWorkList = std::set<uint64_t>;  // For ordering.

}  // namespace

class TraceLifter::Impl {
 public:
  Impl(InstructionLifter *inst_lifter_, TraceManager *manager_);

  // Lift one or more traces starting from `addr`. Calls `callback` with each
  // lifted trace.
  bool Lift(uint64_t addr,
            std::function<void(uint64_t, llvm::Function *)> callback);

  // Reads the bytes of an instruction at `addr` into `state.inst_bytes`.
  bool ReadInstructionBytes(uint64_t addr);

  // Return an already lifted trace starting with the code at address
  // `addr`.
  //
  // NOTE: This is guaranteed to return either `nullptr`, or a function
  //       within `module`.
  llvm::Function *GetLiftedTraceDeclaration(uint64_t addr);

  // Return an already lifted trace starting with the code at address
  // `addr`.
  //
  // NOTE: This is guaranteed to return either `nullptr`, or a function
  //       within `module`.
  llvm::Function *GetLiftedTraceDefinition(uint64_t addr);

  llvm::BasicBlock *GetOrCreateBlock(uint64_t block_pc) {
    auto &block = blocks[block_pc];
    if (!block) {
      block = llvm::BasicBlock::Create(context, "", func);
    }
    return block;
  }

  llvm::BasicBlock *GetOrCreateBranchTakenBlock(void) {
    inst_work_list.insert(inst.branch_taken_pc);
    return GetOrCreateBlock(inst.branch_taken_pc);
  }

  llvm::BasicBlock *GetOrCreateBranchNotTakenBlock(void) {
    inst_work_list.insert(inst.branch_not_taken_pc);
    return GetOrCreateBlock(inst.branch_not_taken_pc);
  }

  llvm::BasicBlock *GetOrCreateNextBlock(void) {
    inst_work_list.insert(inst.next_pc);
    return GetOrCreateBlock(inst.next_pc);
  }

  uint64_t PopTraceAddress(void) {
    auto trace_it = trace_work_list.begin();
    const auto trace_addr = *trace_it;
    trace_work_list.erase(trace_it);
    return trace_addr;
  }

  uint64_t PopInstructionAddress(void) {
    auto inst_it = inst_work_list.begin();
    const auto inst_addr = *inst_it;
    inst_work_list.erase(inst_it);
    return inst_addr;
  }

  const Arch *const arch;
  InstructionLifter &inst_lifter;
  const remill::IntrinsicTable *intrinsics;
  llvm::LLVMContext &context;
  llvm::Module *const module;
  const uint64_t addr_mask;
  TraceManager &manager;

  llvm::Function *func;
  llvm::BasicBlock *block;
  llvm::SwitchInst *switch_inst;
  const size_t max_inst_bytes;
  std::string inst_bytes;
  Instruction inst;
  Instruction delayed_inst;
  DecoderWorkList trace_work_list;
  DecoderWorkList inst_work_list;
  std::map<uint64_t, llvm::BasicBlock *> blocks;
};

TraceLifter::Impl::Impl(InstructionLifter *inst_lifter_, TraceManager *manager_)
    : arch(inst_lifter_->impl->arch),
      inst_lifter(*inst_lifter_),
      intrinsics(inst_lifter.impl->intrinsics),
      context(inst_lifter.impl->word_type->getContext()),
      module(intrinsics->async_hyper_call->getParent()),
      addr_mask(arch->address_size >= 64 ? ~0ULL : (~0ULL >> arch->address_size)),
      manager(*manager_),
      func(nullptr),
      block(nullptr),
      switch_inst(nullptr),
      max_inst_bytes(arch->MaxInstructionSize()) {

  inst_bytes.reserve(max_inst_bytes);
}

// Return an already lifted trace starting with the code at address
// `addr`.
llvm::Function *TraceLifter::Impl::GetLiftedTraceDeclaration(uint64_t addr) {
  auto func = manager.GetLiftedTraceDeclaration(addr);
  if (!func || func->getParent() == module) {
    return func;
  }

  return nullptr;
}

// Return an already lifted trace starting with the code at address
// `addr`.
llvm::Function *TraceLifter::Impl::GetLiftedTraceDefinition(uint64_t addr) {
  auto func = manager.GetLiftedTraceDefinition(addr);
  if (!func || func->getParent() == module) {
    return func;
  }

  CHECK_EQ(&(func->getContext()), &context);

  auto func_type = llvm::dyn_cast<llvm::FunctionType>(
      RecontextualizeType(func->getFunctionType(), context));

  // Handle the different module situation by declaring the trace in
  // this module to be external, with the idea that it will link to
  // another module.
  auto extern_func = module->getFunction(func->getName());
  if (!extern_func || extern_func->getFunctionType() != func_type) {
    extern_func = llvm::Function::Create(
        func_type, llvm::GlobalValue::ExternalLinkage, func->getName(), module);

  } else if (extern_func->isDeclaration()) {
    extern_func->setLinkage(llvm::GlobalValue::ExternalLinkage);
  }

  return extern_func;
}

TraceLifter::~TraceLifter(void) {}

TraceLifter::TraceLifter(InstructionLifter *inst_lifter_,
                         TraceManager *manager_)
    : impl(new Impl(inst_lifter_, manager_)) {}

void TraceLifter::NullCallback(uint64_t, llvm::Function *) {}

// Reads the bytes of an instruction at `addr` into `inst_bytes`.
bool TraceLifter::Impl::ReadInstructionBytes(uint64_t addr) {
  inst_bytes.clear();
  for (size_t i = 0; i < max_inst_bytes; ++i) {
    const auto byte_addr = (addr + i) & addr_mask;
    if (byte_addr < addr) {
      break;  // 32- or 64-bit address overflow.
    }
    uint8_t byte = 0;
    if (!manager.TryReadExecutableByte(byte_addr, &byte)) {
      DLOG(WARNING) << "Couldn't read executable byte at " << std::hex
                    << byte_addr << std::dec;
      break;
    }
    inst_bytes.push_back(static_cast<char>(byte));
  }
  return !inst_bytes.empty();
}

// Lift one or more traces starting from `addr`.
bool TraceLifter::Lift(
    uint64_t addr, std::function<void(uint64_t, llvm::Function *)> callback) {
  return impl->Lift(addr, callback);
}

// Lift one or more traces starting from `addr`.
bool TraceLifter::Impl::Lift(
    uint64_t addr_, std::function<void(uint64_t, llvm::Function *)> callback) {
  auto addr = addr_ & addr_mask;
  if (addr < addr_) {  // Address is out of range.
    LOG(ERROR) << "Trace address " << std::hex << addr_ << " is too big"
               << std::dec;
    return false;
  }

  // Reset the lifting state.
  trace_work_list.clear();
  inst_work_list.clear();
  blocks.clear();
  inst_bytes.clear();
  func = nullptr;
  switch_inst = nullptr;
  block = nullptr;
  inst.Reset();
  delayed_inst.Reset();

  // Get a trace head that the manager knows about, or that we
  // will eventually tell the trace manager about.
  auto get_trace_decl = [=](uint64_t addr) -> llvm::Function * {
    if (auto trace = GetLiftedTraceDeclaration(addr)) {
      return trace;
    }

    if (trace_work_list.count(addr)) {
      const auto target_trace_name = manager.TraceName(addr);
      return DeclareLiftedFunction(module, target_trace_name);
    }

    return nullptr;
  };

  trace_work_list.insert(addr);
  while (!trace_work_list.empty()) {
    const auto trace_addr = PopTraceAddress();

    // Already lifted.
    func = GetLiftedTraceDefinition(trace_addr);
    if (func) {
      continue;
    }

    DLOG(INFO) << "Lifting trace at address " << std::hex << trace_addr
               << std::dec;

    func = get_trace_decl(trace_addr);
    blocks.clear();

    if (!func || !func->isDeclaration()) {
      const auto trace_name = manager.TraceName(trace_addr);
      func = DeclareLiftedFunction(module, trace_name);
    }

    CHECK(func->isDeclaration());

    // Fill in the function, and make sure the block with all register
    // variables jumps to the block that will contain the first instruction
    // of the trace.
    CloneBlockFunctionInto(func);
    auto state_ptr = NthArgument(func, kStatePointerArgNum);

    if (auto entry_block = &(func->front())) {
      auto pc = LoadProgramCounterArg(func);
      auto next_pc_ref =
          inst_lifter.LoadRegAddress(entry_block, state_ptr, kNextPCVariableName);

      // Initialize `NEXT_PC`.
      (void) new llvm::StoreInst(pc, next_pc_ref, entry_block);

      // Branch to the first basic block.
      llvm::BranchInst::Create(GetOrCreateBlock(trace_addr), entry_block);
    }

    CHECK(inst_work_list.empty());
    inst_work_list.insert(trace_addr);

    // Decode instructions.
    while (!inst_work_list.empty()) {
      const auto inst_addr = PopInstructionAddress();

      block = GetOrCreateBlock(inst_addr);
      switch_inst = nullptr;

      // We have already lifted this instruction block.
      if (!block->empty()) {
        continue;
      }

      // Check to see if this instruction corresponds with an existing
      // trace head, and if so, tail-call into that trace directly without
      // decoding or lifting the instruction.
      if (inst_addr != trace_addr) {
        if (auto inst_as_trace = get_trace_decl(inst_addr)) {
          AddTerminatingTailCall(block, inst_as_trace);
          continue;
        }
      }

      // No executable bytes here.
      if (!ReadInstructionBytes(inst_addr)) {
        AddTerminatingTailCall(block, intrinsics->missing_block);
        continue;
      }

      inst.Reset();

      (void) arch->DecodeInstruction(inst_addr, inst_bytes, inst);

      auto lift_status = inst_lifter.LiftIntoBlock(inst, block, state_ptr);
      if (kLiftedInstruction != lift_status) {
        AddTerminatingTailCall(block, intrinsics->error);
        continue;
      }

      // Handle lifting a delayed instruction.
      auto try_delay = arch->MayHaveDelaySlot(inst);
      if (try_delay) {
        delayed_inst.Reset();
        if (!ReadInstructionBytes(inst.delayed_pc) ||
            !arch->DecodeDelayedInstruction(inst.delayed_pc, inst_bytes,
                                            delayed_inst)) {
          LOG(ERROR) << "Couldn't read delayed inst "
                     << delayed_inst.Serialize();
          AddTerminatingTailCall(block, intrinsics->error);
          continue;
        }
      }

      // Functor used to add in a delayed instruction.
      auto try_add_delay_slot = [&](bool on_branch_taken_path,
                                    llvm::BasicBlock *into_block) -> void {
        if (!try_delay) {
          return;
        }
        if (!arch->NextInstructionIsDelayed(inst, delayed_inst,
                                            on_branch_taken_path)) {
          return;
        }
        lift_status = inst_lifter.LiftIntoBlock(
            delayed_inst, into_block, state_ptr, true /* is_delayed */);
        if (kLiftedInstruction != lift_status) {
          AddTerminatingTailCall(block, intrinsics->error);
        }
      };

      // Connect together the basic blocks.
      switch (inst.category) {
        case Instruction::kCategoryInvalid:
        case Instruction::kCategoryError:
          AddTerminatingTailCall(block, intrinsics->error);
          break;

        case Instruction::kCategoryNormal:
        case Instruction::kCategoryNoOp:
          llvm::BranchInst::Create(GetOrCreateNextBlock(), block);
          break;

        // Direct jumps could either be local or could be tail-calls. In the
        // case of a tail call, we'll assume that the trace manager contains
        // advanced knowledge of this, and so when we go to make a block for
        // the targeted instruction, we'll either tail call to the target
        // trace, or we'll just extend out the current trace. Either way, no
        // sacrifice in correctness is made.
        case Instruction::kCategoryDirectJump:
          try_add_delay_slot(true, block);
          llvm::BranchInst::Create(GetOrCreateBranchTakenBlock(), block);
          break;

        case Instruction::kCategoryIndirectJump: {
          try_add_delay_slot(true, block);

          // The trace manager might know about the targets of things like
          // jump tables, so we will let it tell us about those possibilities.
          std::unordered_map<uint64_t, llvm::BasicBlock *> devirt_targets;
          manager.ForEachDevirtualizedTarget(
              inst,
              [&](uint64_t target_addr, DevirtualizedTargetKind target_kind) {
                if (target_kind == DevirtualizedTargetKind::kTraceHead) {
                  auto target_block =
                      llvm::BasicBlock::Create(context, "", func);
                  devirt_targets[target_addr] = target_block;

                  // Always add to the work list. This will cause us to lift
                  // if we haven't, and guarantee that `get_trace_decl` returns
                  // something.
                  trace_work_list.insert(target_addr);
                  auto target_trace = get_trace_decl(target_addr);
                  AddTerminatingTailCall(target_block, target_trace);

                } else {
                  devirt_targets[target_addr] = GetOrCreateBlock(target_addr);
                  inst_work_list.insert(target_addr);
                }
              });

          if (devirt_targets.empty()) {
            AddTerminatingTailCall(block, intrinsics->jump);
            break;
          }

          auto default_case = llvm::BasicBlock::Create(context, "", func);

          auto pc = LoadProgramCounter(block);
          auto pc_type = pc->getType();
          auto dispatcher = llvm::SwitchInst::Create(
              pc, default_case, devirt_targets.size(), block);
          for (auto devirt_target : devirt_targets) {
            dispatcher->addCase(
                llvm::dyn_cast<llvm::ConstantInt>(llvm::ConstantInt::get(
                    pc_type, devirt_target.first, false)),
                devirt_target.second);
          }
          break;
        }

        case Instruction::kCategoryAsyncHyperCall:
          AddCall(block, intrinsics->async_hyper_call);
          goto check_call_return;

        case Instruction::kCategoryIndirectFunctionCall: {
          try_add_delay_slot(true, block);
          const auto fall_through_block =
              llvm::BasicBlock::Create(context, "", func);

          const auto ret_pc_ref =
              LoadReturnProgramCounterRef(fall_through_block);
          const auto next_pc_ref =
              LoadNextProgramCounterRef(fall_through_block);
          llvm::IRBuilder<> ir(fall_through_block);
          ir.CreateStore(ir.CreateLoad(ret_pc_ref), next_pc_ref);
          ir.CreateBr(GetOrCreateNextBlock());

          // The trace manager might know about the targets of things like
          // virtual tables, so we will let it tell us about those possibilities.
          std::unordered_map<uint64_t, llvm::BasicBlock *> devirt_targets;
          manager.ForEachDevirtualizedTarget(
              inst,
              [&](uint64_t target_addr, DevirtualizedTargetKind target_kind) {
                if (target_kind == DevirtualizedTargetKind::kTraceLocal) {
                  LOG(WARNING)
                      << "Ignoring trace-local target in devirtualizable call";
                  return;
                }

                auto target_block = llvm::BasicBlock::Create(context, "", func);
                devirt_targets[target_addr] = target_block;

                // Always add to the work list. This will cause us to lift
                // if we haven't, and guarantee that `get_trace_decl` returns
                // something.
                trace_work_list.insert(target_addr);
                auto target_trace = get_trace_decl(target_addr);
                AddCall(target_block, target_trace);

                llvm::BranchInst::Create(fall_through_block, target_block);
              });

          if (devirt_targets.empty()) {
            AddCall(block, intrinsics->function_call);
            llvm::BranchInst::Create(fall_through_block, block);
            continue;
          }

          auto default_case = llvm::BasicBlock::Create(context, "", func);
          AddCall(default_case, intrinsics->function_call);
          llvm::BranchInst::Create(fall_through_block, default_case);

          auto pc = LoadProgramCounter(block);
          auto pc_type = pc->getType();
          auto dispatcher = llvm::SwitchInst::Create(
              pc, default_case, devirt_targets.size(), block);
          for (auto devirt_target : devirt_targets) {
            dispatcher->addCase(
                llvm::dyn_cast<llvm::ConstantInt>(llvm::ConstantInt::get(
                    pc_type, devirt_target.first, false)),
                devirt_target.second);
          }

          block = fall_through_block;
          continue;
        }

        // In the case of a direct function call, we try to handle the
        // pattern of a call to the next PC as a way of getting access to
        // an instruction pointer. It is the case where a call to the next
        // PC could also be something more like a call to a `noreturn` function
        // and that is OK, because either a user of the trace manager has
        // already told us that the next PC is a trace head (and we'll pick
        // that up when trying to lift it), or we'll just have a really big
        // trace for this function without sacrificing correctness.
        case Instruction::kCategoryDirectFunctionCall: {
          try_add_delay_slot(true, block);
          if (inst.next_pc != inst.branch_taken_pc) {
            trace_work_list.insert(inst.branch_taken_pc);
            auto target_trace = get_trace_decl(inst.branch_taken_pc);
            AddCall(block, target_trace);
          }

          const auto ret_pc_ref = LoadReturnProgramCounterRef(block);
          const auto next_pc_ref = LoadNextProgramCounterRef(block);
          llvm::IRBuilder<> ir(block);
          ir.CreateStore(ir.CreateLoad(ret_pc_ref), next_pc_ref);
          ir.CreateBr(GetOrCreateNextBlock());

          continue;
        }

        // Lift an async hyper call to check if it should do the hypercall.
        // If so, it will jump to the `do_hyper_call` block, otherwise it will
        // jump to the block associated with the next PC. In the case of the
        // `do_hyper_call` block, we assign it to `state.block`, then go
        // to `check_call_return` to add the hyper call into that block,
        // checking if the hyper call returns to the next PC or not.
        case Instruction::kCategoryConditionalAsyncHyperCall: {
          auto do_hyper_call = llvm::BasicBlock::Create(context, "", func);
          llvm::BranchInst::Create(do_hyper_call, GetOrCreateNextBlock(),
                                   LoadBranchTaken(block), block);
          block = do_hyper_call;
          AddCall(block, intrinsics->async_hyper_call);
          goto check_call_return;
        }

        check_call_return:
          do {
            auto pc = LoadProgramCounter(block);
            auto ret_pc =
                llvm::ConstantInt::get(inst_lifter.impl->word_type, inst.next_pc);

            llvm::IRBuilder<> ir(block);
            auto eq = ir.CreateICmpEQ(pc, ret_pc);
            auto unexpected_ret_pc =
                llvm::BasicBlock::Create(context, "", func);
            ir.CreateCondBr(eq, GetOrCreateNextBlock(), unexpected_ret_pc);
            AddTerminatingTailCall(unexpected_ret_pc,
                                   intrinsics->missing_block);
          } while (false);
          break;

        case Instruction::kCategoryFunctionReturn:
          try_add_delay_slot(true, block);
          AddTerminatingTailCall(block, intrinsics->function_return);
          break;

        case Instruction::kCategoryConditionalBranch: {
          auto taken_block = GetOrCreateBranchTakenBlock();
          auto not_taken_block = GetOrCreateBranchNotTakenBlock();

          // If we might need to add delay slots, then try to lift the delayed
          // instruction on each side of the conditional branch, injecting in
          // new blocks (for the delayed instruction) between the branch
          // and its original targets.
          if (try_delay) {
            auto new_taken_block = llvm::BasicBlock::Create(context, "", func);
            auto new_not_taken_block =
                llvm::BasicBlock::Create(context, "", func);

            try_add_delay_slot(true, new_taken_block);
            try_add_delay_slot(false, new_not_taken_block);

            llvm::BranchInst::Create(taken_block, new_taken_block);
            llvm::BranchInst::Create(not_taken_block, new_not_taken_block);

            taken_block = new_taken_block;
            not_taken_block = new_not_taken_block;
          }

          llvm::BranchInst::Create(taken_block, not_taken_block,
                                   LoadBranchTaken(block), block);
          break;
        }
      }
    }

    for (auto &block : *func) {
      if (!block.getTerminator()) {
        AddTerminatingTailCall(&block, intrinsics->missing_block);
      }
    }

    callback(trace_addr, func);
    manager.SetLiftedTraceDefinition(trace_addr, func);
  }

  return true;
}

}  // namespace remill