# Expressions are represented as a tree, where there is a different class for
# each type of node.  For example, an expression (x + 3*4) corresponds to a
# tree:
#     +                                    AExprPlus
#   /   \                                   /   \
#  x    *     -- corresponding to --  AExprVar AExprMul
#      / \                                     /     \
#     3   4                           AExprNumeral  AExprNumeral
# Each expression class is equipped with three methods:
#  eval: takes in a state (dictionary mapping variable names to ints), and
#        produces a value (integer in the case of AExprs or boolean in the
#        case of BExprs)
# to_py: returns a Python string representation of the expression
# cprop: take in a partial state (dictionary mapping *some* variable names to
#        ints), and evaluates the expression to the extent possible, producing
#        either an integer or a residual expression.


# for constant propagation
def make_expr(value):
    """Wrap a value that is either an int or an expression into an expression"""
    if type(value) == int:
        return AExprNumeral(value)
    else:
        return value

class AExprVar:
    """Variables"""
    def __init__(self, name):
        self.name = name
    def eval(self, state):
        return state[self.name]
    def to_py(self):
        return self.name
    def cprop(self, const):
        return const.get(self.name, self)
        
class AExprNumeral:
    """Numerals"""
    def __init__(self, value):
        self.value = value
    def eval(self, state):
        return self.value
    def to_py(self):
        return str(self.value)
    def cprop(self, const):
        return self.value

class BinaryExpr:
    """Abstract class representing binary expressions"""
    def __init__(self, left, right):
        self.left = left
        self.right = right
    def eval(self, state):
        pass
    def to_py(self, state):
        pass
    def cprop(self, const):
        pass

class AExprPlus(BinaryExpr):
    """Addition"""
    def eval(self, state):
        return self.left.eval(state) + self.right.eval(state)
    def to_py(self):
        return self.left.to_py() + " + " + self.right.to_py()
    def cprop(self, const):
        left = self.left.cprop(const)
        right = self.right.cprop(const)
        if type(left) == int and type(right) == int:
            return left + right
        elif left == 0:
            return right
        elif right == 0:
            return left
        else:
            return AExprPlus(make_expr(left), make_expr(right))

class AExprMul(BinaryExpr):
    """Multiplication"""
    def eval(self, state):
        return self.left.eval(state) * self.right.eval(state)
    def to_py(self):
        return self.left.to_py() + " * " + self.right.to_py()
    def cprop(self, const):
        left = self.left.cprop(const)
        right = self.right.cprop(const)
        if type(left) == int and type(right) == int:
            return left * right
        elif left == 0 or right == 0:
            return 0
        elif left == 1:
            return right
        elif right == 1:
            return left
        else:
            return AExprMul(make_expr(left), make_expr(right))

class BExprLt(BinaryExpr):
    """Strictly less-than"""
    def eval(self, state):
        return self.left.eval(state) < self.right.eval(state)
    def to_py(self):
        return self.left.to_py() + " < " + self.right.to_py()
    def cprop(self, const):
        left = self.left.cprop(const)
        right = self.right.cprop(const)
        if type(left) == int and type(right) == int:
            return left < right
        else:
            return BExprLt(make_expr(left), make_expr(right))

class BExprEq(BinaryExpr):
    """Equal to"""
    def eval(self, state):
        return self.left.eval(state) == self.right.eval(state)
    def to_py(self):
        return self.left.to_py() + " == " + self.right.to_py()
    def cprop(self, const):
        left = self.left.cprop(const)
        right = self.right.cprop(const)
        if type(left) == int and type(right) == int:
            return left == right
        else:
            return BExprEq(make_expr(left), make_expr(right))


# Statements are program phrases that can change the state of the program.
# Again, each statement is equipped with three methods:
#  execute: takes in a state, and produces nothing (but may change the state
#           or print something!)
# to_py: take in an indentation level and produces a Python string
#        representation of the statement
# cprop: take in a partial state, and computes a residual statement that is
#        equivalent under the given partial state.  Additionally, the partial state
#        is updated to reflect the computation along the statement.

# for constant propagation
def join_const(m, n):
    """Combine two partial state.  Keep only mappings on which both partial states agree"""
    to_delete = []
    for k in m.keys():
        if not (k in n and m[k] == n[k]):
            del m[k]

class StmtAssign:
    """Variable assignment"""
    def __init__(self, lhs, rhs):
        self.lhs = lhs
        self.rhs = rhs
    def execute(self, state):
        state[self.lhs] = self.rhs.eval(state)
    def to_py(self, indent):
        return ("    " * indent) + self.lhs + " = " + self.rhs.to_py() + "\n"
    def cprop(self, const):
        rhs = self.rhs.cprop(const)
        if type(rhs) == int:
            const[self.lhs] = rhs
        elif self.lhs in const:
            del const[self.lhs]
        return StmtAssign(self.lhs, make_expr(rhs))

