#include "../gc.h"
#include "../util.h"
#include <assert.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct object_header {
  bool mark : 1;
  bool is_compiled_function : 1;
  int rsvd : 1;
  uint16_t untraced_slot_count : 13;
  uintptr_t saved_tag : 3;
  uint16_t value_slot_count : 13;
  struct object_header *next;
  uintptr_t slots[];
};

static_assert(sizeof(struct object_header) == 2 * sizeof(uintptr_t),
              "sizeof(struct object_header) == 2 * sizeof(uintptr_t)");

/**
 * A singly-linked list of live objects.
 */
static struct object_header *live_objects = NULL;

/**
 * The number of bytes of live data, including object headers.
 */
static size_t live_bytes = 0;

/**
 * The number of bytes of live data to reach before collecting.
 */
static size_t limit_bytes = 80 * 1024 * 1024;

/**
 * The root stack is an arbitrary 4KiB. It shouldn't get anywhere near this
 * deep, so this is simply a convenient size.
 */
static struct value *root_stack[4096 / sizeof(uintptr_t)] = {0};
static size_t root_stack_depth = 0;

static struct object_header *hdr(void *ptr) {
  return &((struct object_header *)ptr)[-1];
}

static const struct object_header *hdrc(const void *ptr) {
  return &((const struct object_header *)ptr)[-1];
}

static size_t collect_amount[5] = {0};
static size_t collect_amount_i = 0;

void gc_init(void) {
  for (size_t i = 0; i < sizeof(collect_amount) / sizeof(collect_amount[0]);
       i++)
    collect_amount[i] = (size_t)-1;
}

static void gc_mark(const struct value initial_value) {
  // If the initial value wasn't a pointer to an unmarked object, we can bail
  // out quickly. (This makes our loop invariant simpler.)
  if (!is_ptr(initial_value))
    return;
  struct object *obj = untag_ptr(initial_value);
  if (!obj)
    return;
  struct object_header *header = hdr(obj);
  if (header->mark)
    return;

  // Otherwise, we do a traversal starting at the value, doing pointer reversal
  // to avoid recursion.
  //
  // In simple recursive marking, we store where we came from in the call
  // stack. Each stack frame effectively stores:
  //
  // 1. the pointer to the object we came from
  // 2. the index of the slot we came from in that object
  //
  // In pointer reversal, instead of using that space, we store a pointer to the
  // object we came from in a local variable, prev. We encode the index of the
  // slot we came from by replacing the tag in that slot with TAG_GC_INTERNAL.
  // The tag that was there before is saved in the header's saved_tag field.
  // This allows searching for the slot.
  //
  // From prev and saved_tag, we can reconstruct the value of the slot, so we
  // have the freedom to use the non-tag bits of the slot for whatever we want.
  // If, when we go into an object, we store the _old_ value of prev in that
  // slot, we have space to store the object we're coming from in prev. This
  // way, we don't need an auxiliary stack.
  //
  // We define the invariants in terms of the tricolor abstraction: the heap
  // starts white, objects become gray as we traverse them, and finally become
  // black after they've been traversed.
  //
  // - A white object does not have its mark bit set.
  // - A gray object has its mark bit set, and is either pointed to by obj, or
  //   is in the prev linked list.
  // - A black object has its mark bit set, but is not pointed to by obj, nor is
  //   it in the prev linked list.
  struct object *prev = NULL;
  for (;;) {
    // INVARIANT: obj points to a white object that we're about to make gray.
    // Since it's white, none of its fields should have the tag TAG_GC_INTERNAL,
    // so it shouldn't have a saved tag either.
    assume(!header->mark);
    assume(header->saved_tag == 0);

    // Mark the object, making it gray.
    header->mark = true;

    // Search for a field containing a pointer to a white object.
    size_t i;
    struct value *slot;
    struct object *slot_ptr;
    for (i = 0; i < header->value_slot_count; i++) {
      slot = (struct value *)&header->slots[i];
      if (!is_ptr(*slot))
        continue;
      slot_ptr = untag_ptr(*slot);
      if (!slot_ptr)
        continue;
      if (hdrc(slot_ptr)->mark)
        continue;
      break;
    }

    // If we found one, we traverse into it.
    if (i != header->value_slot_count) {
      header->saved_tag = get_tag(*slot);
      *slot = tag_ptr(prev, TAG_GC_INTERNAL);
      prev = obj;
      obj = slot_ptr;
      header = hdr(obj);
      continue;
    }

    // If didn't find one, we can make this object black by popping the stack.
    // We need to continue popping the stack until we have another gray object
    // in obj.
    for (;;) {
      // INVARIANT: The object is gray, but is about to become black.

      // If this was the object we started traversal with, we're actually done!
      if (!prev)
        return;

      // Compute the value to restore into the slot.
      header = hdr(prev);
      struct value slot_value = tag_ptr(obj, header->saved_tag);

      // Move prev into obj.
      obj = prev;

      // Find the slot of obj with TAG_GC_INTERNAL.
      for (i = 0; i < header->value_slot_count; i++) {
        slot = (struct value *)&header->slots[i];
        // This is get_tag, but that function doesn't allow TAG_GC_INTERNAL.
        if ((slot->bits & 0b111) == TAG_GC_INTERNAL)
          break;
      }
      assume(i != header->value_slot_count);

      // Use the value stored there to restore the old value of prev, and write
      // its old value back to it. This is effectively what's in untag_ptr, but
      // that function doesn't allow TAG_GC_INTERNAL.
      prev = (struct object *)((uintptr_t)slot->bits & ~0b111);
      *slot = slot_value;

      // The old obj is now black. We need to determine if the new obj has any
      // more references to white objects, or if instead it should become black
      // too.
      for (i += 1; i < header->value_slot_count; i++) {
        slot = (struct value *)&header->slots[i];
        if (!is_ptr(*slot))
          continue;
        slot_ptr = untag_ptr(*slot);
        if (!slot_ptr)
          continue;
        if (hdrc(slot_ptr)->mark)
          continue;
        break;
      }

      // If the object had no more references to white objects, we can make it
      // black with this same loop.
      if (i == header->value_slot_count)
        continue;

      // Otherwise, we traverse into the object we found.
      header->saved_tag = get_tag(*slot);
      *slot = tag_ptr(prev, TAG_GC_INTERNAL);
      prev = obj;
      obj = slot_ptr;
      header = hdr(obj);
      break;
    }
  }
}

