#include "blocks.h"
#include "flow.h"
#include "optimise.h"
#define DEPTH (IR->x.depth)
#include "chains.h"
#include "nodes.h"

extern BitSet liveOnEdge(EdgeSet set);
extern void prepareMockAllocate(Block b);
extern int mockAllocate(Chain chain, Block b);

/** Need to remove variables from x-stack that are unreachable
and change across the block. If unreachable but don't change - then 
thats ok.
Unreachable variables are those for which (Machine reach - depth of p-stack) < stack location of variable.
*/
int matchAcrossBlock(Block b) {
    XStack x1, x2;
    int result = 0;
    int i1, i2, size1, size2, delve1, delve2;
    x1 = block_inStack(b);
    x2 = block_outStack(b);
    size1 = xstack_size(x1);
    size2 = xstack_size(x2);
    i1 = size1 - 1;
    i2 = size2 - 1;
    delve1 = DEPTH - block_getStartEDepth(b);
    delve2 = DEPTH - block_getEndEDepth(b);
    while (i1 >= 0 && i2 >= 0) {
        Symbol s1, s2;
        s1 = xstack_item(x1, i1);
        s2 = xstack_item(x2, i2);
        if (s1 != s2) {
            if (compareVirtuals(s1, s2) < 0) {
                if (i2 >= delve2) {
                    xstack_remove(x2, i2);
                    result = 1;   
                }
                --i2;  
            } else {
                if (i1 >= delve1) {
                    xstack_remove(x1, i1);
                    result = 1;
                }
                --i1;
            }
        } else {
            --i1;
            --i2;
        }
    } 
    while (i1 >= delve1) {
        xstack_remove(x1, i1); 
        --i1;
        result = 1;
    }
    while (i2 >= delve2) {
        xstack_remove(x2, i2); 
        --i2;          
        result = 1; 
    }
    return result;
}

void matchLowerStacks() {
   int notDone;
   int i, block_count = flow_size();
   do {
       notDone = 0;
       for(i = 0; i < block_count; i++) {
            notDone |=  matchAcrossBlock(flow_block(i));
       }
   } while(notDone);
}

static void logPropogatedCandidates() {
    int i;
    int block_count = graph_size(flow_graph);
    LOG(("Propogated candidates\n"));
    for(i = 0; i < block_count; i++) {
        LOG(("\nBlock %d rejects ", i)); 
        set_print(logfile, block_reject(flow_block(i)));
        LOG(("\n")); 
    }
}

void propogationOfXstackCandidates() {
    int i, j, k;
    int size = flow_size();
    Block p, c;
    BitSet from, to, live;
    for(k = 0; k < 8; k++) {
        for (i = 0; i < size; i++) {
            BitSet parents, children = flow_children(i);
            if (set_size(children) > 1) {
                p = flow_block(i);
                FOR_EACH(j, children) {
                    c = flow_block(j);
                    p = flow_block(i);
                    from = block_reject(p);
                    to = block_reject(c);
                    live = block_livein(c);
                    set_addSet(to, set_intersection(from, live));
                }
            }
            parents = flow_parents(i);
            if (set_size(parents) > 1) {
                FOR_EACH(j, parents) {
                    c = flow_block(i);
                    p = flow_block(j);
                    from = block_reject(c);
                    to = block_reject(p);
                    live = block_liveout(p);
                    set_addSet(to, set_intersection(from, live));
                }
            }
        }
    }
    if (logfile)
        logPropogatedCandidates();
}

static void chooser1(EdgeSet set) {
    int i, c, p;
    Block block;
    XStack x;
    BitSet all, children, parents, reachable;
    children = edgeset_children(set);
    parents = edgeset_parents(set);
    reachable = set_new(local_variable_count);
    all = set_new(local_variable_count);
    FOR_EACH(i, children) {
        block = flow_block(i);
        set_addSet(all, block_livein(block));
        set_addSet(reachable, block_def(block));
        set_addSet(reachable, block_use(block));
        
    }
    FOR_EACH(i, parents) {
        block = flow_block(i);
        set_addSet(all, block_liveout(block));
        set_addSet(reachable, block_def(block));
        set_addSet(reachable, block_use(block));
    }
    x = xstack_new(all);
    xstack_ensureReachable(x, reachable, IR->x.depth - edgeset_depth(set)); 
    FOR_EACH(c, children) {
        block_setInStack(flow_block(c), x);
    }
    FOR_EACH(p, parents) {
        block_setOutStack(flow_block(p), x);
    }
}

