#include "blocks.h"
#include "flow.h"
#include "optimise.h"
#define stacksize(node) (opsize(node->op) == 8 ? 2 : 1)
#define MAX(x,y) ((x) > (y) ? (x) : (y))

static void permuteTree(Node tree);
static void annotateTree(Node node, int height);
//static void labelBlock(Node forest);
extern optimiser cleanup;
extern void liveness(); 
void processForest(Node forest);
static void printnode(Node p, FILE* out);
static void printsubtree(Node p, char* indent, FILE* out);
FILE* logfile = 0;
extern int do_verify;

static int dump_graph = 0;

int framesize;
int offset;

void write_log(char* format, ...) {
    if (logfile) {
        va_list ap;
        va_start(ap, format);   
        vfprintf(logfile, format, ap);
        va_end(ap);
    } 
}

static ArrayList local_variable_list = NULL;

void makeTemporary(Symbol p) {
    p->x.name = p->name;
    p->sclass = REGISTER;
    p->x.uniqueID = local_variable_count;
    if (local_variable_list == NULL) {
        local_variable_list = list_new(16);
    }
    list_add(local_variable_list, p);
    local_variable_count++;
}

void mkauto(Symbol p) {
    assert(p->sclass == AUTO);
    offset = roundup(offset + p->type->size, p->type->align);
    p->x.offset = -offset;
    p->x.name = stringd(-offset);
}

#define ck(i) return (i) ? 0 : LBURG_MAX

int range(Node p, int lo, int hi) {
    Symbol s = p->syms[0];
    Symbol s1;
    switch (specific(p->op)) {
    case ADDRL+P: ck(s->x.offset >= lo && s->x.offset <= hi);
    case CNST+I:
        ck(s->u.c.v.i  >= lo && s->u.c.v.i  <= hi);
    case CNST+U:  ck(s->u.c.v.u  >= lo && s->u.c.v.u  <= hi);
    case CNST+P:  ck(s->u.c.v.p  == 0  && lo <= 0 && hi >= 0);
    case STACK+P: case COPY+P: 
    case TUCK+I: case TUCK+U: case TUCK+F: case TUCK+P:  
        s1 = p->syms[1];
        ck(s1->u.c.v.i  >= lo && s1->u.c.v.i  <= hi);
    }
    return LBURG_MAX;
}

static Node duplicate(Node n) {
    if (n) { 
        Node new = newnode(n->op, duplicate(n->kids[0]), duplicate(n->kids[1]), n->syms[0]);
        new->syms[1] = n->syms[1];
        return new;
    } else {
        return NULL;
    }
}

static int hasVREG(Node n) {
    Node kid1, kid2;
    if (n->op == VREG + P)
        return 1;
    if (kid1 = n->kids[0]) {
        return hasVREG(kid1) || ((kid2 = n->kids[0]) && hasVREG(kid2));
    } else {
        return 0;   
    }
}

static void split_off_virtuals(Node p, Node prev) {
    if (p->op == VREG + P) {
        Type t = p->syms[0]->type;
        Symbol temp = newtemp(AUTO, ttob(t), t->size);
        Node tempNode = newnode(VREG + P, NULL, NULL, temp);
        Node old = newnode(VREG + P, NULL, NULL, p->syms[0]);
        Node indir = newnode(INDIR + ttob(t), old, NULL, NULL);
        Node asgn = newnode(ASGN + ttob(t), indir, tempNode, intconst( temp->type->size));
        temp->x.branchTemp = 1;
        asgn->syms[1] = intconst(t->align);
        asgn->link = prev->link;
        prev->link = asgn;
        p->syms[0] = temp;
    } else {
        Node kid;
        if (kid = p->kids[0]) {
            split_off_virtuals(kid, prev);
            if (kid = p->kids[1]) {
                split_off_virtuals(kid, prev);
            }        
        }
    }
    
}

static void split_if_complex(Node* kid, Node prev) {
    int op = generic((*kid)->op);
    Symbol temp;
    Node tempNode, asgn ;
    temp = newtemp(AUTO, optype((*kid)->op), opsize((*kid)->op));
    tempNode = newnode(VREG + P, NULL, NULL, temp);
    if ((!hasVREG(*kid)) && (IR->x.recalcNode(tempNode, *kid) == *kid))
        return;
    if (IR->x.makeVirtual(temp)) {
        Node indir = newnode(INDIR + (*kid)->op - op, tempNode, NULL, NULL);
        temp->x.branchTemp = 1;
        asgn = newnode(ASGN + (*kid)->op - op, NULL, NULL, intconst( temp->type->size));
        asgn->kids[1] = newnode(VREG + P, NULL, NULL, temp);
        asgn->kids[0] = *kid;
        asgn->syms[1] = intconst(temp->type->align);
        *kid = indir;
        asgn->link = prev->link;
        prev->link = asgn;
    } else {
        split_off_virtuals(*kid, prev);
    }
}    
        