class StmtIf:
    """Conditional statement"""
    def __init__(self, cond, bthen, belse):
        self.cond = cond
        self.bthen = bthen
        self.belse = belse
    def execute(self, state):
        if (self.cond.eval(state)):
            self.bthen.execute(state)
        else:
            self.belse.execute(state)
    def to_py(self, indent):
        program = ("    " * indent) + "if " + self.cond.to_py() + ":\n"
        program += self.bthen.to_py(indent + 1)
        program += ("    " * indent) + "else:\n"
        program += self.belse.to_py(indent + 1)
        return program
    def cprop(self, const):
        cond = self.cond.cprop(const)
        if type(cond) == bool:
            # if conditional is a constant, just constant propagate the
            # appropriate branch
            if cond:
                return self.bthen.cprop(const)
            else:
                return self.belse.cprop(const)
        else:
            # if conditional is non-constant, propagate the partial state
            # across both branches and then take partial state consisting only
            # of mappings that match.

            # For example, starting in a partial state where x maps to 0,
            # if (y < 3) { x = x + 1 } else { x = 2 }
            # should yield the partial state where x maps to 2 (x is
            # guaranteed to be 2 after executing the statement).
            # However, constant propagation across
            # if (y < 3) { x = x + 1 } else { x = 3 }
            # should yield an empty partial state, because we don't know if x
            # is 2 or 3.
            const_else = const.copy()
            bthen = self.bthen.cprop(const)
            belse = self.belse.cprop(const_else)
            join_const(const, const_else)
            return StmtIf(make_expr(cond), bthen, belse)

class StmtBlock:
    """Sequence of statements"""
    def __init__(self, block):
        self.block = block
    def execute(self, state):
        for stmt in self.block:
            stmt.execute(state)
    def to_py(self, indent):
        return "".join(map(lambda x: x.to_py(indent), self.block))
    def cprop(self, const):
        return StmtBlock(map(lambda x: x.cprop(const), self.block))

class StmtWhile:
    """While loop"""
    def __init__(self, cond, body):
        self.cond = cond
        self.body = body
    def execute(self, state):
        while self.cond.eval(state):
            self.body.execute(state)
    def to_py(self, indent):
        program = ("    " * indent) + "while " + self.cond.to_py() + ":\n"
        program += self.body.to_py(indent + 1)
        return program
    def cprop(self, const):
        # to propagate a partial state across a loop, we repeatedly propagate
        # it across its body, at each step throwing out mappings that don't
        # agree with the original partial state, until eventually the partial
        # state converges.  Since the domain of the partial state is finite,
        # this loop must terminate.
        while True:
            const_copy = const.copy()
            body = self.body.cprop(const_copy)
            previous_size = len(const)
            join_const(const, const_copy)
            if len(const) == previous_size:
                # the partial state hasn't changed -- all done!
                return StmtWhile(make_expr(self.cond.cprop(const)), body)

class StmtPrint:
    """Print to stdout"""
    def __init__(self, expr):
        self.expr = expr
    def execute(self, state):
        print self.expr.eval(state)
    def to_py(self, indent):
        return ("    " * indent) + "print(" + self.expr.to_py() + ")\n"
    def cprop(self, const):
        return StmtPrint(make_expr(self.expr.cprop(const)))

# Examples
mystate = {'x' : 1, 'y' : 2, 'z': 3 }

print "3+4       ==> %d" % AExprPlus(AExprNumeral(3), AExprNumeral(4)).eval(mystate)
print "x+y       ==> %d" % AExprPlus(AExprVar("x"), AExprVar("y")).eval(mystate)
print "7+4 == 11 ==> %r" % BExprEq(AExprPlus(AExprNumeral(7), AExprNumeral(4)), AExprNumeral(11)).eval(mystate)

# x=0; flag=1; while(x<0){ if(flag == 1) {x++} else {flag++}; print x }
program = StmtBlock ([StmtAssign("x", AExprNumeral(0)),
                      StmtAssign("flag", AExprNumeral(1)),
                      StmtWhile (BExprLt(AExprVar("x"), AExprNumeral(10)),
                                 StmtBlock([StmtIf(BExprEq(AExprVar("flag"), AExprNumeral(1)),
                                                   StmtAssign("x", AExprPlus(AExprVar("x"), AExprNumeral(1))),
                                                   StmtAssign("flag", AExprPlus(AExprVar("flag"), AExprNumeral(1)))),
                                            StmtPrint(AExprVar("x"))]))])
# example for constant propagation through a conditional
# x = y*0; if(y == 0) { x++ } else { x = 2 }; y = x
program2 = StmtBlock([StmtAssign("x", AExprMul(AExprVar("y"), AExprNumeral(0))),
                      StmtIf(BExprEq(AExprVar("y"), AExprNumeral(0)),
                             StmtAssign("x", AExprPlus(AExprVar("x"), AExprNumeral(1))),
                            StmtAssign("x", AExprNumeral(1))),
                      StmtAssign("y", AExprVar("x"))])


print "Python translation:"
print program.to_py(0)
print "After constant propagation:"
print program.cprop({}).to_py(0)
print "Execution:"
program.execute(mystate)