#include <rlang.h>
#include "walk.h"

#define SEXP_STACK_INIT_SIZE 256

enum sexp_iterator_type {
  SEXP_ITERATOR_TYPE_node,
  SEXP_ITERATOR_TYPE_pointer,
  SEXP_ITERATOR_TYPE_vector,
  SEXP_ITERATOR_TYPE_atomic
};

enum sexp_iterator_state {
  SEXP_ITERATOR_STATE_done,
  SEXP_ITERATOR_STATE_attrib,
  SEXP_ITERATOR_STATE_tag,
  SEXP_ITERATOR_STATE_car,
  SEXP_ITERATOR_STATE_cdr,
  SEXP_ITERATOR_STATE_elt
};

struct sexp_stack_info {
  r_obj* x;
  enum r_type type;

  const enum sexp_iterator_state* p_state;
  r_obj* const * v_arr;
  r_obj* const * v_arr_end;

  int depth;
  r_obj* parent;
  enum r_sexp_it_relation rel;
  enum r_sexp_it_direction dir;
};

#include "decl/walk-decl.h"


static
const enum sexp_iterator_state node_states[] = {
  SEXP_ITERATOR_STATE_attrib,
  SEXP_ITERATOR_STATE_tag,
  SEXP_ITERATOR_STATE_car,
  SEXP_ITERATOR_STATE_cdr,
  SEXP_ITERATOR_STATE_done
};
static
const enum sexp_iterator_state pointer_states[] = {
  SEXP_ITERATOR_STATE_attrib,
  SEXP_ITERATOR_STATE_tag,
  SEXP_ITERATOR_STATE_cdr,
  SEXP_ITERATOR_STATE_done
};
static
const enum sexp_iterator_state vector_states[] = {
  SEXP_ITERATOR_STATE_attrib,
  SEXP_ITERATOR_STATE_elt,
  SEXP_ITERATOR_STATE_done
};
static
const enum sexp_iterator_state structure_states[] = {
  SEXP_ITERATOR_STATE_attrib,
  SEXP_ITERATOR_STATE_done
};
static
const enum sexp_iterator_state done_state[] = {
  SEXP_ITERATOR_STATE_done
};


struct r_sexp_iterator* r_new_sexp_iterator(r_obj* root) {
  r_obj* shelter = KEEP(r_alloc_list(3));

  r_obj* it = r_alloc_raw(sizeof(struct r_sexp_iterator));
  r_list_poke(shelter, 0, it);
  struct r_sexp_iterator* p_it = r_raw_begin(it);

  struct r_dyn_array* p_stack = r_new_dyn_array(sizeof(struct sexp_stack_info), SEXP_STACK_INIT_SIZE);
  r_list_poke(shelter, 1, p_stack->shelter);

  // Slot 2 holds a pairlist chain of collected attribute pairlists
  // that need protection for the lifetime of the iterator

  enum r_type type = r_typeof(root);
  enum sexp_iterator_type it_type = sexp_iterator_type(type, root);
  bool has_attrib = sexp_has_attrib(type, root);

  struct sexp_stack_info root_info = {
    .x = root,
    .type = type,
    .depth = -1,
    .parent = r_null,
    .rel = R_SEXP_IT_RELATION_root
  };

  if (it_type == SEXP_ITERATOR_TYPE_atomic && !has_attrib) {
    root_info.p_state = NULL;
    root_info.dir = R_SEXP_IT_DIRECTION_leaf;
  } else {
    init_incoming_stack_info(&root_info, it_type, has_attrib);
  }

  r_dyn_push_back(p_stack, &root_info);

  *p_it = (struct r_sexp_iterator) {
    .shelter = shelter,
    .p_stack = p_stack,
    .x = r_null,
    .parent = r_null,
  };

  FREE(1);
  return p_it;
}

/*
 * An incoming node has a state indicating which edge we're at. An
 * outgoing node just need to be visited again and then popped. A
 * leaf node is just visited once and then popped.
 */