static void flattenBranch(Node test, Node prev) {
    if (test->kids[1])
        split_if_complex(&(test->kids[1]), prev);
    split_if_complex(&(test->kids[0]), prev);
}

static void flattenForest(Node forest) {
    Node prev = NULL;
    Node p, test;
    for (p = forest; p; p = p->link) {
        switch(generic(p->op)) {
        case JUMP:
            if (generic(p->kids[0]->op) == ADDRG) {
                 prev = NULL;
                 break;  
            }
        case EQ: case GE: case LT: case GT: case NE: case LE: 
            if (prev == NULL) {
                test = duplicate(p);
                test->syms[1] = p->syms[1];
                test->link = p->link;
                p->link = test;
                p->op = NOOP;
                p->kids[0] = NULL;
                p->kids[1] = NULL;
                p->syms[0] = NULL;
                p->syms[1] = NULL;
                flattenBranch(test, p);
                p = test;
            } else {
                flattenBranch(p, prev);
            }
        case RET:
            prev = NULL;
            break;
        default:
            prev = p;
        }
    }
}

static Node eliminateCSE(Node p) {
    if (p == NULL)
        return NULL;
    if (generic(p->op) == INDIR && p->kids[0]->op == VREG+P) {
        Symbol r = p->kids[0]->syms[0];  
        if (r->u.t.cse) {
            Node n = (IR->x.recalcNode)(p, r->u.t.cse);
            assert(n == p || n == r->u.t.cse);
            return n == p ? n : duplicate(n);
        }
    }
    p->kids[0] = eliminateCSE(p->kids[0]);
    p->kids[1] = eliminateCSE(p->kids[1]);
    return p;
}

static void prepare(Node p) {
    Symbol s;
    if (p == NULL)
        return;
    prepare(p->kids[0]);
    prepare(p->kids[1]);
    s = p->syms[0];
    switch (generic(p->op)) {
    case ADDRF:
        if (s->scope == PARAM)
            break;
        p->op += ADDRL - ADDRF;
    case ADDRL:
        if (s->sclass == REGISTER)
            p->op = VREG+P;
        break;
    }
}

static void swapkids(Node node) { 
    Node swap = node->kids[0]; 
    assert(node->kids[0]);
    assert(node->kids[1]);
    node->kids[0] = node->kids[1]; 
    node->kids[1] = swap; 
}

static void permuteTree(Node tree) {
    if (tree == NULL)
        return;
    switch (generic(tree->op)) {
    case ASGN:
        swapkids(tree);
        break;
    case SUB:
        if (IR->x.rhs.sub) {
            swapkids(tree);
        }
        break;
    case LSH:
        if (IR->x.rhs.lsh) {
            swapkids(tree);
        }
        break;
    case RSH:
        if (IR->x.rhs.rsh) {
            swapkids(tree);
        }
        break;
    case MOD:
        if (IR->x.rhs.mod) {
            swapkids(tree);
        }
        break;
    case DIV:
        if (IR->x.rhs.div) {
            swapkids(tree);
        }
        break;
    case LT:
        if (IR->x.rhs.lt) {
            swapkids(tree);
        }
        break;
    case LE:
        if (IR->x.rhs.le) {
            swapkids(tree);
        }
        break;
    case GT:
        if (IR->x.rhs.gt) {
            swapkids(tree);
        }
        break;
    case GE:
        if (IR->x.rhs.ge) {
            swapkids(tree);
        }
        break;
    case ADD: case MUL: case EQ: case NE: case BAND: case BXOR: case BOR:
        if (specific(tree->kids[0]->op) == CNST + I) {
            swapkids(tree);
        }
    }
    permuteTree(tree->kids[0]);
    permuteTree(tree->kids[1]);
}

static Node lastTemp;

