trs/src/node.py

import os.path
import sys

sys.path.insert(0, os.path.realpath('external'))

from graph_drawing.graph import generate_graph
from graph_drawing.line import generate_line
from graph_drawing.node import Node, Leaf


TYPE_OPERATOR = 1
TYPE_IDENTIFIER = 2
TYPE_INTEGER = 4
TYPE_FLOAT = 8
TYPE_NUMERIC = TYPE_INTEGER | TYPE_FLOAT


# Unary
OP_NEG = 1

# Binary
OP_ADD = 2
OP_SUB = 3
OP_MUL = 4
OP_DIV = 5
OP_POW = 6
OP_MOD = 7

# N-ary (functions)
OP_INT = 8
OP_EXPAND = 9


TYPE_MAP = {
        int: TYPE_INTEGER,
        float: TYPE_FLOAT,
        str: TYPE_IDENTIFIER,
        }

OP_MAP = {
        '+': OP_ADD,
        # Either substitution or negation. Skip the operator sign in 'x' (= 2).
        '-': lambda x: OP_SUB if len(x) > 2 else OP_NEG,
        '*': OP_MUL,
        '/': OP_DIV,
        '^': OP_POW,
        'mod': OP_MOD,
        'int': OP_INT,
        'expand': OP_EXPAND,
        }


class ExpressionNode(Node):
    def __init__(self, *args, **kwargs):
        super(ExpressionNode, self).__init__(*args, **kwargs)
        self.type = TYPE_OPERATOR
        self.op = OP_MAP[args[0]]

        if hasattr(self.op, '__call__'):
            self.op = self.op(args)

    def __str__(self):  # pragma: nocover
        return generate_line(self)

    def graph(self):  # pragma: nocover
        return generate_graph(self)

    def replace(self, node):
        pos = self.parent.nodes.index(self)
        self.parent.nodes[pos] = node
        node.parent = self.parent
        self.parent = None

    def is_power(self):
        return self.op == OP_POW

    def is_nary(self):
        return self.op in [OP_ADD, OP_SUB, OP_MUL]

    def get_order(self):
        if self.is_power() and self[0].is_identifier() \
                and isinstance(self[1], Leaf):
            return (self[0].value, self[1].value, 1)

        for n0, n1 in [(0, 1), (1, 0)]:
            if self[n0].is_numeric() and not isinstance(self[n1], Leaf) \
                    and self[n1].is_power():
                coeff, power = self

                if power[0].is_identifier() and isinstance(power[1], Leaf):
                    return (power[0].value, power[1].value, coeff.value)

    def get_scope(self):
        scope = []

        for child in self:
            if not isinstance(child, Leaf) and child.op == self.op:
                scope += child.get_scope()
            else:
                scope.append(child)

        return scope


class ExpressionLeaf(Leaf):
    def __init__(self, *args, **kwargs):
        super(ExpressionLeaf, self).__init__(*args, **kwargs)

        self.type = TYPE_MAP[type(args[0])]

    def get_order(self):
        if self.is_identifier():
            return (self.value, 1, 1)

    def replace(self, node):
        if not hasattr(self, 'parent'):
            return

        pos = self.parent.nodes.index(self)
        self.parent.nodes[pos] = node
        node.parent = self.parent
        self.parent = None

    def is_identifier(self):
        return self.type & TYPE_IDENTIFIER

    def is_int(self):
        return self.type & TYPE_INTEGER

    def is_float(self):
        return self.type & TYPE_FLOAT

    def is_numeric(self):
        return self.type & TYPE_NUMERIC


if __name__ == '__main__':  # pragma: nocover
    l0 = ExpressionLeaf(3)
    l1 = ExpressionLeaf(4)
    l2 = ExpressionLeaf(5)
    l3 = ExpressionLeaf(7)

    n0 = ExpressionNode('+', l0, l1)
    n1 = ExpressionNode('+', l2, l3)
    n2 = ExpressionNode('*', n0, n1)

    print n2

    N = ExpressionNode

    def rewrite_multiply(node):
        a, b = node[0]
        c, d = node[1]

        ac = N('*', a, c)
        ad = N('*', a, d)
        bc = N('*', b, c)
        bd = N('*', b, d)

        res = N('+', N('+', N('+', ac, ad), bc), bd)

        return res

    possibilities = [
            (n0, lambda (x, y): ExpressionLeaf(x.value + y.value)),
            (n1, lambda (x, y): ExpressionLeaf(x.value + y.value)),
            (n2, rewrite_multiply),
            ]

    print '\n--- after rule 2 ---\n'

    n_, method = possibilities[2]
    new = method(n_)

    print new

    print '\n--- original graph ---\n'

    print n2

    print '\n--- apply rule 0 ---\n'

    n_, method = possibilities[0]
    new = method(n_)
    n_.replace(new)

    print n2

    # Revert rule 0
    new.replace(n_)

    print '\n--- apply rule 1 ---\n'

    n_, method = possibilities[1]
    new = method(n_)
    n_.replace(new)

    print n2