from .astnodes import *

class Visitor:
    """
        Visitor that can visit all nodes in the AST
        and run pre and post functions.
    """
    def visit(self, node, f_pre=None, f_post=None):
        self.f_pre = f_pre
        self.f_post = f_post
        self.do(node)

    def do(self, node):
        # Run visitor:
        if self.f_pre:
            self.f_pre(node)

        # Descent into subnodes:
        if type(node) is Package:
            for decl in node.declarations:
                self.do(decl)
        elif type(node) is Function:
            for s in node.declarations:
                self.do(s)
            self.do(node.body)
        elif type(node) is CompoundStatement:
            for s in node.statements:
                self.do(s)
        elif type(node) is IfStatement:
            self.do(node.condition)
            self.do(node.truestatement)
            if node.falsestatement:
                self.do(node.falsestatement)
        elif type(node) is FunctionCall:
            for arg in node.args:
                self.do(arg)
        elif type(node) is Assignment:
            self.do(node.lval)
            self.do(node.rval)
        elif type(node) is ReturnStatement:
            self.do(node.expr)
        elif type(node) is Binop:
            self.do(node.a)
            self.do(node.b)
        elif type(node) is Unop:
            self.do(node.a)
        elif type(node) is ExpressionStatement:
            self.do(node.ex)
        elif type(node) is TypeCast:
            self.do(node.a)
        elif type(node) is FieldRef:
            self.do(node.base)
        elif type(node) is Deref:
            self.do(node.ptr)
        elif type(node) is Constant:
            self.do(node.value)
        elif type(node) in [VariableUse, Variable, Literal, FunctionType, DefinedType]:
            # Those nodes do not have child nodes.
            pass
        elif type(node) is WhileStatement:
            self.do(node.condition)
            self.do(node.statement)
        else:
            raise Exception('Could not visit "{0}"'.format(node))

        # run post function
        if self.f_post:
            self.f_post(node)