bool r_sexp_next(struct r_sexp_iterator* p_it) {
  struct r_dyn_array* p_stack = p_it->p_stack;
  if (!p_stack->count) {
    return false;
  }

  struct sexp_stack_info* p_info = (struct sexp_stack_info*) r_dyn_last(p_stack);

  if (p_it->skip_incoming) {
    p_it->skip_incoming = false;

    if (p_it->dir == R_SEXP_IT_DIRECTION_incoming) {
      r_dyn_pop_back(p_stack);
      return r_sexp_next(p_it);
    }
  }

  // In the normal case, if we push an "incoming" node on the stack it
  // means that we have already visited it and we are now visiting its
  // children. The root node is signalled with a depth of -1 so it can
  // be visited first before being visited as an incoming node.
  bool root = (p_info->depth == -1);

  if (!root && p_info->dir == R_SEXP_IT_DIRECTION_incoming) {
    return sexp_next_incoming(p_it, p_info);
  }

  r_ssize i = -1;
  if (p_info->v_arr) {
    i = p_info->v_arr_end - p_info->v_arr;
  }
  p_it->x = p_info->x;
  p_it->type = p_info->type;
  p_it->depth = p_info->depth;
  p_it->parent = p_info->parent;
  p_it->rel = p_info->rel;
  p_it->i = i;
  p_it->dir = p_info->dir;

  if (root) {
    ++p_it->depth;
    ++p_info->depth;

    // Incoming visit for the root node
    if (p_it->dir == R_SEXP_IT_DIRECTION_incoming) {
      return true;
    }
  }

  r_dyn_pop_back(p_stack);
  return true;
}

static
bool sexp_next_incoming(struct r_sexp_iterator* p_it,
                        struct sexp_stack_info* p_info) {
  enum sexp_iterator_state state = *p_info->p_state;
  r_obj* x = p_info->x;
  enum r_type type = p_info->type;

  struct sexp_stack_info child = { 0 };
  child.parent = x;
  child.depth = p_info->depth + 1;

  switch (state) {
  case SEXP_ITERATOR_STATE_attrib: {
    // Allocates a new pairlist to avoid direct ATTRIB() access.
    // Could use R_mapAttrib() instead but would require refactoring.
    r_obj* collected = r_attrib_collect(x);
    // Protect collected pairlist in the iterator's shelter chain
    r_obj* chain = r_new_node(collected, r_list_get(p_it->shelter, 2));
    r_list_poke(p_it->shelter, 2, chain);
    child.x = collected;
    child.rel = R_SEXP_IT_RELATION_attrib;
    break;
  }
  case SEXP_ITERATOR_STATE_elt:
    child.x = *p_info->v_arr;
    child.rel = R_SEXP_IT_RELATION_list_elt;
    break;
  case SEXP_ITERATOR_STATE_tag:
    child.x = sexp_node_tag(type, x, &child.rel);
    break;
  case SEXP_ITERATOR_STATE_car:
    child.x = sexp_node_car(type, x, &child.rel);
    break;
  case SEXP_ITERATOR_STATE_cdr:
    child.x = sexp_node_cdr(type, x, &child.rel);
    break;
  case SEXP_ITERATOR_STATE_done:
    r_stop_unreachable();
  }

  child.type = r_typeof(child.x);
  bool has_attrib = sexp_has_attrib(child.type, child.x);
  enum sexp_iterator_type it_type = sexp_iterator_type(child.type, child.x);

  if (it_type == SEXP_ITERATOR_TYPE_atomic && !has_attrib) {
    child.p_state = NULL;
    child.dir = R_SEXP_IT_DIRECTION_leaf;
  } else {
    init_incoming_stack_info(&child, it_type, has_attrib);

    // Push incoming node on the stack so it can be visited again,
    // either to descend its children or to visit it again on the
    // outgoing trip
    r_dyn_push_back(p_it->p_stack, &child);
  }

  // Bump state for next iteration
  if (state == SEXP_ITERATOR_STATE_elt) {
    ++p_info->v_arr;
    if (p_info->v_arr == p_info->v_arr_end) {
      p_info->p_state = done_state;
    }
  } else {
    ++p_info->p_state;
  }

  // Flip incoming to outgoing if we're done visiting children after
  // this iteration. We don't leave a done node on the stack because
  // that would break the invariant that there are remaining nodes to
  // visit when `n > 0` and that the stack can be popped.
  if (*p_info->p_state == SEXP_ITERATOR_STATE_done) {
    p_info->dir = R_SEXP_IT_DIRECTION_outgoing;
  }

  r_ssize i = -1;
  if (child.v_arr) {
    i = child.v_arr_end - child.v_arr;
  }

  p_it->x = child.x;
  p_it->type = child.type;
  p_it->depth = child.depth;
  p_it->parent = child.parent;
  p_it->rel = child.rel;
  p_it->i = i;
  p_it->dir = child.dir;
  return true;
}