static void annotate(Node forest) {
    Node p;
    int height = 0;
    lastTemp = NULL;
    for (p = forest; p; p = p->link) {
        switch (generic(p->op)) {
        case EQ: case GE: case LT: case GT: case NE: case LE: 
            annotateTree(p, height);
            break;
        case JUMP:
            annotateTree(p, height);
            break;    
        case ARG:
            annotateTree(p, height);
            height += stacksize(p);
            break;
        case LABEL: case NOOP:
            break;
        case CALL:
            annotateTree(p, height);
            height = 0;
            break;
        case RET:
            assert(height == 0);
            annotateTree(p, height);
            break;
        case ASGN:
            if(generic(p->kids[0]->op) == CALL) {
                annotateTree(p->kids[0], height);
                height = stacksize(p->kids[0]);
                annotateTree(p->kids[1], height);
                height = 0;
            } else if(generic(p->kids[1]->op) == CALL) {
                annotateTree(p->kids[1], height);
                height = stacksize(p->kids[1]);
                annotateTree(p->kids[0], height);
                height = 0;
            } else if (generic(p->kids[0]->op) == VREG) {
                annotateTree(p->kids[1], height);
                annotateTree(p->kids[0], height + stacksize(p->kids[0]));
            } else {
                annotateTree(p->kids[0], height);
                annotateTree(p->kids[1], height + stacksize(p->kids[0]));    
            }
            break;
        default:
            assert(0);
        }
    }
}

int getEDepth(Node n) {
    assert(generic(n->op) == VREG || generic(n->op) == COPY || generic(n->op) == STACK);
    return n->x.e_depth;    
}

int getTotalDepth(Node n) {
    assert(generic(n->op) == VREG || generic(n->op) == COPY || generic(n->op) == STACK);
    return n->x.l_depth + n->x.e_depth;    
}

int getDelve(Node n) {
    return n->syms[1]->u.c.v.i;    
}

static Symbol ICONSTS[] = { NULL, NULL, NULL, NULL, NULL };

void setDelve(Node n, int delve) {
    assert(delve >= 0);
//    assert(delve <= IR->x.depth);
    if (delve < 5) {
        if (ICONSTS[delve] == NULL)
            ICONSTS[delve] = intconst(delve);
        n->syms[1] = ICONSTS[delve];
    } else {
        n->syms[1] = intconst(delve);
    }
}

static void annotateTree(Node node, int height) {
    int generic = generic(node->op);
    if (generic == VREG) {
        node->x.e_depth = height;
    } else if (node->kids[0]) {
        annotateTree(node->kids[0], height);
        if (node->kids[1]) {
            annotateTree(node->kids[1], height + stacksize(node->kids[0]));
        }
    }
}

static int fallThrough = -1;
static int table_jump = -1;

int local_variable_count;

Symbol localVariableSymbol(int n) {
    return list_get(local_variable_list, n);   
}

void insertBranchEdges();

extern optimiser dumpflow;

void optimise(void) {
    int i;
    char* msg;
    insertBranchEdges();
    flow_close();
    liveness();
    integrity(flow_graph);
    verify("init");
    if (dump_graph) 
        flow_print(stringf("%s/%s.init", TEMPDIR, cfunc->name));
    for(i = 0; i < optimiser_count; i++) {      
        if (optimisersOn[i] > 0) {
            write_log("Calling %s\n", optimisers[i]->name);
            msg = optimisers[i]->function();
            verify(optimisers[i]->name);
            if (dump_graph) 
                flow_print(stringf("%s/%s.%s", TEMPDIR, cfunc->name, optimisers[i]->name));
            if (msg) {
                fatal(optimisers[i]->name, msg, 0);
            }
        }
    }
    fallThrough = -1;
    local_variable_list = NULL;
}

static void setOptimiser(char* name, int on) {
    int j;
    for(j = 0; j < optimiser_count; j++) {
        if (strcmp(name, optimisers[j]->name) == 0) {
            optimisersOn[j] = on;
            return;
        }
    }
    fprintf(stderr, "Cannot find optimiser %s\n", name);
    exit(1);
}

