#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <errno.h>
#include <assert.h>

#include "treetrie.h"
#include "critbit.h"

int tt_dict_size(tt_arena_t *a, tt_node_ptr_t t) {
  if (TT_EMPTY_DICT_P(t)) {
    return 0;
  } else {
    assert(tt_ptr_tag(t) == TT_TAG_DICT);
    return TT_DICT_SIZE(a,t);
  }
}

static inline int bit_ref(tt_atom_t key, unsigned int bit) {
  return key & (1 << bit);
}

tt_node_ptr_t tt_dict_get(tt_arena_t *a, tt_node_ptr_t t, tt_atom_t key) {
  if (TT_EMPTY_DICT_P(t)) {
    return TT_NO_PTR;
  }

  assert(tt_ptr_tag(t) == TT_TAG_DICT);
  t = TT_DICT_ROOT(a,t);

  while (tt_ptr_tag(t) == TT_TAG_NODE) {
    t = bit_ref(key, TT_NODE_INDEX(a, t)) ? TT_NODE_ONE(a, t) : TT_NODE_ZERO(a, t);
  }

  assert(tt_ptr_tag(t) == TT_TAG_LEAF);
  return (TT_LEAF_ATOM(a, t) == key) ? TT_LEAF_TRIE(a, t) : TT_NO_PTR;
}

static tt_node_ptr_t splice_key(tt_arena_t *a,
				tt_atom_t key,
				tt_node_ptr_t trie,
				unsigned int index,
				tt_node_ptr_t sib)
{
  if (bit_ref(key, index)) {
    return tt_cons_node(a, index, sib, RET_IF_NO_PTR(tt_cons_leaf(a, trie, key)));
  } else {
    return tt_cons_node(a, index, RET_IF_NO_PTR(tt_cons_leaf(a, trie, key)), sib);
  }
}

static int set_walk(tt_arena_t *a,
		    tt_atom_t key,
		    tt_node_ptr_t trie,
		    tt_node_ptr_t n,
		    tt_node_ptr_t *result,
		    int *added_count)
{
  switch (tt_ptr_tag(n)) {
    case TT_TAG_LEAF: {
      tt_atom_t differences = key ^ TT_LEAF_ATOM(a, n);
      if (differences == 0) {
	/* Keys identical. If values identical, just return this node. */
	if (TT_LEAF_TRIE(a, n) == trie) {
	  *result = n;
	  return -1;
	} else {
	  *result = tt_cons_leaf(a, trie, key);
	  return -1;
	}
      } else {
	int leading_zeros = __builtin_clz(differences);
	int first_differing_bit = (8 * sizeof(differences)) - leading_zeros - 1;
	*added_count = 1;
	return first_differing_bit;
      }
    }
    case TT_TAG_NODE: {
      int index = TT_NODE_INDEX(a, n);
      if (bit_ref(key, index)) {
	tt_node_ptr_t one = TT_NODE_ONE(a, n);
	int first_differing_bit = set_walk(a, key, trie, one, result, added_count);
	if (first_differing_bit == -1) {
	  if (*result == one) {
	    *result = n;
	  } else if (TT_NO_PTR_P(*result)) {
	    /* Do nothing. */
	  } else {
	    *result = tt_cons_node(a, index, TT_NODE_ZERO(a, n), *result);
	  }
	  return -1;
	} else if (first_differing_bit > index) {
	  return first_differing_bit;
	} else {
	  *result = splice_key(a, key, trie, first_differing_bit, one);
	  if (!TT_NO_PTR_P(*result)) {
	    *result = tt_cons_node(a, index, TT_NODE_ZERO(a, n), *result);
	  }
	  return -1;
	}
      } else {
	tt_node_ptr_t zero = TT_NODE_ZERO(a, n);
	int first_differing_bit = set_walk(a, key, trie, zero, result, added_count);
	if (first_differing_bit == -1) {
	  if (*result == zero) {
	    *result = n;
	  } else if (TT_NO_PTR_P(*result)) {
	    /* Do nothing. */
	  } else {
	    *result = tt_cons_node(a, index, *result, TT_NODE_ONE(a, n));
	  }
	  return -1;
	} else if (first_differing_bit > index) {
	  return first_differing_bit;
	} else {
	  *result = splice_key(a, key, trie, first_differing_bit, zero);
	  if (!TT_NO_PTR_P(*result)) {
	    *result = tt_cons_node(a, index, *result, TT_NODE_ONE(a, n));
	  }
	  return -1;
	}
      }
    }
    default:
      abort();
  }
}

