// Copyright (c) 2025, Zak Fenton
// Zak Fenton's libc is licensed under the Mulan PSL v2. You can use this
// software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
// http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN “AS IS” BASIS, WITHOUT warranties of
// any kind, either express or implied, including but not limited to
// non-infringement, merchantability or fit for a particular purpose.
// See the Mulan PSL v2 for more details.

// NOTE: This is entirely my code, except for the hashmap code being 
// based on ChatGPT output. Don't worry, I'm using better AIs now, and
// will try to mark anywhere with AI generated code in case IP issues
// arise from them (or in case clients want pure AI-free code for policy
// reasons). Btw the AIs are pretty good at that kind of code
#ifndef SIMGC_H
#define SIMGC_H

#include <stddef.h>
#include <stdint.h>
#include <stdbool.h>

#ifndef SIMGC_SYSMALLOC
#include <stdlib.h>
#define SIMGC_SYSMALLOC malloc
#endif
#ifndef SIMGC_SYSFREE
#include <stdlib.h>
#define SIMGC_SYSFREE free
#endif
#ifndef SIMGC_LOCK_TYPE
#define SIMGC_LOCK_TYPE int
#define SIMGC_CREATELOCK(l) l = 1
#define SIMGC_DELETELOCK(l) l = 0
#define SIMGC_LOCK(l) l = 2
#define SIMGC_UNLOCK(l) l = 1
#endif

#ifndef SIMGC_GRANULARITY
#define SIMGC_GRANULARITY 8
#endif

typedef struct simgc simgc_t;
typedef struct simgc_ptrset simgc_ptrset_t;
typedef struct simgc_ptrset_node simgc_ptrset_node_t;
typedef struct simgc_allocation simgc_allocation_t;
typedef struct simgc_thread simgc_thread_t;

struct simgc {
  SIMGC_LOCK_TYPE alloclock; // Locking is not used properly yet but partly implemented as an example
  void* sdata; // Any data which may be needed by any complex SIMGC_SYSMALLOC/SIMGC_SYSFREE macros
  void* udata; // Any data which may be associated with the heap by program or libraries
  simgc_ptrset_t* pointers;
  simgc_allocation_t* allocations;
  simgc_thread_t* threads;
};

struct simgc_ptrset_node {
  void* key;
  simgc_ptrset_node_t* next;
};

struct simgc_ptrset {
  simgc_t* gc;
  simgc_ptrset_node_t** buckets;
  size_t capacity;
  size_t size;
};

struct simgc_allocation {
  simgc_allocation_t* next;
  size_t sizeflg; // Size/flags
  // NOTE: The data will start immediately after this struct
};

struct simgc_thread {
  void* sdata; // The system thread pointer, if applicable
  void* udata; // Userdata, pointer to info assigned by the program or libraries
  void* stackbase; // Base of stack (NOTE: may be numerically higher or lower than "top")
  void* stacktop; // Top of stack (NOTE: may be numerically higher or lower than "top")
  simgc_thread_t* next; // Always NULL for now (until I add locking etc.)
  simgc_t* gc;
  int iomode;
};

int simgc_init(simgc_t* gc);

// simgc_beginio should be called before entering a blocking I/O or pause operation
void simgc_beginio(simgc_t* gc);
// simgc_endio should be called after exiting a blocking I/O or pause operation
void simgc_endio(simgc_t* gc);

simgc_thread_t* simgc_begin_inner(simgc_t* gc, void* stackbase);
void simgc_end_inner(simgc_thread_t* ctx);
#define SIMGC_RUN(c,f) \
  do { \
    intptr_t dummyval = 123; \
    simgc_thread_t* ctx = simgc_begin_inner(c, &dummyval); \
    f; \
    simgc_end_inner(ctx); \
  } while(0)

void* simgc_alloc(simgc_t* gc, size_t sz);

// From ifndef SIMGC_H
#endif

#ifdef SIMGC_IMPLEMENTATION
#ifndef SIMGC_IMPLEMENTATION_ALREADY
#define SIMGC_IMPLEMENTATION_ALREADY

// stdio is only used for debugging
#include <stdio.h>

#define SIMGC_PTRSET_INITIAL_CAPACITY 16


static uint32_t simgc_ptrset_hash(void* ptr) {
  //return (uint32_t)(uintptr_t)ptr;
    uintptr_t val = (uintptr_t)ptr;
    return ((uint32_t)(val ^ (val >> 32))) & 0x7FFFFFFF;
}

simgc_ptrset_t* simgc_ptrset_create(simgc_t* gc) {
    simgc_ptrset_t* set = (simgc_ptrset_t*)SIMGC_SYSMALLOC(sizeof(simgc_ptrset_t));
    if (!set) {
        return NULL;
    }

    set->gc = gc;
    set->capacity = SIMGC_PTRSET_INITIAL_CAPACITY;
    set->size = 0; // TODO: No calloc!
    set->buckets = (simgc_ptrset_node_t**)calloc(set->capacity, sizeof(simgc_ptrset_node_t*));
    if (!set->buckets) {
      SIMGC_SYSFREE(set);
        return NULL;
    }

    return set;
}

