--- a/python/transform.py	Wed Jul 31 21:20:58 2013 +0200
+++ b/python/transform.py	Sun Aug 04 15:10:10 2013 +0200
@@ -7,12 +7,17 @@
 # Standard passes:
 
 class FunctionPass:
+    def __init__(self):
+        self.logger = logging.getLogger('optimize')
+
     def run(self, ir):
         """ Main entry point for the pass """
-        logging.info('Running pass {}'.format(type(self)))
+        self.logger.info('Running pass {}'.format(type(self)))
+        ir.check()
         self.prepare()
         for f in ir.Functions:
             self.onFunction(f)
+        ir.check()
 
     def onFunction(self, f):
         """ Override this virtual method """
@@ -42,34 +47,33 @@
         raise NotImplementedError()
 
 # Usefull transforms:
-class ConstantFolder(InstructionPass):
-    def prepare(self):
-        self.constMap = {}
-
-    def onInstruction(self, i):
-      if type(i) is ImmLoad:
-         self.constMap[i.target] = i.value
-      elif type(i) is BinaryOperator:
-         if i.value1 in self.constMap and i.value2 in self.constMap and i.operation in ['+', '-', '*', '<<']:
-            op = i.operation
-            va = self.constMap[i.value1]
-            vb = self.constMap[i.value2]
-            if op == '+':
-               vr = va + vb
-            elif op == '*':
-               vr = va * vb
-            elif op == '-':
-               vr = va - vb
-            elif op == '<<':
-                vr = va << vb
-            else:
-               vr = None
-               return
-            self.constMap[i.result] = vr
-            i.removeDef(i.result)
-            i2 = ImmLoad(i.result, vr)
-            logging.debug('Replacing {}'.format(i))
-            i.Parent.replaceInstruction(i, i2)
+class ConstantFolder(BasicBlockPass):
+    def onBasicBlock(self, bb):
+        constMap = {}
+        ins = [i for i in bb.Instructions if type(i) in [ImmLoad, BinaryOperator]]
+        for i in ins:
+            if type(i) is ImmLoad:
+                constMap[i.target] = i.value
+            elif type(i) is BinaryOperator:
+                if i.value1 in constMap and i.value2 in constMap and i.operation in ['+', '-', '*', '<<']:
+                    op = i.operation
+                    va = constMap[i.value1]
+                    vb = constMap[i.value2]
+                    if op == '+':
+                       vr = va + vb
+                    elif op == '*':
+                       vr = va * vb
+                    elif op == '-':
+                       vr = va - vb
+                    elif op == '<<':
+                        vr = va << vb
+                    else:
+                        raise NotImplementedError()
+                    constMap[i.result] = vr
+                    i.removeDef(i.result)
+                    i2 = ImmLoad(i.result, vr)
+                    logging.debug('Replacing {} with {}'.format(i, i2))
+                    i.Parent.replaceInstruction(i, i2)
 
 
 class DeadCodeDeleter(BasicBlockPass):