#include "initialize_stack.hpp"
#include <algorithm>
#include <fstream>
#include <math.h>

#define ALLOF(a) begin(a), end(a)

using namespace std;
using namespace IRDB_SDK;
using namespace InitStack;

//
// constructor
//
InitStack_t::InitStack_t(FileIR_t *p_variantIR,
                         const string &p_functions_filename, int p_init_value,
                         bool p_verbose)
    : Transform_t(p_variantIR), // initialize the Transform class so things like
                                // insertAssembly and getFileIR() can be used
      m_init_value(p_init_value), // member variable inits, these will vary
                                  // depending on your transform's objectives
      m_verbose(p_verbose), m_num_transformed(0) {
  // check whether to read in a list of functions to transform
  if (p_functions_filename == "") {
    cout << "Auto-initialize all functions" << endl;
    m_funcs_to_init =
        getFileIR()->getFunctions(); // use all functions from the IR
  } else {
    cout << "Auto-initialize functions specified in: " << p_functions_filename
         << endl;
    readFunctionsFromFile(p_functions_filename); // read functions from file
  }
}

//
// read list of functions to auto-initialize
//
// post conditions: set of functions to auto-initialize
//
void InitStack_t::readFunctionsFromFile(const string &p_filename) {
  // get all functions for readability of the rest of the code
  const auto &all_funcs = getFileIR()->getFunctions();

  // open input file and check for successful open
  ifstream functionsFile(
      p_filename); // can't use auto decl here because of lack of copy
                   // constructor in ifstream class
  if (!functionsFile.is_open())
    throw runtime_error("Cannot open " + p_filename);

  // read each line of the input file.
  auto line = string();
  while (functionsFile >> line) {
    // locate a function with the name read from the file.
    const auto func_it = find_if(ALLOF(all_funcs), [&](const Function_t *f) {
      return f->getName() == line;
    });

    // if found, log and insert it into the set to transform
    if (func_it != end(all_funcs)) {
      auto f = *func_it;
      cout << "Adding " << f->getName() << " to function list" << endl;
      m_funcs_to_init.insert(f);
    }
  }
}

//
// Execute the transform by transforming all to-transform functions
//
// preconditions: the FileIR is read as from the IRDB. valid file listing
// functions to auto-initialize postcondition: instructions added to
// auto-initialize stack for each specified function
//
//
bool InitStack_t::execute() {
  // transform all functions
  for (auto f : m_funcs_to_init)
    initStack(f);

  // #ATTRIBUTE is a convention used to help find useful information in log
  // files
  cout << "#ATTRIBUTE InitStack::num_transformed=" << m_num_transformed << endl;

  return m_num_transformed > 0; // true means success
}

//
// preconditions : f is not NULL
// postconditions: stack auto-initialized if stack frame size > 0
//
void InitStack_t::initStack(Function_t *f) {
  // preconditions
  assert(f != nullptr);

  // check that the frame size is in the area we care about
  const auto frame_size = f->getStackFrameSize();

  // nothing to init.
  if (frame_size == 0)
    return;

  const auto num_locs = static_cast<uint64_t>(ceil(frame_size / 4.0));

  // debug output
  cout << "Function: " << f->getName()
       << " frame size: " << f->getStackFrameSize() << endl;

  // not all functions have an entry point
  const auto entry = f->getEntryPoint();
  if (!entry)
    return;

  // log what we are doing
  cout << "Function: " << f->getName() << " auto-initialize " << dec << num_locs
       << "  stack memory locations (4 bytes at a time) with value = " << hex
       << m_init_value << endl;

  // determine the registers to use on x86-32 or  x86-64
  const auto sp_reg =
      getFileIR()->getArchitectureBitWidth() == 64 ? "rsp" : "esp";
  const auto scratch_reg =
      getFileIR()->getArchitectureBitWidth() == 64 ? "r11" : "ecx";

  // Now, do the dirty work of inserting new assembly to initialize the stack.
  // Insert these instructions at the start of the function (to initialize the
  // stack frame before the function runs) Assume: flags are dead at function
  // entry. Future work: Verify this is true using dead register list. Note: we
  // spill a scratch register into the red zone at 120 bytes passed the end of
  // the frame

  const auto newInsns = insertAssemblyInstructionsBefore(
      entry,
      string() + "    mov [%%1 + %%2], %%3\n"
                 "    mov %%3, -%%4\n"
                 "L1: mov dword [%%1 + %%3 * 4 - 4], %%5\n"
                 "    inc %%3\n"
                 "    jnz 0\n"
                 "    mov %%3, [%%1 + %%2]\n",
      {sp_reg, to_string(-f->getStackFrameSize() - 120), scratch_reg,
       to_string(num_locs), to_string(m_init_value)});

  newInsns[4]->setTarget(newInsns[2]); // Link jnz to L1.

  // bump stats
  m_num_transformed++;
}