#include "config.h"
#include "system.h"
#include "coretypes.h"
#include "backend.h"
#include "tree.h"
#include "gimple-expr.h"
#include "predict.h"
#include "alloc-pool.h"
#include "tree-pass.h"
#include "cgraph.h"
#include "diagnostic.h"
#include "fold-const.h"
#include "gimple-fold.h"
#include "symbol-summary.h"
#include "tree-vrp.h"
#include "ipa-prop.h"
#include "tree-pretty-print.h"
#include "tree-inline.h"
#include "ipa-fnsummary.h"
#include "ipa-utils.h"
#include "tree-ssa-ccp.h"
#include "stringpool.h"
#include "attribs.h"
#include "tree-cfg.h"
#include "gimple.h"
#include "cfg.h" // needed for gimple-iterator.h
#include "gimple-iterator.h"

#include <map>
#include <set>

#include "ipa-type-escape-analysis.h"
#include "ipa-str-reorg-utils.h"
#include "ipa-hello-world.h"


//TODO: place in header file
inline static void
print_function (cgraph_node *cnode)
{
  if (!dump_file)
    return;
  gcc_assert (cnode);
  cnode->get_untransformed_body ();
  dump_function_to_file (cnode->decl, dump_file, TDF_NONE);
}

//TODO: do not use pair.
//This names are unintelligible
typedef std::pair<unsigned /* reads */, unsigned /* writes */> accesses;
struct field_comparator
{
  bool operator()(const fields &left, const fields &right) const
  {
    const_tree left_record = left.first;
    gcc_assert(left_record);
    const_tree right_record = right.first;
    gcc_assert(right_record);
    // Make sure that we are only comparing valid types...
    const enum tree_code tree_code_left_record = TREE_CODE(left_record);
    const enum tree_code tree_code_right_record = TREE_CODE(right_record);
    const bool is_left_record_type = RECORD_TYPE == tree_code_left_record;
    const bool is_right_record_type = RECORD_TYPE == tree_code_right_record;
    gcc_assert(is_left_record_type);
    gcc_assert(is_right_record_type);
    const bool are_left_and_right_valid = is_left_record_type && is_right_record_type;
    gcc_assert(are_left_and_right_valid);

    // Handle typedefs:
    // Get the main variants of the main types
    const_tree main_variant_left = TYPE_MAIN_VARIANT(left_record);
    gcc_assert(main_variant_left);
    const_tree main_variant_right = TYPE_MAIN_VARIANT(right_record);
    gcc_assert(main_variant_right);
    // If they are not equal, we can do a comparison of the types here...
    const bool are_main_variants_equal = main_variant_left == main_variant_right;
    const bool left_type_less_than_right_type = main_variant_left < main_variant_right;
    if (!are_main_variants_equal) return left_type_less_than_right_type;

    // If they are equal, that means that we are comparing fields defined in the same record
    const_tree left_field = left.second;
    gcc_assert(left_field);
    const_tree right_field = right.second;
    gcc_assert(right_field);
    // Make sure that they are valid
    const enum tree_code tree_code_left_field = TREE_CODE(left_field);
    const enum tree_code tree_code_right_field = TREE_CODE(right_field);
    const bool is_left_field_field_decl = FIELD_DECL == tree_code_left_field;
    const bool is_right_field_field_decl = FIELD_DECL == tree_code_right_field;
    const bool are_left_and_right_field_decls = is_left_field_field_decl && is_right_field_field_decl;
    gcc_assert(are_left_and_right_field_decls);

    // Compare on the field offset.
    const_tree left_constant = byte_position(left_field);
    gcc_assert(left_constant);
    const_tree right_constant = byte_position(right_field);
    gcc_assert(right_constant);
    unsigned left_offset = tree_to_uhwi(left_constant);
    unsigned right_offset = tree_to_uhwi(right_constant);
    const bool left_offset_less_than_right_offset = left_offset < right_offset;
    return left_offset_less_than_right_offset;
  }
};