void backendflags(int argc, char *argv[]) {
    int i, j, level = -1;
    for(j = 0; j < optimiser_count;j++) {
        optimisersOn[j] = 0;
    }
    for (i = 1; i < argc; i++) {
        char* arg = argv[i];
        if (strncmp(arg, "-log", 4) == 0) {
            if (arg[4] == '\0') {
                logfile = fopen("lcc.log", "w");
            } else {
                logfile = fopen(arg + 4, "w");
            }
        } else if (strncmp(arg, "-graph", 6) == 0) {
            dump_graph = 1;
        } else if (strncmp(arg, "-verify", 7) == 0) {
            do_verify = 1;
        } else if (strncmp(arg, "-O", 2) == 0) {
            if (arg[3]) {
                if (level == -1)
                    level = 0;
                setOptimiser(arg + 2, 1);
            } else {
                level = arg[2] - '0';
                if (level < 0 || level > 3) {
                    fprintf(stderr, "Illegal optimisation level: %s\n", arg + 2);
                    exit(1);
                }
            }
        } else if (strncmp(arg, "-X", 2) == 0) {
            setOptimiser(arg + 2, -1);
        }
    }
    /* Set optimisation level to default level of two if no -O flag */
    if (level == -1)
        level = 2;
    for(j = 0; j < optimiser_count;j++) {
        if ((optimisers[j]->level & (1 << level)) && (optimisersOn[j] == 0)) {
           optimisersOn[j] = 1;
        }
    }
    for(j = 0; j < optimiser_count;j++) {
        if (optimisersOn[j] == 1)
            write_log("Optimiser %s on.\n", optimisers[j]->name);
    }
}

void encode() {
//    flow_applyToBlocks(&encodeBlock);
}

Node popParam(Node destination, int sizeAndType, Node link, int depth) {
    Node n, tos = newnode(STACK + P, NULL, NULL, destination->syms[0]);
    setDelve(tos, 1);
    tos->x.l_depth = depth;
    tos = newnode(INDIR + sizeAndType, tos, NULL, NULL);
    n = newnode(ASGN + sizeAndType, tos, destination, NULL);
    n->link = link;
    assert(n->count == 0);
    return n;    
}

Node stackParameters(Symbol caller[], Symbol callee[]) {
    Node last;
    Block b;
    Node intro = NULL;
    int i, sizeAndType, depth = 0;
    for (i = 0; callee[i]; i++) {
        Node leftNode;
        Symbol p = callee[i];
        Symbol q = caller[i];
        p->scope = LOCAL;
        assert(q);
        sizeAndType = ttob(p->type);
        depth += roundup(p->type->size, IR->ptrmetric.size) / IR->ptrmetric.size;
        if (IR->x.makeVirtual(p)) {
            leftNode = newnode(VREG + P, NULL, NULL, p);
            leftNode->x.l_depth = depth-1;
            leftNode->x.e_depth = 1;
            q->sclass = REGISTER;
            q->type = p->type;
        } else {
            mkauto(p);
            leftNode = newnode(ADDRL + P + sizeop(IR->ptrmetric.size), NULL, NULL, p);
        }
        intro = popParam(leftNode, sizeAndType, intro, depth);
    }
    assert(!caller[i]);
    flow_graph = graph_new();
    if (!intro) {
        intro = newnode(NOOP, NULL, NULL, NULL);
    }
    for (last = intro; last->link; last = last->link);
    b = block_new(0, 0, intro, last);
    block_setStartLDepth(b, depth);
    fallThrough = flow_addBlock(b);
    return intro;
}

Node stackgen(Node forest) {
    Node p;
    assert(forest);
    for (p = forest; p; p = p->link) {
        p->x.listed = 1;
        prepare(p);
        eliminateCSE(p);
        permuteTree(p);
        assert(p->x.registered == 0);
    }
    flattenForest(forest);
    annotate(forest);
    processForest(forest);
    return forest;
}


#define LABEL_BLOCK(label) label->sclass

static ArrayList branches = NULL;

typedef struct _branch {
    int from;
    Symbol to;
} *Branch;

void insertBranchEdges() {
    Branch b;
    if(branches) {
        int i = list_size(branches);
        while(i) {
            i--;   
            b = (Branch)list_get(branches, i);
            addEdge(b->from, LABEL_BLOCK(b->to));
        }
        branches = NULL;
    }
}

void newBranch(int block, Symbol to) {
    Branch item;    
    if (branches == NULL) {
        branches = list_new(100);
    }
    item = allocate(sizeof(struct _branch), FUNC);
    item->from = block;
    item->to = to;
    list_add(branches, item);
}