static void chooser2(EdgeSet set) {
    int i, c, p;
    Block block;
    XStack x;
    BitSet rejects, all, children, parents, reachable;
    reachable = set_new(local_variable_count);
    rejects = set_new(local_variable_count);
    all = set_new(local_variable_count);
    children = edgeset_children(set);
    parents = edgeset_parents(set);
    FOR_EACH(i, children) {
        block = flow_block(i);
        set_addSet(rejects, block_reject(block));
        set_addSet(all, block_livein(block));
        set_addSet(reachable, block_def(block));
        set_addSet(reachable, block_use(block));
    }
    FOR_EACH(i, parents) {
        block = flow_block(i);
        set_addSet(rejects, block_reject(block));
        set_addSet(all, block_liveout(block));
        set_addSet(reachable, block_def(block));
        set_addSet(reachable, block_use(block));
    }
    set_removeSet(all, rejects);
    x = xstack_new(all);
    xstack_ensureReachable(x, reachable, IR->x.depth - edgeset_depth(set)); 
    FOR_EACH(c, children) {
        block_setInStack(flow_block(c), x);
    }
    FOR_EACH(p, parents) {
        block_setOutStack(flow_block(p), x);
    }
}

static void initialiseXStacks() {
    int i, block_count = graph_size(flow_graph);
    XStack x_empty = xstack_new(set_new(local_variable_count));
    for(i = 0; i < block_count; i++) {
        Block b = flow_block(i);
        block_setInStack(b, x_empty);
        block_setOutStack(b, x_empty);
    }
}

void determineXStacks(void (* chooser)(EdgeSet)) {
    int i;
    ArrayList edge_sets = flow_edgeSets();
    int size = list_size(edge_sets);
    for(i = 0; i < size; i++) {
        chooser((EdgeSet)list_get(edge_sets, i));
    }
}

/** Comparison routine for ordering chains. Chains that are local to block have higher priority than those
that are not. If locality is equal then by refernce count, it thats equal, then by name.*/
static int chain_compare_locality_then_refCount(const void* v1, const void* v2) {
    const Chain* chain1 = v1;
    const Chain* chain2 = v2;    
    if (chain_livein(*chain1) || chain_liveout(*chain1)) {
        if (chain_livein(*chain2) || chain_liveout(*chain2)) {
            return compareVirtuals(chain_symbol(*chain1), chain_symbol(*chain2));
        } else {
            return 1;    
        }
    } else {
        if (chain_livein(*chain2) || chain_liveout(*chain2)) {
            return -1;
        } else {
            return compareVirtuals(chain_symbol(*chain1), chain_symbol(*chain2));
        }
    }
}

ArrayList* ordering;

void generateAndSortChains() {
    int i;
    Block b;
    ArrayList chains;
    ordering = allocate(sizeof (Chain*) * flow_size(), FUNC);
    for(i = 0; i < flow_size(); i++) {
        b = flow_block(i);
        chains = generateChains(b);
        list_sort(chains, chain_compare_locality_then_refCount);
        ordering[i] = chains;
    }    
}


/** Generates reject sets - solely live-in or out and ordered by ref count */
static void find_rejects(Block b) {
    Node p;
    int i, count;
    BitSet reject;
    ArrayList chains = ordering[block_id(b)];
    reject = block_reject(b);
    for(p = block_firstTemp(b); p; p = p->x.next) {
        p->x.scratch = getTotalDepth(p);
    }
    LOG(("Mock allocation for block %d\n", block_id(b)));
    for (i = 0, count = list_size(chains); i < count; i++) {
        Chain c = (Chain)list_get(chains, i);
        int ac = mockAllocate(c, b);
        if (chain_liveout(c) || chain_livein(c)) {
            int id = chain_symbol(c)->x.uniqueID;
            if (ac < 0)
                set_add(reject, id);
        }
    }
    LOG(("Block %d rejects ", block_id(b))); 
    if (logfile) set_print(logfile, reject);
    LOG(("\n")); 
}

static float* priority = NULL;

static int compareByPriority(const void* v1, const void* v2) {
    Block b1 = *(Block*)v1;
    Block b2 = *(Block*)v2;
    float f1, f2;
    f1 = priority[block_id(b1)];
    f2 = priority[block_id(b2)];
    if (f1 == f2)
        return block_length(b2) - block_length(b1);
    else
        return f1 > f2 ? -1 : 1;
}