typedef std::map<fields, accesses, field_comparator> field_access_counter;
typedef std::set<const_tree> record_set;

enum access_code { READ_ACCESS, WRITE_ACCESS };

void count_access_for_types_in_expr(const_tree expr, const record_set &non_escaping_records, field_access_counter &counter, const enum access_code access);

void
count_access_for_type_in_component_ref(const_tree component_ref, const record_set &non_escaping_records, field_access_counter &counter, const enum access_code access)
{
  log("types in component_ref\n");
  gcc_assert(component_ref);
  enum tree_code tree_code_component_ref = TREE_CODE(component_ref);
  const bool is_component_ref = COMPONENT_REF == tree_code_component_ref;
  gcc_assert(is_component_ref);

  const_tree _struct = TREE_OPERAND(component_ref, 0);
  gcc_assert(_struct);

  log ("going in recursion\n");
  count_access_for_types_in_expr(_struct, non_escaping_records, counter, READ_ACCESS);
  const_tree tree_type_struct = TREE_TYPE(_struct);
  gcc_assert(tree_type_struct);
  enum tree_code tree_type_struct_code = TREE_CODE(tree_type_struct);
  const bool is_record_type = RECORD_TYPE == tree_type_struct_code;
  //TODO: Also write something for UNION_TYPE
  switch (tree_type_struct_code)
  {
    case UNION_TYPE: return; break;
    case RECORD_TYPE: break;
    default: gcc_unreachable(); break;
  }
  gcc_assert(is_record_type);

  //FIXME: Future proofing or making things more difficult to read?
  bool in_set = 
#if __cplusplus > 201703L
          non_escaping_records.contains(tree_type_struct)
#else
	  non_escaping_records.find(tree_type_struct) != non_escaping_records.end()
#endif
	  ;

  in_set |= 
#if __cplusplus > 201703L
          non_escaping_records.contains(TYPE_MAIN_VARIANT(tree_type_struct))
#else
	  non_escaping_records.find(TYPE_MAIN_VARIANT(tree_type_struct)) != non_escaping_records.end()
#endif
	  ;


  log("%s is in non_escaping_records ? %s\n", get_type_name(tree_type_struct), in_set ? "true" : "false");
  log("access is %s\n", access == READ_ACCESS ? "read_access" : "write_access");
  log("%s is in non_escaping set %s\n", get_type_name(tree_type_struct), in_set ? "true" : "false");
  if (!in_set) return;

  const_tree field = TREE_OPERAND(component_ref, 1);
  gcc_assert(field);
  enum tree_code tree_code_field = TREE_CODE(field);
  const bool is_field_decl = FIELD_DECL == tree_code_field;
  gcc_assert(is_field_decl);

  const std::pair<const_tree, const_tree> struct_field_pair = std::make_pair(tree_type_struct, field);
  accesses &access_counter = counter[struct_field_pair];

  switch (access)
  {
    case READ_ACCESS:
     //TODO: do not use pair.
     //This names are unintelligible
     access_counter.first++;
     log("%s.%s read %d\n", get_type_name(tree_type_struct), get_field_name(field), access_counter.first);
    break;
    case WRITE_ACCESS:
     //TODO: do not use pair.
     //This names are unintelligible
      access_counter.second++;
      log("%s.%s write %d\n", get_type_name(tree_type_struct), get_field_name(field), access_counter.second);
    break;
    default:
    gcc_unreachable();
    break;
  }
}

static inline void
is_addr_expr_p(const_tree expr)
{
  gcc_assert(expr);
  const enum tree_code code = TREE_CODE(expr);
  const bool is_addr_expr = ADDR_EXPR == code;
  gcc_assert(is_addr_expr);
}


void
count_access_for_type_in_addr_expr(const_tree expr, const record_set &non_escaping_records, field_access_counter &counter, const enum access_code access)
{
  is_addr_expr_p(expr);
  const_tree op0 = TREE_OPERAND(expr, 0);
  count_access_for_types_in_expr(op0, non_escaping_records, counter, access);
}

