--- a/python/c3/visitor.py	Fri Jul 05 11:18:58 2013 +0200
+++ b/python/c3/visitor.py	Fri Jul 05 13:00:03 2013 +0200
@@ -1,49 +1,58 @@
 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):
+    """ 
+        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:
-      self.f1(node)
+      if self.f_pre:
+            self.f_pre(node)
+
       # Descent into subnodes:
       if type(node) is Package:
-         for s in node.scope:
-            self.visit(s)
+            for decl in node.declarations:
+                self.do(decl)
       elif type(node) is Function:
-         for s in node.scope:
-            self.visit(s)
-         self.visit(node.typ)
-         self.visit(node.body)
+            for s in node.declarations:
+                self.do(s)
+            self.do(node.body)
       elif type(node) is CompoundStatement:
          for s in node.statements:
-            self.visit(s)
+            self.do(s)
       elif type(node) is IfStatement:
-         self.visit(node.condition)
-         self.visit(node.truestatement)
-         self.visit(node.falsestatement)
+         self.do(node.condition)
+         self.do(node.truestatement)
+         self.do(node.falsestatement)
       elif type(node) is FunctionCall:
          for arg in node.args:
-            self.visit(arg)
+            self.do(arg)
       elif type(node) is Assignment:
-         self.visit(node.lval)
-         self.visit(node.rval)
+         self.do(node.lval)
+         self.do(node.rval)
       elif type(node) is ReturnStatement:
-         self.visit(node.expr)
+         self.do(node.expr)
       elif type(node) is Binop:
-         self.visit(node.a)
-         self.visit(node.b)
+         self.do(node.a)
+         self.do(node.b)
       elif type(node) is Constant:
-         self.visit(node.value)
+         self.do(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)
+         self.do(node.condition)
+         self.do(node.statement)
       else:
-         raise Exception('UNK visit "{0}"'.format(node))
-      self.f2(node)
+           raise Exception('Could not visit "{0}"'.format(node))
 
+      # run post function
+      if self.f_post:
+            self.f_post(node)
+