#include "blocks.h"
#include "optimise.h"

static void doSet(EdgeSet set);
void popAll(Block b, int* temps, int count);
void pushAll(Block b, int* temps, int count);

static char* xxxx()
{
    // Get edge sets.
    EdgeSet* set;
    for (set = flow_edgeSets(); *set; set++) {
        doSet(*set);   
    }
    return NULL;
}

BitSet liveOnEdge(EdgeSet set) {
    BitSet children, parents;
    int i;
    Block b;
    children = edgeset_children(set);
    parents = edgeset_parents(set);
    i = FIRST_IN(children);
    b = flow_block(i);
    BitSet live = set_clone(block_livein(b));
    for (; i >= 0; i = set_next(children, i+1)) {
        b = flow_block(i);
        set_addSet(live, block_livein(b));
    }
    FOR_EACH(i, parents) {
        b = flow_block(i);
        set_addSet(live, block_liveout(b));
    }
    return live;
}

static void doSet(EdgeSet set) {
    BitSet children, parents;
    int i, j, *array;
    Block b;
    children = edgeset_children(set);
    parents = edgeset_parents(set);
    BitSet live = liveOnEdge(set);
    if (set_empty(live))
        return;
    array = allocate(sizeof(int) * set_size(live), FUNC);
    j = 0;
    FOR_EACH(i, live) {
        array[j++] = i;
    }
    assert(j == set_size(live));
    FOR_EACH(i, children) {
        b = flow_block(i);
        popAll(b , array, j);
    }
    FOR_EACH(i, parents) {
        b = flow_block(i);
        pushAll(b, array, j);
    }
}

static Node virtualNode(int op, Symbol s, int e_depth, int l_depth) {
    Node v = newnode(op, NULL, NULL, s);
    v->x.e_depth = e_depth;
    v->x.l_depth = l_depth;
    return v;
}

static Node TOS(Symbol s, int e_depth, int l_depth) {
    Node tos = virtualNode(STACK + P, s, e_depth, l_depth);
    setDelve(tos, 1);
    return tos;
}

static Node dropTop(int temp, int l_depth) {
    Symbol s = localVariableSymbol(temp);
    int sizeAndType = ttob(s->type);
    return newnode(INDIR + sizeAndType, TOS(s, 0, l_depth), NULL, NULL);
}

static Node storeTop(int temp, int l_depth) {
    Symbol s = localVariableSymbol(temp);
    int sizeAndType = ttob(s->type);
    Node new;
    new = newnode(sizeAndType + ASGN, NULL, NULL, NULL);
    new->kids[1] = virtualNode(VREG + P, s, 1, l_depth - 1);
    new->kids[1]->x.defn = 1;
    new->kids[0] = newnode(INDIR + sizeAndType, TOS(s, 0, l_depth), NULL, NULL);
    return new;
}

static Node storeToStack(Node value, int index, int depth) {
    Node new = newnode(value->op + ASGN - generic(value->op), NULL, NULL, NULL);
    new->kids[0] = value;
    Node stk = virtualNode(STACK + P, value->syms[0], 1, depth);
    stk->x.defn = 1;
     setDelve(stk, index);
    new->kids[1] = stk;
   return new;    
}

static Node zero(int index, int depth) {
    Symbol s = intconst(0);
    Node value = newnode(CNST + I + sizeop(4), NULL, NULL, s);
    return storeToStack(value, index, depth + 1);
}

static Node storeAtDepth(int temp, int index, int depth) {
    Symbol s = localVariableSymbol(temp);
    int sizeAndType = ttob(s->type);
    Node value = newnode(INDIR + sizeAndType, virtualNode(VREG + P, s, 0, depth), NULL, NULL);
    return storeToStack(value, index, depth);
}

void popAll(Block b, int* temps, int count) {
    BitSet defd, _block_use;
    Node forest = NULL;
    if (count == 0)
        return;
    int i, depth;
    defd = block_def(b);
    _block_use = block_use(b);
    depth = block_getStartLDepth(b);
    BitSet livein = block_livein(b);
    for (i = count - 1; i >= 0; i--) {
        depth++;
        Node ins;
        if (set_contains(livein, temps[i])) {
            ins = storeTop(temps[i], depth);
            set_remove(livein, temps[i]);
            set_add(defd, temps[i]);
            set_remove(_block_use, temps[i]);
        } else {
            // Not live so just drop...
            ins = dropTop(temps[i], depth);
        }
        assert(ins->link == NULL);
        ins->link = forest;
        forest = ins;
    }
    block_setStartLDepth(b, depth);
    block_insertAtHead(b, forest);
}

extern void increaseTreeDepth(Node n, int depth);

int block_useBy(Node n) {
    Node kid, x;
    int result, gop = generic(n->op);
    if (gop == STACK) {
        return 1;
    } else {
        if (kid = n->kids[0]) {
            result = block_useBy(n->kids[0]);
            if (kid = n->kids[1]) {
                result += block_useBy(n->kids[1]);
            }
            return result;
        } else {
            return 0;    
        }
    }
}

void pushAll(Block b, int* temps, int count) {
//    verify(b);
    Node forest, last;
    if (count == 0)
        return;
    int i, depth, initialDepth;
    switch(generic(block_lastNode(b)->op)) {
        Node prev;
    case LT: case JUMP: case NE: case EQ: case GT: case LE: case GE:
        depth = initialDepth = block_useBy(block_lastNode(b));
        break;
    default:
        depth = initialDepth = 0;
    }
    forest = last = NULL;
    BitSet liveout = block_liveout(b);
    for (i = count - 1; i >= 0; i--) {
        Node ins;
        if (set_contains(liveout, temps[i])) {
            ins = storeAtDepth(temps[i], initialDepth + 1, depth);
            set_remove(liveout, temps[i]);
        } else {
            // Not live so just insert rubbish.
            ins = zero(initialDepth + 1, depth);
        }
        if (last) {
            last->link = ins;
        } else {
            forest = ins;
        }
        last = ins;
        depth++; 
    }
    block_insertAtTail(b, forest);
}

optimiser optimist = {
    "optimist",
    "Schedules all local variables across block-edges, leaving intra-block optimisation to sort out the mess!",
    0, // Experimental only
    &xxxx
};