static inline void
is_array_expr_p(const_tree expr)
{
  gcc_assert(expr);
  const enum tree_code code = TREE_CODE(expr);
  const bool is_addr_expr = ARRAY_REF == code;
  gcc_assert(is_addr_expr);
}

void
count_access_for_type_in_array_expr(const_tree expr, const record_set &non_escaping_records, field_access_counter &counter, const enum access_code access)
{
  is_array_expr_p(expr);
  const_tree op0 = TREE_OPERAND(expr, 0);
  const_tree op1 = TREE_OPERAND(expr, 1);
  count_access_for_types_in_expr(op0, non_escaping_records, counter, READ_ACCESS);
  count_access_for_types_in_expr(op1, non_escaping_records, counter, READ_ACCESS);
}

void
count_access_for_types_in_expr(const_tree expr, const record_set &non_escaping_records, field_access_counter &counter, const enum access_code access)
{
  log("types in expr\n");
  gcc_assert(expr);
  enum tree_code tree_code_expr = TREE_CODE(expr);
  switch (tree_code_expr)
  {
    case COMPONENT_REF: count_access_for_type_in_component_ref(expr, non_escaping_records, counter, access); break;
    case ADDR_EXPR: count_access_for_type_in_addr_expr(expr, non_escaping_records, counter, access); break;
    case ARRAY_REF : count_access_for_type_in_array_expr(expr, non_escaping_records, counter, access); break;
    default: break;
  }
}

void
count_access_for_types_in_lhs(gimple *stmt, const record_set &non_escaping_records, field_access_counter &counter)
{
  gcc_assert(stmt);
  const enum gimple_code gimple_code_stmt = gimple_code(stmt);
  const bool is_assign = GIMPLE_ASSIGN == gimple_code_stmt;
  const bool is_call = GIMPLE_CALL == gimple_code_stmt;
  const bool is_valid = is_assign || is_call;
  gcc_assert(is_valid);


  const_tree lhs = is_assign ? gimple_assign_lhs (stmt) : gimple_call_lhs (stmt);
  /* GIMPLE_CALL might have a lhs null.
   * E.g.
   * foo()
   */
  if (!lhs) return;

  count_access_for_types_in_expr(lhs, non_escaping_records, counter, WRITE_ACCESS);
}

void
count_access_for_types_in_rhs(gimple *stmt, const record_set &non_escaping_records, field_access_counter &counter)
{
  gcc_assert(stmt);
  const enum gimple_code gimple_code_stmt = gimple_code(stmt);
  const bool is_assign = GIMPLE_ASSIGN == gimple_code_stmt;
  gcc_assert(is_assign);

  log("types in rhs\n");
  const enum gimple_rhs_class gimple_rhs_class_stmt = gimple_assign_rhs_class(stmt);
  switch (gimple_rhs_class_stmt)
  {
    case GIMPLE_TERNARY_RHS:
    {
      log("gimple_ternary_rhs\n");
    }
    /* fall through */
    case GIMPLE_BINARY_RHS:
    {
      log("gimple_binary_rhs\n");
    }
    /* fall through */
    case GIMPLE_SINGLE_RHS:
    case GIMPLE_UNARY_RHS:
    {
      const_tree rhs1 = gimple_assign_rhs1(stmt);
      gcc_assert(rhs1);
      count_access_for_types_in_expr(rhs1, non_escaping_records, counter, READ_ACCESS);
      log("gimple_single_rhs\n");
    }
    break;
    default:
      gcc_unreachable();
    break;
  }
}

inline static void
is_gimple_assign_p(gimple *stmt)
{
  gcc_assert(stmt);
  const enum gimple_code code = gimple_code(stmt);
  const bool is_assign = GIMPLE_ASSIGN == code;
  gcc_assert(is_assign);
}

inline static void
is_gimple_call_p(gimple *stmt)
{
  gcc_assert(stmt);
  const enum gimple_code code = gimple_code(stmt);
  const bool is_call = GIMPLE_CALL == code;
  gcc_assert(is_call);

}