void matchStacks(Block b, XStack x, int x_size, ArrayList l, int end) {
    int i1, i2, l_size, p_depth;
    BitSet liveout;
    Symbol t1, t2;
    Node p, forest, *tail;
    LOG(("Matching stacks at %s of block %d\n", end ? "end" : "start", block_id(b)));
    if (logfile) {
        int i;
        if (!end) {
            LOG(("XStack: "));
            xstack_print(x, logfile);
            LOG((" to depth %d\n", x_size));
        }
        LOG(("LStack:  ~ "));
        for (i = list_size(l) - 1; i > 0; i--)
            fprintf(logfile, "%s, ", ((Symbol)list_get(l, i))->name);
        if (i == 0)
            fprintf(logfile, "%s", ((Symbol)list_get(l, i))->name); 
        LOG(("\n"));
        if (end) {
            LOG(("XStack: "));
            xstack_print(x, logfile);
            LOG((" to depth %d\n", x_size));
        }
    }
    liveout = block_liveout(b);
    i1 = i2 = 0;
    forest = NULL;
    tail = &forest;
    p_depth = end ? block_getEndEDepth(b) : block_getStartEDepth(b);
//    x_size = xstack_size(x);
    l_size = list_size(l);
    while (i1 < x_size && i2 < l_size) {
        t1 = xstack_item(x, i1);
        t2 = (Symbol)list_get(l, i2);
        if (t1 == t2) {
            i1++;
            i2++;
        } else {
            if (compareVirtuals(t1 , t2) < 0) {
                // X symbol doesn't match...
                i1++;
                if (end) {
                    if (set_contains(liveout, t1->x.uniqueID)) {
                        p = loadTemp(t1, i1 + p_depth);
                    } else {
                        p = dummyToStack(t1, i1 + p_depth);  
                    }
                    p_depth--;
                    p->link = forest;
                    forest = p;
                } else {
                    p = storeTemp(t1, i1 + p_depth);
                    p_depth--;
                    *tail = p;
                    tail = &(p->link);
                }
             } else {
                // L symbol doesn't match...
                fprintf(stderr, "L stack symbol doesn't match X-stack\n");
                assert(0);
            }
        }
    }
    while (i1 < x_size) {
       t1 = xstack_item(x, i1);
       i1++;
       if (end) {
            if (set_contains(liveout, t1->x.uniqueID)) {
                p = loadTemp(t1, i1 + p_depth);
            } else {
                p = dummyToStack(t1, i1 + p_depth);  
            }
            p_depth--;
            p->link = forest;
            forest = p;
       } else {
            p = storeTemp(t1, i1 + p_depth);
            p_depth--;
            *tail = p;
            tail = &(p->link);
        }         
    }
    if (end) {
        if (forest)
            block_insertAtTail(b, forest);
    } else {
        if (forest)
            block_insertAtHead(b, forest);
    }
}

void lstackAllocation(Block b) {
    XStack start = block_inStack(b);
    XStack end = block_outStack(b);
    int i, l1, l2, x1, x2, count;
    BitSet live, def_use;
    ArrayList lstack1, lstack2, chains;
    l1 = x1 = xstack_size(start);
    l2 = x2 = xstack_size(end);
//    if (x1 == 0 && x2 == 0)
//        return;
    // Find upper limit of common part - Above this all variables must be handled by l-stack.
    do {
            l1--;
            l2--;
    } while (l1 >= 0 && l2 >= 0 && xstack_item(start, l1) == xstack_item(end, l2));
    l1++;
    l2++;
    def_use = set_union(block_use(b), block_def(b));
    for(i = l1; i < x1; i++) {
        if (set_contains(def_use, xstack_item(start, i)->x.uniqueID)) {
                l1 = i + 1;
                l2 = x2 - x1 + l1;
        }
    }
    // Get chains for this block
    chains = ordering[block_id(b)];
    lstack1 = list_new(4);
    lstack2 = list_new(4);  
    LOG(("Allocation for block %d\n", block_id(b)));
    for (i = 0, count = list_size(chains); i < count; i++) {
        Chain c = list_get(chains, i);
//        allocate chain c in block, maintaining l-stacks at either end...  
//        if (set_contains(common, chain_symbol(c)->x.uniqueID)) {
        if (((!chain_livein(c)) || xstack_contains(start, chain_symbol(c))) &&
            ((chain_liveout(c) != 1) || xstack_contains(end, chain_symbol(c)))) {
            chain_allocate(c, b, lstack1, lstack2);
        } else {         
              LOG(("Chain for %s not allocated\n", chain_symbol(c)->name));
 //           assert(set_contains(def_use, chain_symbol(c)->x.uniqueID));
        }
    }
    block_setStartLDepth(b, block_getStartLDepth(b) + list_size(lstack1));
    block_setEndLDepth(b, list_size(lstack2));
    matchStacks(b, start, l1, lstack1, 0);
    matchStacks(b, end, l2, lstack2, 1);
    live = block_livein(b);
    for (i = 0; i < x1; i++) {
        set_remove(live, xstack_item(start, i)->x.uniqueID);
    }
    live = block_liveout(b);
    for (i = 0; i < x2; i++) {
        set_remove(live, xstack_item(end, i)->x.uniqueID);
    }
}

static char* xxxx1(void) {
     initialiseXStacks();
    generateAndSortChains();
    determineXStacks(chooser1);    
    matchLowerStacks();
    flow_applyToBlocks(lstackAllocation);
    return NULL;
}
 
static char* xxxx2(void) {
    initialiseXStacks();
    generateAndSortChains();
    flow_applyToBlocks(find_rejects);
    propogationOfXstackCandidates();
    determineXStacks(chooser2);    
    matchLowerStacks();
    flow_applyToBlocks(lstackAllocation);
    return NULL;
}

// Global optimiser - invokes fast global algorithm
optimiser global1 = {
    "global1",
    "Global register allocation",
    12, // This is faster and simpler so choose this as default.
    &xxxx1
};

// Global optimiser - invokes current best global algorithm, although worse than global1 in some cases.
optimiser global2 = {
    "global2",
    "Variant global register allocation",
    0,
    &xxxx2
};