/* context-switch-pactice.cc
 * This program shows an example of how we can make two threads and
 * switch between them.  It uses the context_switch function, which
 * does a stack switch to another thread; this is implemented
 * in assembly because it breaks the standard rules for C/C++ functions
 * (mainly; changing to a different stack frame mid-function!).
 */

#include <iostream>

using namespace std;


// A thread is a stack plus the current top of the stack
typedef struct Thread {
    char stack[8192];
    char *saved_rsp;
} Thread;


/* This sets up a thread so that the first time it is switched to,
 * it "resumes" by starting to run the specified func.
 */
static const int kNumRegistersToSave = 6;
Thread create_thread(void (*func)()) {
    Thread new_thread;
    void *stack_top = new_thread.stack + sizeof(new_thread.stack);

    // Make it look like this thread was about to start func,
    // and then context switched.  In other words, once
    // the saved registers are popped, ret should take us to
    // the start of func.
    *(void **)((char *)stack_top - sizeof(void *)) = (void *)func;

    // Move the stack pointer downwards by num registers to make it appear like
    // we pushed those registers on when we previously context switched
    new_thread.saved_rsp = (char *)stack_top - sizeof(void *) - kNumRegistersToSave * sizeof(long);
    return new_thread;
}

/* This function context_switch is written in assembly!  It switches from running 'current'
 * to running 'next'.  To do this, it first pushes 'current's registers onto the stack
 * to save them, then updates 'current' with the new stack pointer.  Then it loads 'next's
 * stack pointer and from then on we are operating on 'next's stack, not current's!  So we
 * proceed to pop off all the registers we left on 'next's' stack from before, and then
 * return, and this will return back to the function that 'next' had been executing. 
 * 'current' will be resumed when another thread calls context_switch to switch back to
 * its stack.
 */
extern "C" void context_switch(Thread& current, Thread& next);

asm(R"(.text
	.global context_switch
context_switch:
	pushq %rbp
	pushq %rbx
	pushq %r12
	pushq %r13
	pushq %r14
	pushq %r15
	movq %rsp,0x2000(%rdi)	    # saved current's stack pointer
    movq 0x2000(%rsi),%rsp	    # load next's stack pointer
	popq %r15
	popq %r14
	popq %r13
	popq %r12
	popq %rbx
	popq %rbp
	ret     # return back to the function **this stack** was running
)");


Thread main_thread; // refers to "default thread" (running main)
Thread other_thread;

void other_func() {
    context_switch(other_thread, main_thread);
    cout << "D" << endl;
    context_switch(other_thread, main_thread);
    cout << "A" << endl;
}

int main(int argc, char *argv[]) {
    // Initialize other_thread to run other_func
    other_thread = create_thread(other_func);
    
    cout << "B" << endl;
    context_switch(main_thread, other_thread);
    cout << "C" << endl;
    context_switch(main_thread, other_thread);
    return 0;
}