void
count_access_for_types_in_assign(gimple *stmt, const record_set &non_escaping_records, field_access_counter &counter)
{
  is_gimple_assign_p(stmt);

  count_access_for_types_in_rhs(stmt, non_escaping_records, counter);
  count_access_for_types_in_lhs(stmt, non_escaping_records, counter);
}

static void
count_access_for_types_in_call_rhs(gimple *stmt, const record_set &non_escaping_records, field_access_counter &counter)
{
  is_gimple_call_p(stmt);
  unsigned args = gimple_call_num_args (stmt);



  for (unsigned i = 0; i < args; i++)
  {
    const_tree arg = gimple_call_arg (stmt, i);
    count_access_for_types_in_expr(arg, non_escaping_records, counter, READ_ACCESS);
  }

}

void
count_access_for_types_in_call(gimple *stmt, const record_set &non_escaping_records, field_access_counter &counter)
{
  is_gimple_call_p(stmt);
  count_access_for_types_in_lhs(stmt, non_escaping_records, counter);

  /* TODO: We need to iterate over each argument and find out if it is a read */
  count_access_for_types_in_call_rhs(stmt, non_escaping_records, counter);
}

void
count_access_for_types_in_stmt(gimple *stmt, const record_set &non_escaping_records, field_access_counter &counter)
{
   gcc_assert(stmt);
   const enum gimple_code gimple_code_stmt = gimple_code(stmt);
   switch(gimple_code_stmt)
   {
     case GIMPLE_ASSIGN:
       count_access_for_types_in_assign(stmt, non_escaping_records, counter);
     break;
     case GIMPLE_CALL:
       count_access_for_types_in_call(stmt, non_escaping_records, counter);
     break;
     default:
     break;
   }
}

void
count_access_for_types_in_bb(basic_block bb, const record_set &non_escaping_records, field_access_counter &counter)
{
  gcc_assert(bb);
  for (gimple_stmt_iterator gsi = gsi_start_bb(bb); !gsi_end_p(gsi); gsi_next(&gsi))
  {
    gimple *stmt = gsi_stmt(gsi);
    count_access_for_types_in_stmt(stmt, non_escaping_records, counter);
  }
}

void
count_access_for_types_in_function(cgraph_node *cnode, const record_set &non_escaping_records, field_access_counter &counter)
{
  gcc_assert(cnode);
  tree decl = cnode->decl;
  gcc_assert(decl);
  function *func = DECL_STRUCT_FUNCTION (decl);
  gcc_assert(func);
  push_cfun(func);
  basic_block bb = NULL;
  FOR_EACH_BB_FN (bb, func)
  {
    gcc_assert(bb); 
    count_access_for_types_in_bb(bb, non_escaping_records, counter);
  }
  pop_cfun();
}

void
init_field_access_counter(field_access_counter &counter, const record_set &non_escaping_records)
{
  for (auto it = non_escaping_records.cbegin(); it != non_escaping_records.cend(); ++it)
  {
    const_tree record = *it;
    gcc_assert(record);
    enum tree_code tree_code_record_type = TREE_CODE(record);
    const bool is_record_type = RECORD_TYPE == tree_code_record_type;
    gcc_assert(is_record_type);

    for (tree field = TYPE_FIELDS (record); field; field = DECL_CHAIN (field))
    {
      gcc_assert(field);
      const std::pair<const_tree, const_tree> struct_field_pair = std::make_pair(record, field);
      counter[struct_field_pair] = { 0, 0 };
    }
  }
}

/* INFO:
 * Yes, I know we are returning a std::map.
 * Bad pattern? Maybe, but this will only be called once
 * and I rather pass by value because that allows
 * me to have a pure function and not worry about
 * garbage collection
 *
 * TODO: I'd like to change type_map for a std::map
 * TODO: I'd like to make this a template that can work
 * for std::map and std::set
 */