void simgc_ptrset_free(simgc_ptrset_t* set) {
  simgc_t* gc = set->gc;
    for (size_t i = 0; i < set->capacity; ++i) {
        simgc_ptrset_node_t* node = set->buckets[i];
        while (node) {
            simgc_ptrset_node_t* next = node->next;
            SIMGC_SYSFREE(node);
            node = next;
        }
    }
    SIMGC_SYSFREE(set->buckets);
    SIMGC_SYSFREE(set);
}

static bool simgc_ptrset_resize(simgc_ptrset_t* set, size_t new_capacity) {
    simgc_ptrset_node_t** new_buckets = (simgc_ptrset_node_t**)calloc(new_capacity, sizeof(simgc_ptrset_node_t*));
    if (!new_buckets) {
        return false;
    }

    for (size_t i = 0; i < set->capacity; ++i) {
        simgc_ptrset_node_t* node = set->buckets[i];
        while (node) {
            simgc_ptrset_node_t* next = node->next;
            uint32_t hash = simgc_ptrset_hash(node->key);
            size_t index = hash % new_capacity;

            node->next = new_buckets[index];
            new_buckets[index] = node;
            node = next;
        }
    }

    SIMGC_SYSFREE(set->buckets);
    set->buckets = new_buckets;
    set->capacity = new_capacity;

    return true;
}

bool simgc_ptrset_add(simgc_ptrset_t* set, void* ptr) {
    if ((set->size + 1) * 2 > set->capacity) {
        if (!simgc_ptrset_resize(set, set->capacity * 2)) {
            return false;
        }
    }

    uint32_t hash = simgc_ptrset_hash(ptr);
    size_t index = ((uint64_t)hash) % ((uint64_t)(set->capacity));
    simgc_ptrset_node_t* node = set->buckets[index];

    while (node) {
        if (node->key == ptr) {
            return true;  // Already in the set
        }
        node = node->next;
    }

    simgc_ptrset_node_t* new_node = (simgc_ptrset_node_t*)SIMGC_SYSMALLOC(sizeof(simgc_ptrset_node_t));
    if (!new_node) {
        return false;
    }

    new_node->key = ptr;
    new_node->next = set->buckets[index];
    set->buckets[index] = new_node;
    set->size++;

    return true;
}

bool simgc_ptrset_remove(simgc_ptrset_t* set, void* ptr) {
    uint32_t hash = simgc_ptrset_hash(ptr);
    size_t index = hash % set->capacity;
    simgc_ptrset_node_t* node = set->buckets[index];
    simgc_ptrset_node_t* prev = NULL;

    while (node) {
        if (node->key == ptr) {
            if (prev) {
                prev->next = node->next;
            } else {
                set->buckets[index] = node->next;
            }
            SIMGC_SYSFREE(node);
            set->size--;
            return true;
        }
        prev = node;
        node = node->next;
    }

    return false;
}

bool simgc_ptrset_has(simgc_ptrset_t* set, void* ptr) {
  //fprintf(stderr, "LOOKING FOR %p\n", ptr);fflush(stderr);
    uint32_t hash = simgc_ptrset_hash(ptr);
    //fprintf(stderr, "HASH %x CAPACITY %llu\n", hash, set->capacity);fflush(stderr);
    size_t index = hash % set->capacity;
    //fprintf(stderr, "INDEX %p\n", index);fflush(stderr);
    simgc_ptrset_node_t* node = set->buckets[index];

    while (node) {
      //fprintf(stderr, "CHECKING %p\n", node);fflush(stderr);
        if (node->key == ptr) {
          //fprintf(stderr, "FOUND %p\n", ptr);fflush(stderr);
            return true;
        }
        node = node->next;
    }

  fprintf(stderr, "NOT FOUND %p\n", ptr);fflush(stderr);
    return false;
}

int simgc_init(simgc_t* gc) {
  SIMGC_CREATELOCK(gc->alloclock);
  gc->sdata = NULL;
  gc->udata = NULL;
  //gc->allocations = NULL;
  gc->pointers = simgc_ptrset_create(gc);
  gc->threads = NULL;
}

simgc_thread_t* simgc_begin_inner(simgc_t* gc, void* stackbase) {
  simgc_thread_t* t = SIMGC_SYSMALLOC(sizeof(simgc_thread_t));
  if (t == NULL) {
    return NULL;
  }
  t->sdata = NULL;
  t->udata = NULL;
  t->stackbase = stackbase;
  t->stacktop = NULL;
  t->gc = gc;
  t->iomode = 0;
  
  t->next = gc->threads;
  gc->threads = t;
  
  return t;
}