void processForest(Node forest) {
    int depth = 0;
    int startDepth = 0;
    Node first, previous, tree;
    Symbol label = NULL;
    int block;
    first = previous = NULL;
    tree = forest;
    while (tree) {
//        fprintf(stderr, "op %d\n", generic(tree->op));
        switch(generic(tree->op)) {
        case LABEL:
            if (first) {
                block = flow_addBlock(block_new(startDepth, depth, first, previous));
                if (fallThrough >= 0) {
                    addEdge(fallThrough, block);
                }
                if (label) {
                    LABEL_BLOCK(label) = block;
                }
                fallThrough = block;
            }
            label = tree->syms[0];
            first = tree;
            startDepth = depth;
            break;
        case LT: case GT: case LE: case GE: case EQ: case NE: case JUMP:
//            fprintf(stderr, "branch or jump\n");
            if (first) {
                block = flow_addBlock(block_new(startDepth, depth, first, tree));
            } else {
                block = flow_addBlock(block_new(depth, depth, tree, tree));
            }
            if (label) {    
                LABEL_BLOCK(label) = block;
                label = NULL;    
            }
            if (fallThrough >= 0) {
                addEdge(fallThrough, block);
            }
            if (generic(tree->op) == JUMP) {
//                fprintf(stderr, "jump\n");
                if (generic(tree->kids[0]->op) == ADDRG) {
                    newBranch(block, tree->kids[0]->syms[0]);
                } else {
                    table_jump = block;
                }
                fallThrough = -1;
            } else {
//                fprintf(stderr, "branch\n");
                fallThrough = block;
                newBranch(block, tree->syms[0]);
            }
            first = NULL;
            startDepth = depth;
            break;
        case CALL:
            depth = 0;
            goto end;
        case ARG:
            depth += NODE_SIZE(tree);
            goto end;
        case ASGN:
            if (generic(tree->kids[0]->op) == CALL)
                depth = 0;
            goto end;
        default:
            end:
            if (first == NULL)
                first = tree;
        }
        previous = tree;
        tree = tree->link;
    }
    if (first) {
        block = flow_addBlock(block_new(startDepth, depth, first, previous));
        if (label) {        
            label->sclass = block;
            label = NULL;    
        }
        if (fallThrough >= 0) {
            addEdge(fallThrough, block);
        }
        fallThrough = block;
    }
}

void swtch(Symbol s) {
   newBranch(table_jump, s);
}


void printTree(Node p, FILE* out) {
    printsubtree(p, "", out);
}

void prettytree(Node p) {
    printsubtree(p, "", stderr);
}

static void printsubtree(Node p, char* indent, FILE* out) {
    if (p->kids[1]) {
        printnode(p, out);
        fprint(out, indent); 
        fprint(out, "| '-");
        printsubtree(p->kids[0], stringf("%s|   ", indent), out);
        fprint(out, indent); 
        fprint(out, "'---");
        printsubtree(p->kids[1], stringf("%s    ", indent), out);        
    } else if (p->kids[0]) {
        printnode(p, out);
        fprintf(out, indent); 
        fprint(out, "  '-");
        printsubtree(p->kids[0], stringf("%s    ", indent), out);        
    } else {
        printnode(p, out);    
    }
}

static void printnode(Node p, FILE* out) {
/*    if (p->op == VREG+P && p->syms[0]) {
        fprint(stderr, "VREGP(%s)", p->syms[0]->name);
        return;
    } */
    fprint(out, "%s ", opname(p->op));
    switch (generic(p->op)) {
    case VREG:
        fprint(out, "%s ", p->syms[0]->name);
        fprint(out, "e=%d l=%d", getEDepth(p), p->x.l_depth);
        fprint(out, " ref: %g", p->syms[0]->ref);
        break;
    case STACK: case COPY:
        if (p->syms[1])
            fprint(out, " %d e=%d l=%d", getDelve(p), getEDepth(p), p->x.l_depth);
        else 
            fprint(out, " NO DELVE e=%d l=%d", getEDepth(p), p->x.l_depth);
        if (p->syms[0])
            fprint(out, " %s", p->syms[0]->name);
        break;
    case TUCK:
        fprint(out, " %d", p->syms[1]->u.c.v.i);
        if (p->syms[0])
            fprint(out, " %s", p->syms[0]->name);
        break;
    case CNST: 
    case ADDRG: case ADDRF: case ADDRL:
        if (p->syms[0])
            fprint(out, "%s", p->syms[0]->name);
        break;
    case LABEL:     
        fprint(out, "%s", p->syms[0]->name);
        fprint(out, " ref: %g", p->syms[0]->ref);
        break;
    case RET: case NOOP:
    case CVF: case CVI: case CVP: case CVU: case JUMP: 
    case ARG: case BCOM: case NEG: case INDIR:
    case CALL:
    case ASGN: case BOR: case BAND: case BXOR: case RSH: case LSH:
    case ADD: case SUB:  case DIV: case MUL: case MOD:
        break;
    case EQ: case NE: case GT: case GE: case LE: case LT:
        fprint(out, " l%s", p->syms[0]->name);
        break;
    default: assert(0);
    }
    fprint(out, "\n");
}