static inline
void init_incoming_stack_info(struct sexp_stack_info* p_info,
                              enum sexp_iterator_type it_type,
                              bool has_attrib) {
  p_info->dir = R_SEXP_IT_DIRECTION_incoming;

  switch (it_type) {
  case SEXP_ITERATOR_TYPE_atomic:
    p_info->p_state = structure_states;
    break;
  case SEXP_ITERATOR_TYPE_node:
    p_info->p_state = node_states + !has_attrib;
    break;
  case SEXP_ITERATOR_TYPE_pointer:
    p_info->p_state = pointer_states + !has_attrib;
    break;
  case SEXP_ITERATOR_TYPE_vector:
    p_info->v_arr = r_vec_cbegin(p_info->x);
    p_info->v_arr_end = p_info->v_arr + r_length(p_info->x);
    p_info->p_state = vector_states + !has_attrib;
    break;
  }
}

static inline
enum sexp_iterator_type sexp_iterator_type(enum r_type type,
                                           r_obj* x) {
  switch (type) {
  case R_TYPE_closure:
  case R_TYPE_environment:
  case R_TYPE_promise:
  case R_TYPE_pairlist:
  case R_TYPE_call:
  case R_TYPE_dots:
    return SEXP_ITERATOR_TYPE_node;
  case R_TYPE_pointer:
    return SEXP_ITERATOR_TYPE_pointer;
  case R_TYPE_list:
  case R_TYPE_expression:
  case R_TYPE_character:
    if (r_length(x)) {
      return SEXP_ITERATOR_TYPE_vector;
    } else {
      return SEXP_ITERATOR_TYPE_atomic;
    }
  default:
    return SEXP_ITERATOR_TYPE_atomic;
  }
}
static inline
bool sexp_has_attrib(enum r_type type, r_obj* x) {
  // Strings have private data stored in attributes
  if (type == R_TYPE_string) {
    return false;
  } else {
    return r_attrib_has_any(x);
  }
}
static inline
r_obj* sexp_node_car(enum r_type type,
                     r_obj* x,
                     enum r_sexp_it_relation* p_rel) {
  switch (type) {
  case R_TYPE_closure:     *p_rel = R_SEXP_IT_RELATION_function_fmls; return r_fn_formals(x);
  case R_TYPE_environment: *p_rel = R_SEXP_IT_RELATION_environment_frame; return FRAME(x);
  case R_TYPE_promise:     *p_rel = R_SEXP_IT_RELATION_promise_value; return PRVALUE(x);
  case R_TYPE_pairlist:
  case R_TYPE_call:
  case R_TYPE_dots:        *p_rel = R_SEXP_IT_RELATION_node_car; return CAR(x);
  case R_TYPE_pointer:
  default:                 *p_rel = -1; return r_null;
  }
}
static inline
r_obj* sexp_node_cdr(enum r_type type,
                     r_obj* x,
                     enum r_sexp_it_relation* p_rel) {
  switch (type) {
  case R_TYPE_closure:     *p_rel = R_SEXP_IT_RELATION_function_body; return r_fn_body(x);
  case R_TYPE_environment: *p_rel = R_SEXP_IT_RELATION_environment_enclos; return r_env_parent(x);
  case R_TYPE_promise:     *p_rel = R_SEXP_IT_RELATION_promise_expr; return PREXPR(x);
  case R_TYPE_pointer:     *p_rel = R_SEXP_IT_RELATION_pointer_prot; return EXTPTR_PROT(x);
  case R_TYPE_pairlist:
  case R_TYPE_call:
  case R_TYPE_dots:        *p_rel = R_SEXP_IT_RELATION_node_cdr; return CDR(x);
  default:                 *p_rel = -1; return r_null;
  }
}
static inline
r_obj* sexp_node_tag(enum r_type type,
                     r_obj* x,
                     enum r_sexp_it_relation* p_rel) {
  switch (type) {
  case R_TYPE_closure:     *p_rel = R_SEXP_IT_RELATION_function_env; return r_fn_env(x);
  case R_TYPE_environment: *p_rel = R_SEXP_IT_RELATION_environment_hashtab; return HASHTAB(x);
  case R_TYPE_promise:     *p_rel = R_SEXP_IT_RELATION_promise_env; return PRENV(x);
  case R_TYPE_pointer:     *p_rel = R_SEXP_IT_RELATION_pointer_tag; return EXTPTR_TAG(x);
  case R_TYPE_pairlist:
  case R_TYPE_call:
  case R_TYPE_dots:        *p_rel = R_SEXP_IT_RELATION_node_tag; return TAG(x);
  default:                 *p_rel = -1; return r_null;
  }
}