tt_node_ptr_t tt_dict_singleton(tt_arena_t *a,
				tt_atom_t key,
				tt_node_ptr_t trie)
{
  return tt_cons_dict(a, RET_IF_NO_PTR(tt_cons_leaf(a, trie, key)), 1);
}

tt_node_ptr_t tt_dict_set(tt_arena_t *a,
			  tt_node_ptr_t t,
			  tt_atom_t key,
			  tt_node_ptr_t trie)
{
  if (TT_EMPTY_DICT_P(t)) {
    return tt_dict_singleton(a, key, trie);
  }

  assert(tt_ptr_tag(t) == TT_TAG_DICT);

  {
    tt_node_ptr_t old_root = TT_DICT_ROOT(a,t);
    int added_count = 0;
    tt_node_ptr_t result = TT_NO_PTR;
    int first_differing_bit = set_walk(a, key, trie, old_root, &result, &added_count);
    if (first_differing_bit != -1) {
      result = splice_key(a, key, trie, first_differing_bit, old_root);
    }
    if (TT_NO_PTR_P(result)) {
      return TT_NO_PTR;
    } else if (result == old_root) {
      return t;
    } else {
      return tt_cons_dict(a, result, TT_DICT_SIZE(a,t) + added_count);
    }
  }
}

tt_node_ptr_t tt_dict_remove1(tt_arena_t *a,
			      tt_node_ptr_t n,
			      tt_atom_t key,
			      int *removed_count)
{
  switch (tt_ptr_tag(n)) {
    case TT_TAG_LEAF: {
      if (TT_LEAF_ATOM(a,n) == key) {
	*removed_count = 1;
	return TT_EMPTY_DICT;
      } else {
	return n;
      }
    }
    case TT_TAG_NODE: {
      int index = TT_NODE_INDEX(a,n);
      if (bit_ref(key, index)) {
	tt_node_ptr_t n1 =
	  RET_IF_NO_PTR(tt_dict_remove1(a, TT_NODE_ONE(a,n), key, removed_count));
	if (TT_EMPTY_DICT_P(n1)) {
	  return TT_NODE_ZERO(a,n);
	} else {
	  return tt_cons_node(a, index, TT_NODE_ZERO(a,n), n1);
	}
      } else {
	tt_node_ptr_t n1 =
	  RET_IF_NO_PTR(tt_dict_remove1(a, TT_NODE_ZERO(a,n), key, removed_count));
	if (TT_EMPTY_DICT_P(n1)) {
	  return TT_NODE_ONE(a,n);
	} else {
	  return tt_cons_node(a, index, n1, TT_NODE_ONE(a,n));
	}
      }
    }
    default:
      abort();
  }
}

tt_node_ptr_t tt_dict_remove(tt_arena_t *a,
			     tt_node_ptr_t t,
			     tt_atom_t key)
{
  if (TT_EMPTY_DICT_P(t)) {
    return TT_EMPTY_DICT;
  }

  assert(tt_ptr_tag(t) == TT_TAG_DICT);

  {
    tt_node_ptr_t old_root = TT_DICT_ROOT(a,t);
    int removed_count = 0;
    tt_node_ptr_t n =
      RET_IF_NO_PTR(tt_dict_remove1(a, old_root, key, &removed_count));
    if (TT_EMPTY_DICT_P(n)) {
      return TT_EMPTY_DICT;
    } else if (n == old_root) {
      return t;
    } else {
      return tt_cons_dict(a, n, TT_DICT_SIZE(a,t) - removed_count);
    }
  }
}

int tt_dict_foreach1(tt_arena_t *a,
		     tt_node_ptr_t n,
		     void *context,
		     int (*f)(void *context, tt_atom_t key, tt_node_ptr_t trie))
{
  switch (tt_ptr_tag(n)) {
    case TT_TAG_LEAF: {
      return f(context, TT_LEAF_ATOM(a,n), TT_LEAF_TRIE(a,n));
    }
    case TT_TAG_NODE: {
      if (!tt_dict_foreach1(a, TT_NODE_ZERO(a,n), context, f)) return 0;
      return tt_dict_foreach1(a, TT_NODE_ONE(a,n), context, f);
    }
    default:
      assert(0);
  }
}

int tt_dict_foreach(tt_arena_t *a,
		    tt_node_ptr_t t,
		    void *context,
		    int (*f)(void *, tt_atom_t key, tt_node_ptr_t trie))
{
  if (TT_EMPTY_DICT_P(t)) {
    return 1;
  } else {
    assert(tt_ptr_tag(t) == TT_TAG_DICT);
    return tt_dict_foreach1(a, TT_DICT_ROOT(a,t), context, f);
  }
}

