from .astnodes import *

class Visitor:
   """ Visitor that visits all nodes in the ast and runs the function 'f' """
   def __init__(self, f1, f2):
      self.f1 = f1
      self.f2 = f2
   def visit(self, node):
      # Run visitor:
      self.f1(node)
      # Descent into subnodes:
      if type(node) is Package:
         for s in node.scope:
            self.visit(s)
      elif type(node) is Function:
         for s in node.scope:
            self.visit(s)
         self.visit(node.typ)
         self.visit(node.body)
      elif type(node) is CompoundStatement:
         for s in node.statements:
            self.visit(s)
      elif type(node) is IfStatement:
         self.visit(node.condition)
         self.visit(node.truestatement)
         self.visit(node.falsestatement)
      elif type(node) is FunctionCall:
         for arg in node.args:
            self.visit(arg)
      elif type(node) is Assignment:
         self.visit(node.lval)
         self.visit(node.rval)
      elif type(node) is ReturnStatement:
         self.visit(node.expr)
      elif type(node) is Binop:
         self.visit(node.a)
         self.visit(node.b)
      elif type(node) is Constant:
         self.visit(node.value)
      elif type(node) in [EmptyStatement, VariableUse, Variable, Literal, FunctionType]:
         # Those nodes do not have child nodes.
         pass
      elif type(node) is WhileStatement:
         self.visit(node.condition)
         self.visit(node.dostatement)
      else:
         raise Exception('UNK visit "{0}"'.format(node))
      self.f2(node)