const char* r_sexp_it_direction_as_c_string(enum r_sexp_it_direction dir) {
  switch (dir) {
  case R_SEXP_IT_DIRECTION_leaf: return "leaf";
  case R_SEXP_IT_DIRECTION_incoming: return "incoming";
  case R_SEXP_IT_DIRECTION_outgoing: return "outgoing";
  default: r_stop_unreachable();
  }
}

const char* r_sexp_it_relation_as_c_string(enum r_sexp_it_relation rel) {
  switch (rel) {
  case R_SEXP_IT_RELATION_root: return "root";
  case R_SEXP_IT_RELATION_attrib: return "attrib";

  case R_SEXP_IT_RELATION_node_car: return "node_car";
  case R_SEXP_IT_RELATION_node_cdr: return "node_cdr";
  case R_SEXP_IT_RELATION_node_tag: return "node_tag";

  case R_SEXP_IT_RELATION_symbol_string: return "symbol_string";
  case R_SEXP_IT_RELATION_symbol_value: return "symbol_value";
  case R_SEXP_IT_RELATION_symbol_internal: return "symbol_internal";

  case R_SEXP_IT_RELATION_function_fmls: return "function_fmls";
  case R_SEXP_IT_RELATION_function_body: return "function_body";
  case R_SEXP_IT_RELATION_function_env: return "function_env";

  case R_SEXP_IT_RELATION_environment_frame: return "environment_frame";
  case R_SEXP_IT_RELATION_environment_enclos: return "environment_enclos";
  case R_SEXP_IT_RELATION_environment_hashtab: return "environment_hashtab";

  case R_SEXP_IT_RELATION_promise_value: return "promise_value";
  case R_SEXP_IT_RELATION_promise_expr: return "promise_expr";
  case R_SEXP_IT_RELATION_promise_env: return "promise_env";

  case R_SEXP_IT_RELATION_pointer_prot: return "pointer_prot";
  case R_SEXP_IT_RELATION_pointer_tag: return "pointer_tag";

  case R_SEXP_IT_RELATION_list_elt: return "list_elt";
  case R_SEXP_IT_RELATION_character_elt: return "character_elt";
  case R_SEXP_IT_RELATION_expression_elt: return "expression_elt";

  case R_SEXP_IT_RELATION_none: r_stop_internal("r_sexp_it_relation_as_c_string",
                                                "Found `R_SEXP_IT_RELATION_none`.");
  default: r_stop_unreachable();
  }
}

const char* r_sexp_it_raw_relation_as_c_string(enum r_sexp_it_raw_relation rel) {
  switch (rel) {
  case R_SEXP_IT_RAW_RELATION_root: return "root";
  case R_SEXP_IT_RAW_RELATION_attrib: return "attrib";
  case R_SEXP_IT_RAW_RELATION_node_car: return "node_car";
  case R_SEXP_IT_RAW_RELATION_node_cdr: return "node_cdr";
  case R_SEXP_IT_RAW_RELATION_node_tag: return "node_tag";
  case R_SEXP_IT_RAW_RELATION_vector_elt: return "vector_elt";
  default: r_stop_unreachable();
  }
}
