/********** A really minimal coroutine package for C **********
 * By Armin Rigo
 */

#include "stacklet.h"

#include <stddef.h>
#include <assert.h>
#include <string.h>

/************************************************************
 * platform specific code
 */

/* The default stack direction is downwards, 0, but platforms
 * can redefine it to upwards growing, 1.
 */
#define STACK_DIRECTION 0   

#include "slp_platformselect.h"

#if STACK_DIRECTION != 0
#  error "review this whole code, which depends on STACK_DIRECTION==0 so far"
#endif

/************************************************************/

/* #define DEBUG_DUMP */

#ifdef DEBUG_DUMP
#include <stdio.h>
#endif

/************************************************************/

struct stacklet_s {
    /* The portion of the real stack claimed by this paused tealet. */
    char *stack_start;                /* the "near" end of the stack */
    char *stack_stop;                 /* the "far" end of the stack */

    /* The amount that has been saved away so far, just after this struct.
     * There is enough allocated space for 'stack_stop - stack_start'
     * bytes.
     */
    ptrdiff_t stack_saved;            /* the amount saved */

    /* Internally, some stacklets are arranged in a list, to handle lazy
     * saving of stacks: if the stacklet has a partially unsaved stack,
     * this points to the next stacklet with a partially unsaved stack,
     * creating a linked list with each stacklet's stack_stop higher
     * than the previous one.  The last entry in the list is always the
     * main stack.
     */
    struct stacklet_s *stack_prev;
};

void *(*_stacklet_switchstack)(void*(*)(void*, void*),
                               void*(*)(void*, void*), void*) = NULL;
void (*_stacklet_initialstub)(struct stacklet_thread_s *,
                              stacklet_run_fn, void *) = NULL;

struct stacklet_thread_s {
    struct stacklet_s *g_stack_chain_head;  /* NULL <=> running main */
    char *g_current_stack_stop;
    char *g_current_stack_marker;
    struct stacklet_s *g_source;
    struct stacklet_s *g_target;
};

/***************************************************************/

static void g_save(struct stacklet_s* g, char* stop
#ifdef DEBUG_DUMP
                   , int overwrite_stack_for_debug
#endif
                   )
{
    /* Save more of g's stack into the heap -- at least up to 'stop'

       In the picture below, the C stack is on the left, growing down,
       and the C heap on the right.  The area marked with xxx is the logical
       stack of the stacklet 'g'.  It can be half in the C stack (its older
       part), and half in the heap (its newer part).

       g->stack_stop |________|
                     |xxxxxxxx|
                     |xxx __ stop       .........
                     |xxxxxxxx|    ==>  :       :
                     |________|         :_______:
                     |        |         |xxxxxxx|
                     |        |         |xxxxxxx|
      g->stack_start |        |         |_______| g+1

     */
    ptrdiff_t sz1 = g->stack_saved;
    ptrdiff_t sz2 = stop - g->stack_start;
    assert(stop <= g->stack_stop);

    if (sz2 > sz1) {
        char *c = (char *)(g + 1);
#if STACK_DIRECTION == 0
        memcpy(c+sz1, g->stack_start+sz1, sz2-sz1);
#  ifdef DEBUG_DUMP
        if (overwrite_stack_for_debug)
          memset(g->stack_start+sz1, 0xdb, sz2-sz1);
#  endif
#else
        xxx;
#endif
        g->stack_saved = sz2;
    }
}

/* Allocate and store in 'g_source' a new stacklet, which has the C
 * stack from 'old_stack_pointer' to 'g_current_stack_stop'.  It is
 * initially completely unsaved, so it is attached to the head of the
 * chained list of 'stack_prev'.
 */
static int g_allocate_source_stacklet(void *old_stack_pointer,
                                      struct stacklet_thread_s *thrd)
{
    struct stacklet_s *stacklet;
    ptrdiff_t stack_size = (thrd->g_current_stack_stop -
                            (char *)old_stack_pointer);

    thrd->g_source = malloc(sizeof(struct stacklet_s) + stack_size);
    if (thrd->g_source == NULL)
        return -1;

    stacklet = thrd->g_source;
    stacklet->stack_start = old_stack_pointer;
    stacklet->stack_stop  = thrd->g_current_stack_stop;
    stacklet->stack_saved = 0;
    stacklet->stack_prev  = thrd->g_stack_chain_head;
    thrd->g_stack_chain_head = stacklet;
    return 0;
}

/* Save more of the C stack away, up to 'target_stop'.
 */
static void g_clear_stack(struct stacklet_s *g_target,
                          struct stacklet_thread_s *thrd)
{
    struct stacklet_s *current = thrd->g_stack_chain_head;
    char *target_stop = g_target->stack_stop;

    /* save and unlink tealets that are completely within
       the area to clear. */
    while (current != NULL && current->stack_stop <= target_stop) {
        struct stacklet_s *prev = current->stack_prev;
        current->stack_prev = NULL;
        if (current != g_target) {
            /* don't bother saving away g_target, because
               it would be immediately restored */
            g_save(current, current->stack_stop
#ifdef DEBUG_DUMP
                   , 1
#endif
                   );
        }
        current = prev;
    }

    /* save a partial stack */
    if (current != NULL && current->stack_start < target_stop)
        g_save(current, target_stop
#ifdef DEBUG_DUMP
               , 1
#endif
               );

    thrd->g_stack_chain_head = current;
}

/* This saves the current state in a new stacklet that gets stored in
 * 'g_source', and save away enough of the stack to allow a jump to
 * 'g_target'.
 */
