#!/usr/bin/env python

"""

Simplify AST structures for easier type propagation and analysis.

Copyright (C) 2006 Paul Boddie <paul@boddie.org.uk>

This software is free software; you can redistribute it and/or

modify it under the terms of the GNU General Public License as

published by the Free Software Foundation; either version 2 of

the License, or (at your option) any later version.

This software is distributed in the hope that it will be useful,

but WITHOUT ANY WARRANTY; without even the implied warranty of

MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

GNU General Public License for more details.

You should have received a copy of the GNU General Public

License along with this library; see the file LICENCE.txt

If not, write to the Free Software Foundation, Inc.,

51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA

"""

from compiler.visitor import ASTVisitor

import compiler.ast

from simplified import *

class Simplifier(ASTVisitor):

    """

    A simplifying visitor for AST nodes.

    Covered: And, AssAttr, AssList, AssName, AssTuple, Assign, AugAssign, Break,

             CallFunc, Class, Compare, Const, Continue, Dict, Discard, For,

             From, Function, Getattr, Global, If, Import, Invert, Keyword,

             Lambda, List, Module, Name, Not, Or, Pass, Raise, Return, Slice,

             Stmt, Subscript, TryExcept, TryFinally, Tuple, While, UnaryAdd,

             UnarySub.

    Missing: Add, Assert, Backquote, Bitand, Bitor, Bitxor, Decorators, Div,

             Ellipsis, Exec, FloorDiv, LeftShift, ListComp, ListCompFor,

             ListCompIf, Mod, Mul, Power, Print, Printnl, RightShift, Sliceobj,

             Sub, Yield.

    """

    def __init__(self):

        ASTVisitor.__init__(self)

        self.result = None              # The resulting tree.

        self.subprograms = []           # Subprograms outside the tree.

        self.current_subprograms = []   # Current subprograms being processed.

    # Generic visitor methods.

    def default(self, node, *args):

        raise ValueError, node.__class__

    def dispatch(self, node, *args):

        return ASTVisitor.dispatch(self, node, *args)

    def dispatches(self, nodes, *args):

        results = []

        for node in nodes:

            results.append(self.dispatch(node, *args))

        return results

    def dispatch_or_none(self, node, *args):

        if node is not None:

            return self.dispatch(node, *args)

        else:

            return LoadName(name="None")

    # Placeholder or deletion transformations.

    def visitStmt(self, stmt):

        return self.dispatches(stmt.nodes)

    def visitPass(self, pass_):

        return Pass(pass_)

    def visitDiscard(self, discard):

        return self.dispatch(discard.expr)

    # Relatively trivial transformations.

    def visitModule(self, module):

        self.result = Module(module)

        self.result.code = self.dispatch(module.node)

        return self.result

    def visitClass(self, class_):

        result = Class(class_, name=class_.name, bases=class_.bases)

        result.code = self.dispatch(class_.code)

        return result

    def visitGetattr(self, getattr):

        result = LoadAttr(getattr, name=getattr.attrname)

        result.expr = self.dispatch(getattr.expr)

        return result

    def visitKeyword(self, keyword):

        result = Keyword(keyword, name=keyword.name)

        result.expr = self.dispatch(keyword.expr)

        return result

    def visitGlobal(self, global_):

        result = Global(global_, names=global_.names)

        return result

    def visitImport(self, import_):

        result = Assign(import_)

        code = []

        for path, alias in import_.names:

            importer = Import(name=path)

            top = alias or path.split(".")[0]

            code.append(StoreName(expr=importer, name=top))

        result.code = code

        return result

    def visitFrom(self, from_):

        result = Assign(from_)

        code = []

        code.append(StoreTemp(expr=Import(name=from_.modname)))

        for name, alias in from_.names:

            code.append(StoreName(expr=LoadAttr(expr=LoadTemp(), name=name), name=(alias or name)))

        result.code = code

        return result

    def visitName(self, name):

        result = LoadName(name, name=name.name)

        return result

    def visitConst(self, const):

        result = LoadConst(const, value=const.value)

        return result

    def visitReturn(self, return_):

        result = Return(return_)

        result.expr = self.dispatch(return_.value)

        return result

    def visitBreak(self, break_):

        result = Return(break_)

        return result

    def visitContinue(self, continue_):

        result = Invoke(continue_, same_frame=1, produces_result=0, star=None, dstar=None, args=[])

        result.expr = LoadRef(ref=self.current_subprograms[-1])

        return result

    def visitRaise(self, raise_):

        result = Raise(raise_, expr=self.dispatch(raise_.expr1), traceback=None)

        if raise_.expr2 is not None:

            result.args = [self.dispatch(raise_.expr2)]

        if raise_.expr3 is not None:

            result.traceback = self.dispatch(raise_.expr3)

        return result

    def visitIf(self, if_):

        result = If(if_, else_=[])

        tests = []

        for compare, stmt in if_.tests:

            # Produce something like...

            # expr.__true__() ? body

            test = Conditional(else_=[], test=Invoke(

                expr=LoadAttr(expr=self.dispatch(compare), name="__true__"),

                params=[], star=None, dstar=None))

            test.body = self.dispatch(stmt)

            tests.append(test)

        result.tests = tests

        if if_.else_ is not None:

            result.else_ = self.dispatch(if_.else_)

        return result

    def _visitBuiltin(self, builtin, name):

        result = Invoke(builtin, expr=LoadName(name=name))

        result.args = self.dispatches(builtin.nodes)

        return result

    def visitTuple(self, tuple):

        return self._visitBuiltin(tuple, "Tuple")

    def visitList(self, list):

        return self._visitBuiltin(list, "List")

    def visitDict(self, dict):

        result = Invoke(dict, expr=LoadName(name="Dict"))

        args = []

        for key, value in dict.items:

            tuple = Invoke(expr=LoadName(name="Tuple"), star=None, dstar=None)

            tuple.args = [self.dispatch(key), self.dispatch(value)]

            args.append(tuple)

        result.args = args

        return result

    # Logical and comparison operations.

    comparison_methods = {

        "==" : "__eq__", "!=" : "__ne__", "<" : "__lt__", "<=" : "__le__",

        ">=" : "__ge__", ">" : "__gt__", "is" : None, "is not" : None

        }

    def visitCompare(self, compare):

        # Make a subprogram for the expression and record it outside the main tree.

        subprogram = Subprogram(name=hex(id(compare)), acquire_locals=1, returns_value=1, params=[], star=None, dstar=None)

        self.current_subprograms.append(subprogram)

        # In the subprogram, make instructions which invoke a method on the

        # first operand of each operand pair and, if appropriate, return with

        # the value from that method.

        nodes = []

        last = compare.ops[-1]

        previous = self.dispatch(compare.expr)

        for op in compare.ops:

            op_name, node = op

            expr = self.dispatch(node)

            method_name = self.comparison_methods[op_name]

            if method_name:

                invocation = Invoke(expr=LoadAttr(expr=previous, name=method_name), params=[expr], star=None, dstar=None)

            elif op_name == "is":

                invocation = Invoke(expr=LoadName(name="__is__"), params=[previous, expr], star=None, dstar=None)

            elif op_name == "is not":

                invocation = Not(expr=Invoke(expr=LoadName(name="__is__"), params=[previous, expr], star=None, dstar=None))

            else:

                raise NotImplementedError, op_name

            nodes.append(StoreTemp(expr=invocation))

            if op is not last:

                nodes.append(Conditional(test=Not(expr=LoadTemp()), body=[Return(expr=LoadTemp())]))

                nodes.append(ReleaseTemp())

            else:

                nodes.append(Return(expr=LoadTemp()))

        subprogram.code = nodes

        self.current_subprograms.pop()

        self.subprograms.append(subprogram)

        # Make an invocation of the subprogram.

        result = Invoke(compare, same_frame=1, star=None, dstar=None, args=[])

        result.expr = LoadRef(ref=subprogram)

        return result

    def visitAnd(self, and_):

        # Make a subprogram for the expression and record it outside the main tree.

        subprogram = Subprogram(name=hex(id(and_)), acquire_locals=1, returns_value=1, params=[], star=None, dstar=None)

        self.current_subprograms.append(subprogram)

        # In the subprogram, make instructions which store each operand, test

        # for each operand's truth status, and if appropriate return from the

        # subprogram with the value of the operand.

        nodes = []

        last = and_.nodes[-1]

        for node in and_.nodes:

            expr = self.dispatch(node)

            if node is not last:

                nodes.append(StoreTemp(expr=expr))

                invocation = Invoke(expr=LoadAttr(expr=LoadTemp(), name="__true__"), params=[], star=None, dstar=None)

                nodes.append(Conditional(test=Not(expr=invocation), body=[Return(expr=LoadTemp())]))

                nodes.append(ReleaseTemp())

            else:

                nodes.append(Return(expr=expr))

        subprogram.code = nodes

        self.current_subprograms.pop()

        self.subprograms.append(subprogram)

        # Make an invocation of the subprogram.

        result = Invoke(and_, same_frame=1, star=None, dstar=None, args=[])

        result.expr = LoadRef(ref=subprogram)

        return result

    def visitOr(self, or_):

        # Make a subprogram for the expression and record it outside the main tree.

        subprogram = Subprogram(name=hex(id(or_)), acquire_locals=1, returns_value=1, params=[], star=None, dstar=None)

        self.current_subprograms.append(subprogram)

        # In the subprogram, make instructions which store each operand, test

        # for each operand's truth status, and if appropriate return from the

        # subprogram with the value of the operand.

        nodes = []

        last = or_.nodes[-1]

        for node in or_.nodes:

            expr = self.dispatch(node)

            if node is not last:

                nodes.append(StoreTemp(expr=expr))

                invocation = Invoke(expr=LoadAttr(expr=LoadTemp(), name="__true__"),

                    params=[], star=None, dstar=None)

                nodes.append(Conditional(test=invocation, body=[Return(expr=LoadTemp())]))

                nodes.append(ReleaseTemp())

            else:

                nodes.append(Return(expr=expr))

        subprogram.code = nodes

        self.current_subprograms.pop()

        self.subprograms.append(subprogram)

        # Make an invocation of the subprogram.

        result = Invoke(or_, same_frame=1, star=None, dstar=None, args=[])

        result.expr = LoadRef(ref=subprogram)

        return result

    def visitNot(self, not_):

        result = Not(not_, expr=Invoke(expr=LoadAttr(expr=self.dispatch(not_.expr), name="__true__"),

            params=[], star=None, dstar=None))

        return result

    # Operators.

    def visitUnaryAdd(self, unaryadd):

        return Invoke(unaryadd, expr=LoadAttr(expr=self.dispatch(unaryadd.expr), name="__pos__"), args=[])

    def visitUnarySub(self, unarysub):

        return Invoke(unarysub, expr=LoadAttr(expr=self.dispatch(unarysub.expr), name="__neg__"), args=[])

    def visitInvert(self, invert):

        return Invoke(invert, expr=LoadAttr(expr=self.dispatch(invert.expr), name="__invert__"), args=[])

    # Assignments.

    augassign_methods = {

        "+=" : "__iadd__", "-=" : "__isub__", "*=" : "__imul__", "/=" : "__idiv__",

        "%=" : "__imod__", "**=" : "__ipow__", "<<=" : "__ilshift__", ">>=" : "__irshift__",

        "&=" : "__iand__", "^=" : "__ixor__", "|=" : "__ior__"

        }

    def visitAugAssign(self, augassign):

        result = Assign(augassign)

        expr = self.dispatch(augassign.expr)

        # Simple augmented assignment: name += expr

        if isinstance(augassign.node, compiler.ast.Name):

            name = augassign.node

            node = self.dispatch(name)

            get_incremented = StoreTemp(

                expr=Invoke(expr=LoadAttr(expr=node, name=self.augassign_methods[augassign.op]), args=[expr])

                )

            store = StoreName(expr=LoadTemp(), name=name.name)

            result.code = [get_incremented, store, ReleaseTemp()]

        # Complicated augmented assignment: lvalue.attr += expr

        elif isinstance(augassign.node, compiler.ast.Getattr):

            # <lvalue> -> setattr(<lvalue>, getattr(<lvalue>, "attr").__xxx__(expr))

            getattr = augassign.node

            store_expr = StoreTemp(index="expr", expr=self.dispatch(getattr.expr))

            node_attr = LoadAttr(expr=LoadTemp(index="expr"), name=getattr.attrname)

            get_incremented = StoreTemp(

                expr=Invoke(expr=LoadAttr(expr=node_attr, name=self.augassign_methods[augassign.op]), args=[expr])

                )

            store = StoreAttr(expr=LoadTemp(), lvalue=LoadTemp(index="expr"), name=getattr.attrname)

            result.code = [store_expr, get_incremented, store, ReleaseTemp(index="expr"), ReleaseTemp()]

        # Complicated augassign using slices: lvalue[lower:upper] += expr

        elif isinstance(augassign.node, compiler.ast.Slice):

            # <lvalue>, <lower>, <upper> -> <lvalue>.__setslice__(<lower>, <upper>, <lvalue>.__getslice__(<lower>, <upper>).__xxx__(expr))

            slice = augassign.node

            store_expr = StoreTemp(index="expr", expr=self.dispatch(slice.expr))

            store_lower = StoreTemp(index="lower", expr=self.dispatch_or_none(slice.lower))

            store_upper = StoreTemp(index="upper", expr=self.dispatch_or_none(slice.upper))

            node_slice = self._visitSlice(slice, LoadTemp(index="expr"), LoadTemp(index="lower"), LoadTemp(index="upper"), "OP_APPLY")

            get_incremented = StoreTemp(

                expr=Invoke(expr=LoadAttr(expr=node_slice, name=self.augassign_methods[augassign.op]), args=[expr])

                )

            store = self._visitSlice(slice, LoadTemp(index="expr"), LoadTemp(index="lower"), LoadTemp(index="upper"), "OP_ASSIGN", LoadTemp())

            result.code = [store_expr, store_lower, store_upper, get_incremented, store,

                ReleaseTemp(index="expr"), ReleaseTemp(index="lower"), ReleaseTemp(index="upper"), ReleaseTemp()]

        # Complicated augassign using subscripts: lvalue[subs] += expr

        elif isinstance(augassign.node, compiler.ast.Subscript):

            # <lvalue>, <subs> -> <lvalue>.__setitem__(<subs>, <lvalue>.__getitem__(<subs>).__xxx__(expr))

            subscript = augassign.node

            store_expr = StoreTemp(index="expr", expr=self.dispatch(subscript.expr))

            subs = self._visitSubscriptSubs(subscript.subs)

            store_subs = StoreTemp(index="subs", expr=subs)

            node_subscript = self._visitSubscript(subscript, LoadTemp(index="expr"), LoadTemp(index="subs"), "OP_APPLY")

            get_incremented = StoreTemp(

                expr=Invoke(expr=LoadAttr(expr=node_subscript, name=self.augassign_methods[augassign.op]), args=[expr])

                )

            store = self._visitSubscript(subscript, LoadTemp(index="expr"), LoadTemp(index="subs"), "OP_ASSIGN", LoadTemp())

            result.code = [store_expr, store_subs, get_incremented, store, ReleaseTemp(index="expr"), ReleaseTemp(index="subs"), ReleaseTemp()]

        else:

            raise NotImplementedError, augassign.node.__class__

        return result

    def visitAssign(self, assign):

        result = Assign(assign)

        store = StoreTemp(expr=self.dispatch(assign.expr))

        release = ReleaseTemp()

        result.code = [store] + self.dispatches(assign.nodes, 0) + [release]

        return result

    def visitAssList(self, asslist, in_sequence=0):

        if not in_sequence:

            expr = LoadTemp()

        else:

            expr = Invoke(expr=LoadAttr(expr=LoadTemp(), name="next"))

        result = Assign(asslist)

        store = StoreTemp(expr=Invoke(expr=LoadAttr(name="__iter__", expr=expr)))

        release = ReleaseTemp()

        result.code = [store] + self.dispatches(asslist.nodes, 1) + [release]

        return result

    visitAssTuple = visitAssList

    def _visitAssNameOrAttr(self, node, in_sequence):

        if not in_sequence:

            return LoadTemp()

        else:

            return Invoke(expr=LoadAttr(expr=LoadTemp(), name="next"))

    def visitAssName(self, assname, in_sequence=0):

        expr = self._visitAssNameOrAttr(assname, in_sequence)

        result = StoreName(assname, name=assname.name, expr=expr)

        return result

    def visitAssAttr(self, assattr, in_sequence=0):

        expr = self._visitAssNameOrAttr(assattr, in_sequence)

        lvalue = self.dispatch(assattr.expr)

        result = StoreAttr(assattr, name=assattr.attrname, lvalue=lvalue, expr=expr)

        return result

    def _visitSlice(self, slice, expr, lower, upper, flags, value=None):

        if flags == "OP_ASSIGN":

            args = [value]

            result = Invoke(expr=LoadAttr(expr=expr, name="__setslice__"))

        elif flags == "OP_APPLY":

            args = []

            result = Invoke(expr=LoadAttr(expr=expr, name="__getslice__"))

        elif flags == "OP_DELETE":

            args = []

            result = Invoke(expr=LoadAttr(expr=expr, name="__delslice__"))

        else:

            raise NotImplementedError, flags

        # Add the dimensions.

        args.insert(0, lower)

        args.insert(1, upper)

        result.original = slice

        result.args = args

        return result

    def visitSlice(self, slice, in_sequence=0):

        value = self._visitAssNameOrAttr(slice, in_sequence)

        return self._visitSlice(slice, self.dispatch(slice.expr), self.dispatch_or_none(slice.lower),

            self.dispatch_or_none(slice.upper), slice.flags, value)

    def _visitSubscript(self, subscript, expr, subs, flags, value=None):

        if flags == "OP_ASSIGN":

            args = [value]

            result = Invoke(expr=LoadAttr(expr=expr, name="__setitem__"))

        elif flags == "OP_APPLY":

            args = []

            result = Invoke(expr=LoadAttr(expr=expr, name="__getitem__"))

        elif flags == "OP_DELETE":

            args = []

            result = Invoke(expr=LoadAttr(expr=expr, name="__delitem__"))

        else:

            raise NotImplementedError, flags

        # Add the dimensions.

        args.insert(0, subs)

        result.original = subscript

        result.args = args

        return result

    def _visitSubscriptSubs(self, subs):

        if len(subs) == 1:

            return self.dispatch(subs[0])

        else:

            return Invoke(expr=LoadName(name="Tuple"), args=self.dispatches(subs))

    def visitSubscript(self, subscript, in_sequence=0):

        value = self._visitAssNameOrAttr(subscript, in_sequence)

        subs = self._visitSubscriptSubs(subscript.subs)

        return self._visitSubscript(subscript, self.dispatch(subscript.expr), subs, subscript.flags, value)

    # Invocation and subprogram transformations.

    def _visitFunction(self, function, subprogram):

        if function.flags & 4 != 0: has_star = 1

        else: has_star = 0

        if function.flags & 8 != 0: has_dstar = 1

        else: has_dstar = 0

        ndefaults = len(function.defaults)

        npositional = len(function.argnames) - has_star - has_dstar

        if has_star: star = function.argnames[npositional]

        else: star = None

        if has_dstar: dstar = function.argnames[npositional + has_star]

        else: dstar = None

        params = []

        for i in range(0, npositional - ndefaults):

            params.append((function.argnames[i], None))

        # NOTE: Fix/process defaults.

        for i in range(0, ndefaults):

            default = function.defaults[i]

            if default is not None:

                params.append((function.argnames[npositional - ndefaults + i], self.dispatch(default)))

            else:

                params.append((function.argnames[npositional - ndefaults + i], default))

        subprogram.params = params

        subprogram.star = star

        subprogram.dstar = dstar

        self.subprograms.append(subprogram)

    def visitFunction(self, function):

        # Make a subprogram for the function and record it outside the main

        # tree.

        subprogram = Subprogram(name=function.name, returns_value=1, star=None, dstar=None)

        self.current_subprograms.append(subprogram)

        subprogram.code = self.dispatch(function.code)

        self.current_subprograms.pop()

        self._visitFunction(function, subprogram)

        # Make a definition of the function associating it with a name.

        result = Assign(function)

        load = LoadRef(ref=subprogram)

        store = StoreName(name=function.name)

        result.code = [load, store]

        return result

    def visitLambda(self, lambda_):

        # Make a subprogram for the function and record it outside the main

        # tree.

        subprogram = Subprogram(name=hex(id(lambda_)), returns_value=1, star=None, dstar=None)

        self.current_subprograms.append(subprogram)

        subprogram.code = [Return(expr=self.dispatch(lambda_.code))]

        self.current_subprograms.pop()

        self._visitFunction(lambda_, subprogram)

        # Get the subprogram reference to the lambda.

        return LoadRef(lambda_, ref=subprogram)

    def visitCallFunc(self, callfunc):

        result = Invoke(callfunc, same_frame=0, star=None, dstar=None)

        result.args = self.dispatches(callfunc.args)

        if callfunc.star_args is not None:

            result.star = self.dispatch(callfunc.star_args)

        if callfunc.dstar_args is not None:

            result.dstar = self.dispatch(callfunc.dstar_args)

        result.expr = self.dispatch(callfunc.node)

        return result

    def visitWhile(self, while_):

        # Make a subprogram for the block and record it outside the main tree.

        subprogram = Subprogram(name=hex(id(while_)), acquire_locals=1, returns_value=0, params=[], star=None, dstar=None)

        self.current_subprograms.append(subprogram)

        # Include a conditional statement in the subprogram.

        test = Conditional(else_=[])

        test.test = Invoke(expr=LoadAttr(expr=self.dispatch(while_.test), name="__true__"),

            params=[], star=None, dstar=None)