/* Compile with something like this:
   gcc testlightning.c `pkg-config --cflags --libs glib-2.0` -llightning -o testlightning
*/

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

/* for hash table */
#include <glib.h>

#include <lightning.h>

/* Replacement memory management functions */

static GHashTable *allocated_blocks = NULL;
static int block_cnt = 0;
static size_t bytes_allocated = 0;
static int verbose_alloc = 1;

#define ADD_BLOCK(p,s) g_hash_table_insert(allocated_blocks,p,(gpointer) s)
#define REMOVE_BLOCK(p) g_hash_table_remove(allocated_blocks,p)
#define ALLOCATED_BLOCK(p) g_hash_table_contains(allocated_blocks,p)
#define ALLOCATED_BLOCK_SIZE(p) ((size_t) g_hash_table_lookup(allocated_blocks,p))

static void init_alloc(void)
{
  if (allocated_blocks == NULL) {
    allocated_blocks = g_hash_table_new(&g_direct_hash, &g_direct_equal);
  }
}

typedef struct report_data_ {FILE *stream; size_t sum;} report_data_t;
  
static void printline_alloc(gpointer p, gpointer ps, report_data_t *data)
{
  fprintf(data->stream, "%p [%8u bytes]\n",p, (size_t) ps);
  data->sum += (size_t)ps;
}

static void report_alloc(char *msg)
{
  report_data_t data = {.stream = stderr,.sum = 0};
  if (block_cnt) {
    fprintf(stderr,"%s\nAllocated blocks (%d) are:\n",msg,block_cnt);
     g_hash_table_foreach(allocated_blocks, (GHFunc)printline_alloc, (gpointer)&data);
     fprintf(stderr,"Total size = %u [%u] bytes in %d blocks.\n",data.sum,bytes_allocated,block_cnt);
  } else
    fprintf(stderr,"%s\nNo allocated blocks.\n",msg);
}

static gpointer my_alloc(size_t size)
{
  gpointer p;

  p = malloc(size);
  if (verbose_alloc)
    fprintf(stderr,"ALLOC   %p[%8d]\n",p,size);

  if (p == NULL) {
    fprintf(stderr,"Out of memory (alloc(%d)).\n",size);
    report_alloc("(alloc)");
    exit(1);
  }
  assert(ADD_BLOCK(p,size));
  ++block_cnt;
  bytes_allocated += size;
  return p;
}

static gpointer my_realloc(gpointer p, size_t size)
{
  gpointer newp = NULL;

  if (!ALLOCATED_BLOCK(p)) {
    fprintf(stderr,"*** realloc: Pointer %p was not allocated ***\n",p);
    report_alloc("(realloc)");
  } else {
    bytes_allocated -= ALLOCATED_BLOCK_SIZE(p);
  }
  newp = realloc(p,size);

  if (newp == NULL) {
    fprintf(stderr,"Out of memory (realloc(%p,%d)).\n",p,size);
    report_alloc("(realloc)");
    exit(1);
  }
  if (verbose_alloc)
    fprintf(stderr,"REALLOC %p[%d] to %p[%d]\n",p,ALLOCATED_BLOCK_SIZE(p),newp,size);
  bytes_allocated += size;
  assert(REMOVE_BLOCK(p));
  assert(ADD_BLOCK(newp,size));
  return newp;
}

static void my_free(gpointer p)
{
  if (verbose_alloc)
    fprintf(stderr,"FREE    %p[%8d]\n",p,ALLOCATED_BLOCK_SIZE(p));
  if (!ALLOCATED_BLOCK(p)) {
    fprintf(stderr,"*** free: Pointer %p was not allocated ***\n",p);
    report_alloc("(free)");
  } else {
    bytes_allocated -= ALLOCATED_BLOCK_SIZE(p);
    assert(REMOVE_BLOCK(p));
    --block_cnt;
    free(p);
  }
}

static jit_state_t *_jit;

typedef void (*pvfi)(int);    /* Pointer to Void Function of Int */

int main(int argc, char *argv[])
{
  pvfi          myFunction;             /* ptr to generated code */
  jit_node_t    *start, *end;           /* a couple of labels */
  jit_node_t    *in;                    /* to get the argument */

  init_alloc();

  init_jit(argv[0]);

  jit_set_memory_functions(my_alloc,my_realloc,my_free);

  _jit = jit_new_state();

  report_alloc("(after calling jit_new_state)");
  
  start = jit_note(__FILE__, __LINE__);
  jit_prolog();
  in = jit_arg();
  jit_getarg(JIT_R1, in);
  jit_pushargi((jit_word_t)"generated %d bytes\n");
  jit_ellipsis();
  jit_pushargr(JIT_R1);
  jit_finishi(printf);
  jit_ret();
  jit_epilog();
  end = jit_note(__FILE__, __LINE__);

  report_alloc("(before calling jit_emit)");

  myFunction = jit_emit();

  report_alloc("(after jit_emit returned)");

  /* call the generated code, passing its size as argument */
  myFunction((char*)jit_address(end) - (char*)jit_address(start));
  jit_clear_state();

  jit_disassemble();

  jit_destroy_state();

  report_alloc("(after jit_destroy_state returned)");

  /* 2nd example */

  _jit = jit_new_state();

  report_alloc("(after calling jit_new_state (2nd time))");
  
  jit_prolog();
  jit_ret();
  jit_epilog();
  report_alloc("(before calling jit_emit)");

  myFunction = jit_emit();

  printf("2nd example:\n");
  jit_disassemble();

  report_alloc("(after jit_emit returned)");

  jit_destroy_state();

  report_alloc("(after jit_destroy_state returned)");

  finish_jit();

  report_alloc("(after finish_jit returned)");

  return 0;
}