static void *g_save_state(void *old_stack_pointer, void *rawthrd)
{
    struct stacklet_thread_s *thrd = (struct stacklet_thread_s *)rawthrd;
    if (g_allocate_source_stacklet(old_stack_pointer, thrd) < 0)
        return NULL;
    g_clear_stack(thrd->g_target, thrd);
    return thrd->g_target->stack_start;
}

/* This saves the current state in a new stacklet that gets stored in
 * 'g_source', but returns NULL, to not do any restoring yet.
 */
static void *g_initial_save_state(void *old_stack_pointer, void *rawthrd)
{
    struct stacklet_thread_s *thrd = (struct stacklet_thread_s *)rawthrd;
    if (g_allocate_source_stacklet(old_stack_pointer, thrd) == 0)
        g_save(thrd->g_source, thrd->g_current_stack_marker
#ifdef DEBUG_DUMP
               , 0
#endif
               );
    return NULL;
}

/* Save away enough of the stack to allow a jump to 'g_target'.
 */
static void *g_destroy_state(void *old_stack_pointer, void *rawthrd)
{
    struct stacklet_thread_s *thrd = (struct stacklet_thread_s *)rawthrd;
    thrd->g_source = EMPTY_STACKLET_HANDLE;
    g_clear_stack(thrd->g_target, thrd);
    return thrd->g_target->stack_start;
}

/* Restore the C stack by copying back from the heap in 'g_target',
 * and free 'g_target'.
 */
static void *g_restore_state(void *new_stack_pointer, void *rawthrd)
{
    /* Restore the heap copy back into the C stack */
    struct stacklet_thread_s *thrd = (struct stacklet_thread_s *)rawthrd;
    struct stacklet_s *g = thrd->g_target;
    ptrdiff_t stack_saved = g->stack_saved;

    assert(new_stack_pointer == g->stack_start);
#if STACK_DIRECTION == 0
    memcpy(g->stack_start, g+1, stack_saved);
#else
    memcpy(g->stack_start - stack_saved, g+1, stack_saved);
#endif
    thrd->g_current_stack_stop = g->stack_stop;
    free(g);
    return EMPTY_STACKLET_HANDLE;
}

static void g_initialstub(struct stacklet_thread_s *thrd,
                          stacklet_run_fn run, void *run_arg)
{
    struct stacklet_s *result;

    /* The following call returns twice! */
    result = (struct stacklet_s *) _stacklet_switchstack(g_initial_save_state,
                                                         g_restore_state,
                                                         thrd);
    if (result == NULL && thrd->g_source != NULL) {
        /* First time it returns.  Only g_initial_save_state() has run
           and has created 'g_source'.  Call run(). */
        thrd->g_current_stack_stop = thrd->g_current_stack_marker;
        result = run(thrd->g_source, run_arg);

        /* Then switch to 'result'. */
        thrd->g_target = result;
        _stacklet_switchstack(g_destroy_state, g_restore_state, thrd);

        assert(!"stacklet: we should not return here");
        abort();
    }
    /* The second time it returns. */
}

/************************************************************/

stacklet_thread_handle stacklet_newthread(void)
{
    struct stacklet_thread_s *thrd;

    if (_stacklet_switchstack == NULL) {
        /* set up the following global with an indirection, which is needed
           to prevent any inlining */
        _stacklet_initialstub = g_initialstub;
        _stacklet_switchstack = slp_switch;
    }

    thrd = malloc(sizeof(struct stacklet_thread_s));
    if (thrd != NULL)
        memset(thrd, 0, sizeof(struct stacklet_thread_s));
    return thrd;
}

void stacklet_deletethread(stacklet_thread_handle thrd)
{
    free(thrd);
}

stacklet_handle stacklet_new(stacklet_thread_handle thrd,
                             stacklet_run_fn run, void *run_arg)
{
    long stackmarker;
    assert((char *)NULL < (char *)&stackmarker);
    if (thrd->g_current_stack_stop <= (char *)&stackmarker)
        thrd->g_current_stack_stop = ((char *)&stackmarker) + 1;

    thrd->g_current_stack_marker = (char *)&stackmarker;
    _stacklet_initialstub(thrd, run, run_arg);
    return thrd->g_source;
}

stacklet_handle stacklet_switch(stacklet_thread_handle thrd,
                                stacklet_handle target)
{
    long stackmarker;
    if (thrd->g_current_stack_stop <= (char *)&stackmarker)
        thrd->g_current_stack_stop = ((char *)&stackmarker) + 1;

    thrd->g_target = target;
    _stacklet_switchstack(g_save_state, g_restore_state, thrd);
    return thrd->g_source;
}

void stacklet_destroy(stacklet_thread_handle thrd, stacklet_handle target)
{
    /* remove 'target' from the chained list 'unsaved_stack', if it is there */
    struct stacklet_s **pp = &thrd->g_stack_chain_head;
    for (; *pp != NULL; pp = &(*pp)->stack_prev)
        if (*pp == target) {
            *pp = target->stack_prev;
            break;
        }
    free(target);
}

char **_stacklet_translate_pointer(stacklet_handle context, char **ptr)
{
  char *p = (char *)ptr;
  long delta;
  if (context == NULL)
    return ptr;
  delta = p - context->stack_start;
  if (((unsigned long)delta) < ((unsigned long)context->stack_saved)) {
      /* a pointer to a saved away word */
      char *c = (char *)(context + 1);
      return (char **)(c + delta);
  }
  if (((unsigned long)delta) >=
      (unsigned long)(context->stack_stop - context->stack_start)) {
      /* out-of-stack pointer!  it's only ok if we are the main stacklet
         and we are reading past the end, because the main stacklet's
         stack stop is not exactly known. */
      assert(delta >= 0);
      assert(((long)context->stack_stop) & 1);
  }
  return ptr;
}