struct tt_dictset_union_context {
  tt_arena_t *a;
  tt_node_ptr_t result;
};

static int tt_dictset_union_f(void *context, tt_atom_t key, tt_node_ptr_t ignored_trie) {
  struct tt_dictset_union_context *c = (struct tt_dictset_union_context *) context;
  tt_node_ptr_t new_result = tt_grab(c->a, tt_dict_set(c->a, c->result, key, TT_EMPTY));
  tt_drop(c->a, c->result);
  c->result = new_result;
  return !TT_NO_PTR_P(c->result);
}

tt_node_ptr_t tt_dictset_union(tt_arena_t *a, tt_node_ptr_t s1, tt_node_ptr_t s2) {
  struct tt_dictset_union_context context;
  context.a = a;
  context.result = TT_NO_PTR; /* grab'd */

  if (tt_dict_size(a, s1) < tt_dict_size(a, s2)) {
    tt_node_ptr_t tmp = s2;
    s2 = s1;
    s1 = tmp;
  }
  /* Now s1 is larger, s2 is smaller. */
  context.result = tt_grab(a, s1);
  tt_dict_foreach(a, s2, &context, tt_dictset_union_f);
  return context.result;
}

struct tt_dictset_intersection_context {
  tt_arena_t *a;
  tt_node_ptr_t s1;
  tt_node_ptr_t result;
};

static int tt_dictset_intersection_f(void *context, tt_atom_t key, tt_node_ptr_t ignored_trie) {
  struct tt_dictset_intersection_context *c = (struct tt_dictset_intersection_context *) context;
  if (!TT_NO_PTR_P(tt_dict_get(c->a, c->s1, key))) {
    tt_node_ptr_t new_result = tt_grab(c->a, tt_dict_set(c->a, c->result, key, TT_EMPTY));
    tt_drop(c->a, c->result);
    c->result = new_result;
  }
  return !TT_NO_PTR_P(c->result);
}

tt_node_ptr_t tt_dictset_intersection(tt_arena_t *a, tt_node_ptr_t s1, tt_node_ptr_t s2) {
  struct tt_dictset_intersection_context context;
  if (tt_dict_size(a, s1) < tt_dict_size(a, s2)) {
    tt_node_ptr_t tmp = s2;
    s2 = s1;
    s1 = tmp;
  }
  /* Now s1 is larger, s2 is smaller. */
  context.a = a;
  context.s1 = s1;
  context.result = TT_EMPTY_DICT; /* grab'd */
  tt_dict_foreach(a, s2, &context, tt_dictset_intersection_f);
  return context.result;
}

struct tt_dictset_difference_context {
  tt_arena_t *a;
  tt_node_ptr_t s2;
  tt_node_ptr_t result;
};

static int tt_dictset_difference_f1(void *context, tt_atom_t key, tt_node_ptr_t ignored_trie) {
  struct tt_dictset_difference_context *c = (struct tt_dictset_difference_context *) context;
  if (!TT_NO_PTR_P(tt_dict_get(c->a, c->s2, key))) {
    tt_node_ptr_t new_result = tt_grab(c->a, tt_dict_set(c->a, c->result, key, TT_EMPTY));
    tt_drop(c->a, c->result);
    c->result = new_result;
  }
  return !TT_NO_PTR_P(c->result);
}

static int tt_dictset_difference_f2(void *context, tt_atom_t key, tt_node_ptr_t ignored_trie) {
  struct tt_dictset_difference_context *c = (struct tt_dictset_difference_context *) context;
  tt_node_ptr_t new_result = tt_grab(c->a, tt_dict_remove(c->a, c->result, key));
  tt_drop(c->a, c->result);
  c->result = new_result;
  return !TT_NO_PTR_P(c->result);
}

tt_node_ptr_t tt_dictset_difference(tt_arena_t *a, tt_node_ptr_t s1, tt_node_ptr_t s2) {
  struct tt_dictset_difference_context context;
  context.a = a;
  context.s2 = s2;
  if (tt_dict_size(a, s1) < tt_dict_size(a, s2)) {
    context.result = TT_EMPTY_DICT;
    tt_dict_foreach(a, s1, &context, tt_dictset_difference_f1);
  } else {
    context.result = tt_grab(a, s1);
    tt_dict_foreach(a, s2, &context, tt_dictset_difference_f2);
  }
  return context.result;
}