field_access_counter
count_access_for_types_in_linking_unit(const record_set &non_escaping_records)
{
  field_access_counter counter;
  init_field_access_counter(counter, non_escaping_records);
  cgraph_node *cnode = NULL;
  FOR_EACH_DEFINED_FUNCTION(cnode)
  {
    gcc_assert(cnode);
    log("printing a defined function\n");
    print_function(cnode);
    count_access_for_types_in_function(cnode, non_escaping_records, counter);
  }
  return counter;
}

bool
_calculate_non_escaping_records(const_tree const &type, escaping_info *info, record_set *non_escaping_records)
{
  
  gcc_assert(info);
  log("NOW: %s is escaping %s\n", get_type_name(type), info->is_escaping ? "true" : "false");
  if (info->is_escaping) return true;
  gcc_assert(non_escaping_records);
  
  enum tree_code tree_code_type = TREE_CODE(type);
  const bool is_record_type = RECORD_TYPE == tree_code_type;
  log("NOW: %s is record %s\n", get_type_name(type), is_record_type ? "true" : "false");
  if (!is_record_type) return true;

  non_escaping_records->insert(type);
  return true;
}

/*
 * Yes, again, we are passing a std::set
 * as a value. I don\t care too much since
 * this is only called once...
 */
static record_set
calculate_non_escaping_records(type_map &escaping_type_info)
{
  // We are going to have to traverse the type_map...
  // This is why I don't really like hash_set...
  record_set non_escaping_records;
  escaping_type_info.traverse<record_set*, _calculate_non_escaping_records>(&non_escaping_records);
  return non_escaping_records;
}

static field_access_counter
count_field_accesses()
{
  
  type_map escaping_types;
  calculate_escaping_types(escaping_types);
  const record_set non_escaping_records = calculate_non_escaping_records(escaping_types); 
  field_access_counter counter = count_access_for_types_in_linking_unit(non_escaping_records);
  return counter;
}

record_field_set
get_fields_to_reorg()
{
  const field_access_counter counter = count_field_accesses();
  record_field_set records_and_fields_to_reorg;
  for (auto it = counter.cbegin(); it != counter.cend(); ++it)
  {
     fields record_field_pair = it->first;
     accesses counter = it->second;
     const_tree record = record_field_pair.first;
     const_tree field = record_field_pair.second;
     // We are interested in reads == 0;
     unsigned reads = counter.first;
     log("final count for %s.%s = %d\n", get_type_name(record), get_field_name(field), reads);
     if (0 != reads) continue;

     records_and_fields_to_reorg.insert(record_field_pair);
  }

  return records_and_fields_to_reorg;
}

void
print_record_field_set(const record_field_set &to_reorg)
{
  log_2("I am about to print\n");
  for (auto it = to_reorg.cbegin(); it != to_reorg.cend(); ++it)
  {
    fields record_field_pair = *it;
    const_tree record = record_field_pair.first;
    const_tree field = record_field_pair.second;
    // TODO: Some fields / records might be anonymous
    log_2("will eliminate %s.%s\n", get_type_name(record), get_field_name(field));
  }
}


static unsigned int
iphw_execute()
{
  const record_field_set to_reorg = get_fields_to_reorg();
  log("I am about to enter print function\n");
  print_record_field_set(to_reorg);
  return 0;
}

namespace {
const pass_data pass_data_ipa_hello_world =
{
  SIMPLE_IPA_PASS,
  "hello-world",
  OPTGROUP_NONE,
  TV_NONE,
  (PROP_cfg | PROP_ssa),
  0,
  0,
  0,
  0,
};

class pass_ipa_hello_world : public simple_ipa_opt_pass
{
public:
  pass_ipa_hello_world (gcc::context *ctx)
    : simple_ipa_opt_pass(pass_data_ipa_hello_world, ctx)
  {}

  virtual bool gate(function*) { return flag_ipa_hello_world; }
  virtual unsigned execute (function*) { return iphw_execute(); }
};
} // anon namespace

simple_ipa_opt_pass*
make_pass_ipa_hello_world (gcc::context *ctx)
{
  return new pass_ipa_hello_world (ctx);
}