static void gc_sweep(void) {
  struct object_header **head = &live_objects, *header;
  while ((header = *head)) {
    if ((*head)->mark) {
      (*head)->mark = false;
      (*head)->saved_tag = 0;
      head = &header->next;
    } else {
      *head = header->next;
      const size_t total_bytes =
          sizeof(struct object_header) +
          (header->value_slot_count + header->untraced_slot_count) *
              sizeof(uintptr_t);
      free(header);
      live_bytes -= total_bytes;
    }
  }
}

static const char *string_of_gc_collect_reason(enum gc_collect_reason reason) {
  switch (reason) {
  case GC_COLLECT_MANUAL:
    return "GC_COLLECT_MANUAL";
  case GC_COLLECT_LIMIT:
    return "GC_COLLECT_LIMIT";
  case GC_COLLECT_OOM:
    return "GC_COLLECT_OOM";
  default:
    return "???";
  }
}

void gc_collect(enum gc_collect_reason reason) {
  debugf("gc_collect(%s)... ", string_of_gc_collect_reason(reason));
  for (size_t i = 0; i < root_stack_depth; i++)
    gc_mark(*root_stack[i]);
  const size_t before = live_bytes;
  gc_sweep();
  const size_t after = live_bytes;
  fprintf(stderr, "collected %zu bytes\n", before - after);

  collect_amount[collect_amount_i++] = before - after;
  const size_t collect_amount_count =
      sizeof(collect_amount) / sizeof(collect_amount[0]);
  collect_amount_i %= collect_amount_count;
  size_t i;
  for (i = 0; i < collect_amount_count; i++)
    if (collect_amount[i] > before / 50)
      break;
  if (i == collect_amount_count)
    todo("OOM: GC was not productive");
}

struct object *gc_alloc(size_t value_slot_count, size_t untraced_slot_count) {
  assume(value_slot_count < (1 << 13));
  assume(untraced_slot_count < (1 << 13));

  const size_t total_bytes =
      sizeof(struct object_header) +
      (value_slot_count + untraced_slot_count) * sizeof(uintptr_t);
  if ((live_bytes += total_bytes) >= limit_bytes)
    gc_collect(GC_COLLECT_LIMIT);
  struct object_header *header = malloc(total_bytes);
  if (!header) {
    gc_collect(GC_COLLECT_OOM);
    header = malloc(total_bytes);
    if (!header)
      todo("OOM");
  }
  memset(header, 0, total_bytes);
  header->value_slot_count = value_slot_count;
  header->untraced_slot_count = untraced_slot_count;
  header->next = live_objects;
  live_objects = header;
  return (struct object *)&header->slots[0];
}

struct object *gc_alloc_compiled_function(size_t value_slot_count,
                                          size_t untraced_slot_count) {
  struct object *obj = gc_alloc(value_slot_count, untraced_slot_count);
  hdr(obj)->is_compiled_function = true;
  return obj;
}

struct object *gc_alloc_hashtable_eq(size_t value_slot_count,
                                     size_t untraced_slot_count) {
  return gc_alloc(value_slot_count, untraced_slot_count);
}

struct value gc_read_value_slot(const struct object *obj, size_t slot_index) {
  assume(slot_index < hdrc(obj)->value_slot_count);
  struct value *slots = (struct value *)obj;
  struct value value = slots[slot_index];
  return value;
}

uintptr_t gc_read_untraced_slot(const struct object *obj, size_t slot_index) {
  assume(slot_index < hdrc(obj)->untraced_slot_count);
  todo("gc_read_untraced_slot");
}

uint8_t gc_read_untraced_byte(const struct object *obj, size_t byte_index) {
  todo("gc_read_untraced_byte");
}

void gc_write_value_slot(struct object *obj, size_t slot_index,
                         struct value value) {
  struct value *slots = (struct value *)obj;
  slots[slot_index] = value;
}

void gc_write_untraced_slot(struct object *obj, size_t slot_index,
                            uintptr_t value) {
  assume(slot_index < hdrc(obj)->untraced_slot_count);
  todo("gc_write_untraced_slot");
}

void gc_write_untraced_byte(struct object *obj, size_t byte_index,
                            uint8_t byte) {
  todo("gc_write_untraced_byte");
}

void gc_root_push(struct value *root) {
  assume(root_stack_depth < sizeof(root_stack));
  root_stack[root_stack_depth++] = root;
}

void gc_root_pop(void) {
  assume(root_stack_depth > 0);
  root_stack_depth--;
}

void gc_debug(void) { fprintf(stderr, "live bytes = %zu\n", live_bytes); }