#include"stack.h"

#include"ast.h"
#include"ntc.h"
#include"scope.h"
#include<assert.h>
#include<stdlib.h>
#include"x86/arch.h"

struct Spill2StackState {
	AST *targetTLC;
	ScopeItem *target; // can be NULL
	size_t stackGrowth;
};
static void spill2stack_visitor(AST **aptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
	struct Spill2StackState *this = ud;
	
	if(tlc != this->targetTLC) {
		// Don't do anything.
		return;
	}
	
	AST *a = *aptr;
	
	if(a == tlc) {
		a->chunk.stackReservation += this->stackGrowth;
	} else if(a->nodeKind == AST_EXPR_VAR) {
		
		if(a->exprVar.thing == this->target) {
			// DO THE SPILL
			
			ASTExprStackPointer *rsp = calloc(1, sizeof(*rsp));
			rsp->nodeKind = AST_EXPR_STACK_POINTER;
			rsp->type = type_u(arch_gpr_size());
			
			ASTExprPrimitive *offset = calloc(1, sizeof(*offset));
			offset->nodeKind = AST_EXPR_PRIMITIVE;
			offset->type = rsp->type;
			offset->val = -this->stackGrowth; // This will be affected by the other part of this pass, so we must reverse
			offset->stackGrowth = true;
			
			ASTExprBinaryOp *bop = calloc(1, sizeof(*bop));
			bop->nodeKind = AST_EXPR_BINARY_OP;
			bop->type = rsp->type;
			bop->operator = BINOP_ADD;
			bop->operands[0] = (AST*) rsp;
			bop->operands[1] = (AST*) offset;
			
			ASTExprUnaryOp *deref = calloc(1, sizeof(*deref));
			deref->nodeKind = AST_EXPR_UNARY_OP;
			deref->type = a->expression.type;
			deref->operator = UNOP_DEREF;
			deref->operand = (AST*) bop;
			
			*aptr = (AST*) deref;
		}
		
	} else if(a->nodeKind == AST_EXPR_PRIMITIVE && a->exprPrim.stackGrowth) {
		
		// Guaranteed to not require more dumbification
		a->exprPrim.val += this->stackGrowth;
		
	}
}

void ast_spill_to_stack(AST *tlc, ScopeItem *vte) {
	assert(tlc->nodeKind == AST_CHUNK);
	assert(vte != NULL);
	assert(vte->kind == SCOPEITEM_VAR);
	
	if(ntc_get_int("pdbg")) {
		fprintf(stderr, "### SPILLING %s ###\n", vte->data.var.name);
	}
	
	struct Spill2StackState state;
	memset(&state, 0, sizeof(state));
	state.target = vte;
	state.targetTLC = tlc;
	state.stackGrowth = (type_size(vte->type) + 7) & ~7;
	
	generic_visitor(&tlc, NULL, NULL, tlc, tlc, &state, spill2stack_visitor, NULL);
	
	size_t vteIndex = 0;
	while(tlc->chunk.vars[vteIndex] != vte) vteIndex++;
	
	memmove(&tlc->chunk.vars[vteIndex], &tlc->chunk.vars[vteIndex + 1], sizeof(*tlc->chunk.vars) * (tlc->chunk.varCount - vteIndex - 1));
	tlc->chunk.varCount--;
}

void ast_grow_stack_frame(AST *tlc, size_t bytes) {
	if(bytes == 0) {
		return;
	}
	
	struct Spill2StackState state;
	memset(&state, 0, sizeof(state));
	state.target = NULL; // spill nothing
	state.targetTLC = tlc;
	state.stackGrowth = bytes;
	
	generic_visitor(&tlc, NULL, NULL, tlc, tlc, &state, spill2stack_visitor, NULL);
}