#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

struct data {
  void *start;
  size_t len;
  size_t refcount;
};

struct node {
  struct data data;

  struct node *left;
  struct node *right;
};

static struct node *ROOT = 0;
static int TOTAL = 0;

static struct node *lookup_interval(struct node *root, void *ptr) {
  if (!root) return 0;

  if (ptr < root->data.start)
    return lookup_interval(root->left, ptr);

  if (ptr < (root->data.start + root->data.len))
    return root;

  // our last hope is something on the right
  return lookup_interval(root->right, ptr);
}

// exact lookup, ignoring range len
static struct node **lookup(void *ptr) {
  struct node **node = &ROOT;
  while (*node) {
    if ((*node)->data.start == ptr)
      return node;
    if (ptr > (*node)->data.start)
      node = &((*node)->right);
    else
      node = &((*node)->left);
  }
  return node;
}

static struct node *insert(void *ptr, size_t len) {
  struct node **loc = lookup(ptr);
  assert(!(*loc));
  *loc = calloc(1, sizeof(struct node));
  (*loc)->data.start = ptr;
  (*loc)->data.len = len;
  return *loc;
}

static void remove_root(struct node **tree) {
  if (!(*tree)) {
    return;
  } else if (!((*tree)->right)) {
    *tree = (*tree)->left;
  } else if (!((*tree)->left)) {
    *tree = (*tree)->right;
  } else {
    // swap *tree and (*tree)->right, then delete the root in tree->right
    struct data root_data = (*tree)->data;
    (*tree)->data = (*tree)->right->data;
    (*tree)->right->data = root_data;
    remove_root(&((*tree)->right));
  }
}

static struct node *remove_node(void *ptr) {
  remove_root(lookup(ptr));
}

void *refcount_calloc(unsigned long count, unsigned long size) {
  void *data = calloc(count, size);
  TOTAL += 1;
  printf("Allocated: %p ; now tracking %d regions\n", data, TOTAL);
  insert(data, count * size);
  return data;
}

void refcount_write(void *new, void *old) {
  // printf("Refcount write new: %p     old: %p\n", new, old);
  if (new) {
    struct node *new_node = lookup_interval(ROOT, new);
    if (new_node) {
      // printf("Incrementing refcount ...\n");
      new_node->data.refcount++;
    }
  }
  if (old) {
    struct node *old_node = lookup_interval(ROOT, old);
    if (old_node) {
      // printf("Decrementing refcount ...\n");
      old_node->data.refcount--;
      if (!(old_node->data.refcount)) {
        TOTAL -= 1;
        printf("Freeing! Left: %d\n", TOTAL);
        remove_node(old_node->data.start);
        free(old);
      }
    }
  }
}

void refcount_returning(void *ptr) {
  if (!ptr) return;
  struct node *ptr_node = lookup_interval(ROOT, ptr);
  if (!ptr_node) return;
  ptr_node->data.refcount--;
}