void simgc_end_inner(simgc_thread_t* ctx) {
  if (ctx == NULL) {
    return;
  }
  simgc_t* gc = ctx->gc;
  simgc_thread_t* prev = NULL;
  simgc_thread_t* thr = gc->threads;
  while (thr != NULL) {
    if (thr == ctx) {
      if (prev == NULL) {
        gc->threads = ctx->next;
      } else {
        prev->next = ctx->next;
      }
      SIMGC_SYSFREE(ctx);
      return;
    }
    thr = thr->next;
  }
}

void* simgc_alloc(simgc_t* gc, size_t sz) {
  if (gc == NULL) {
    return NULL;
  }
  while ((sz & SIMGC_GRANULARITY) != 0) sz++;
  simgc_allocation_t* alloc = SIMGC_SYSMALLOC(sizeof(simgc_allocation_t) + sz);
  if (alloc == NULL) {
    return NULL;
  }
  void* ptr = alloc+1;
  alloc->sizeflg = sz;
  alloc->next = gc->allocations;
  gc->allocations = alloc;
  
  simgc_ptrset_add(gc->pointers, ptr);
  return ptr;
}

simgc_allocation_t* simgc_lookup(simgc_t* gc, void* ptr) {
  if ((((uintptr_t)ptr) > 1000) && simgc_ptrset_has(gc->pointers, ptr)) {
    return ptr-sizeof(simgc_allocation_t);
  } else {
    return NULL;
  }
}

int simgc_recmark(simgc_t* gc, void* ptr);
int simgc_recmark(simgc_t* gc, void* ptr) {
  fprintf(stderr, "MARK %p\n", ptr);fflush(stderr);
  int n = 0;
  simgc_allocation_t* alloc = simgc_lookup(gc, ptr);
  if (alloc != NULL) {
    fprintf(stderr, "SPECULATIVE HIT %p\n", ptr);fflush(stderr);
  }
  if (alloc == NULL || alloc->sizeflg & 1) {
    fprintf(stderr, "MISS %p\n", ptr);fflush(stderr);
    return 0;
  }
  fprintf(stderr, "HIT %p\n", ptr);fflush(stderr);
  n = 1;
  alloc->sizeflg = alloc->sizeflg | 1;
  size_t sz = (alloc->sizeflg >> 3) << 3;
  size_t i;
  for (i = 0; i < sz; i += sizeof(void*)) {
    n += simgc_recmark(gc, *((void**)(ptr + i)));
  }
  return n;
}

int simgc_markstack(simgc_t* gc, void* stackbase, void* stacktop) {
  int n = 0;
  // The stack generally grows top-down, but here we simplify it
  if (((uintptr_t)stackbase) > ((uintptr_t)stacktop)) {
    void* x = stacktop;
    stacktop = stackbase;
    stackbase = x;
  }
  while ((((uintptr_t)stackbase)%SIMGC_GRANULARITY) != 0) stackbase++;
  while ((((uintptr_t)stacktop)%SIMGC_GRANULARITY) != 0) stacktop--;
  void* p = stackbase;
  while (((uintptr_t)p) < ((uintptr_t)stacktop)) {
    n += simgc_recmark(gc, *((void**)p));
    p += sizeof(void*);
  }
  return n;
}

int simgc_quicksweep(simgc_t* gc) {
  int n = 0;
  simgc_allocation_t* prev = NULL;
  simgc_allocation_t* alloc = gc->allocations;
  while (alloc != NULL) {
    if (alloc->sizeflg & 1) {
      alloc->sizeflg = (alloc->sizeflg >> 1) << 1; // Clear the "marked" bit
      prev = alloc;
      alloc = alloc->next;
    } else {
      simgc_allocation_t* next = alloc->next;
      if (prev == NULL) {
        gc->allocations = alloc->next;
      } else {
        prev->next = alloc->next;
      }
      simgc_ptrset_remove(gc->pointers, alloc);
      SIMGC_SYSFREE(alloc);
      alloc = next;
      n++;
    }
  }
  return n;
}

simgc_thread_t* simgc_threadbystack(simgc_t* gc, void* stacktop) {
  return gc->threads;
}

void simgc_beginio(simgc_t* gc) {
  simgc_thread_t* t;
  t = simgc_threadbystack(gc, &t);
  if (t == NULL) {
    return;
  }
  t->iomode = 1;
}

void simgc_endio(simgc_t* gc) {
  simgc_thread_t* t;
  t = simgc_threadbystack(gc, &t);
  if (t == NULL) {
    return;
  }
  t->iomode = 0;
}

int simgc_reclaim(simgc_t* gc) {
  intptr_t dummyval = 123;
  void* stacktop = &dummyval;
  simgc_thread_t* t = simgc_threadbystack(gc, stacktop);
  if (t == NULL) {
    return -1;
  }
  t->stacktop = stacktop;
  int n = simgc_markstack(gc, t->stackbase, t->stacktop);
  fprintf(stderr, "Marked %d\n", n);
  int o = simgc_quicksweep(gc);
  fprintf(stderr, "Swept %d\n", o);
  return o;
}

// From ifndef SIMGC_IMPLEMENTATION_ALREADY:
#endif
// From ifdef SIMGC_IMPLEMENTATION
#endif