from cfg.CFG_Node import (
    CFG_Node,
    CFG_CALL,
    CFG_RETURN,
    CFG_DIAMOND,
)

import compiler
import syntax


# -------------------------------------------------
# Expressions — NO CFG NODES
# -------------------------------------------------

class CONST(compiler.CONST):
    def cfa(self, pred, end):
        return pred


class ID(compiler.ID):
    def cfa(self, pred, end):
        return pred


class AOP(compiler.AOP):
    def cfa(self, pred, end):
        return pred


class COMP(compiler.COMP):
    def cfa(self, pred, end):
        return pred


class EQOP(compiler.EQOP):
    def cfa(self, pred, end):
        return pred


class LOP(compiler.LOP):
    def cfa(self, pred, end):
        return pred


# -------------------------------------------------
# Statements
# -------------------------------------------------

class ASSIGN(compiler.ASSIGN):
    def cfa(self, pred, end):
        n = CFG_Node(self)
        pred.add_child(n)
        return n


class SEQ(compiler.SEQ):
    def cfa(self, pred, end):
        mid = self.exp1.cfa(pred, end)
        if mid is None:
            return None
        return self.exp2.cfa(mid, end)


class IF(compiler.IF):
    def cfa(self, pred, end):
        cond = CFG_DIAMOND(self)
        pred.add_child(cond)

        join = CFG_Node()

        then_end = self.exp1.cfa(cond, end)
        else_end = self.exp2.cfa(cond, end)

        if then_end is not None:
            then_end.add_child(join)
        if else_end is not None:
            else_end.add_child(join)

        return join


class WHILE(compiler.WHILE):
    def cfa(self, pred, end):
        cond = CFG_DIAMOND(self)
        pred.add_child(cond)

        body_end = self.body.cfa(cond, end)
        if body_end is not None:
            body_end.add_child(cond)

        after = CFG_Node()
        cond.add_child(after)

        return after


# -------------------------------------------------
# Functions / calls
# -------------------------------------------------

class CALL(compiler.CALL):
    def cfa(self, pred, end):
        n = CFG_CALL(self)
        pred.add_child(n)
        return n


class DECL(compiler.DECL):
    def cfa(self, pred, end):
        entry = CFG_Node(self)
        pred.add_child(entry)
        return self.body.cfa(entry, end)


class LET(compiler.LET):
    def cfa(self, pred, end):
        current = pred

        decls = self.decl if isinstance(self.decl, list) else [self.decl]
        for d in decls:
            current = d.cfa(current, end)
            if current is None:
                return None

        return self.body.cfa(current, end)


class RETURN(syntax.EXPRESSION):
    def cfa(self, pred, end):
        n = CFG_RETURN(self)
        pred.add_child(n)
        n.add_child(end)   # direct jump to global END
        return None        # no fallthrough