#include "all.h"
#include <stdbool.h>
#include <stddef.h>

static void mark_changed(bool *ch) { *ch = true; }

static bool is_arith(ast_type_e t) { return t == AST_INC || t == AST_DEC; }

static bool is_ptr(ast_type_e t) {
  return t == AST_PTR_INC || t == AST_PTR_DEC;
}

static int node_signed_value(const ast_node_t *n) {
  if (n->type == AST_DEC || n->type == AST_PTR_DEC)
    return -n->val;
  return n->val;
}

static void set_arith_node(ast_node_t *n, int v) {
  if (v > 0) {
    n->type = AST_INC;
    n->val = v;
  } else {
    n->type = AST_DEC;
    n->val = -v;
  }
}

static void set_ptr_node(ast_node_t *n, int v) {
  if (v > 0) {
    n->type = AST_PTR_INC;
    n->val = v;
  } else {
    n->type = AST_PTR_DEC;
    n->val = -v;
  }
}

static void optimize_list(ast_node_t **head, bool *ch);

static void optimize_loop_node(ast_node_t *n, bool *ch) {
  if (!n || n->type != AST_LOOP)
    return;

  if (n->children.len == 0 || n->children._data[0] == NULL) {
    n->children.len = 0;
    return;
  }

  optimize_list(&n->children._data[0], ch);

  if (n->children._data[0] == NULL) {
    n->children.len = 0;
    return;
  }

  if (n->children.len == 1) {
    ast_node_t *c = n->children._data[0];

    if ((c->type == AST_INC || c->type == AST_DEC) && c->val == 1) {
      n->type = AST_CLEAR;
      n->val = 0;
      n->children.len = 0;
      n->children._data[0] = NULL;
      mark_changed(ch);
    }
  }
}

static void optimize_list(ast_node_t **head, bool *ch) {
  ast_node_t **p = head;

  while (*p) {
    ast_node_t *n = *p;

    if (n->type == AST_LOOP) {
      optimize_loop_node(n, ch);

      if (n->type == AST_LOOP && n->children.len == 0) {
        *p = n->next;
        mark_changed(ch);
        continue;
      }
    }

    if (*p && (*p)->next) {
      ast_node_t *a = *p;
      ast_node_t *b = a->next;

      if (is_arith(a->type) && is_arith(b->type)) {
        int v = node_signed_value(a) + node_signed_value(b);

        if (v == 0) {
          a->next = b->next;
          mark_changed(ch);
          continue;
        }

        set_arith_node(a, v);
        a->next = b->next;
        mark_changed(ch);
        continue;
      }

      if (is_ptr(a->type) && is_ptr(b->type)) {
        int v = node_signed_value(a) + node_signed_value(b);

        if (v == 0) {
          a->next = b->next;
          mark_changed(ch);
          continue;
        }

        set_ptr_node(a, v);
        a->next = b->next;
        mark_changed(ch);
        continue;
      }
    }

    p = &(*p)->next;
  }
}

ast_node_t *optimize_ast(ast_node_t *ast) {
  bool ch;

  do {
    ch = false;
    optimize_list(&ast, &ch);
  } while (ch);

  return ast;
}