Codegen for unary, ifexpr bugfix, mandelbrot sample

chapter6.py

@@ -555,6 +555,11 @@ class LLVMCodeGenerator(object):
     def _codegen_VariableExprAST(self, node):
         return self.func_symtab[node.name]
 
+    def _codegen_UnaryExprAST(self, node):
+        operand = self._codegen(node.operand)
+        func = self.module.get_global('unary{0}'.format(node.op))
+        return self.builder.call(func, [operand], 'unop')
+
     def _codegen_BinaryExprAST(self, node):
         lhs = self._codegen(node.lhs)
         rhs = self._codegen(node.rhs)
@@ -603,6 +608,9 @@ class LLVMCodeGenerator(object):
         self.builder.function.basic_blocks.append(else_bb)
         self.builder.position_at_start(else_bb)
         else_val = self._codegen(node.else_expr)
+
+        # Emission of else_val could have modified the current basic block.
+        else_bb = self.builder.block
         self.builder.branch(merge_bb)
 
         # Emit the merge ('ifcnt') block
@@ -926,14 +934,77 @@ class TestEvaluator(unittest.TestCase):
         self.assertEqual(e.evaluate('10 % 5'), 5)
         self.assertEqual(e.evaluate('100 % 5.5'), 94.5)
 
+    def test_custom_unop(self):
+        e = KaleidoscopeEvaluator()
+        e.evaluate('def unary!(a) 0 - a')
+        e.evaluate('def unary^(a) a * a')
+        self.assertEqual(e.evaluate('!10'), -10)
+        self.assertEqual(e.evaluate('^10'), 100)
+        self.assertEqual(e.evaluate('!^10'), -100)
+        self.assertEqual(e.evaluate('^!10'), 100)
+
+    def test_mixed_ops(self):
+        e = KaleidoscopeEvaluator()
+        e.evaluate('def unary!(a) 0 - a')
+        e.evaluate('def unary^(a) a * a')
+        e.evaluate('def binary %(a b) a - b')
+        self.assertEqual(e.evaluate('!10 % !20'), 10)
+        self.assertEqual(e.evaluate('^(!10 % !20)'), 100)
+
+
+def generate_mandelbrot():
+    e = KaleidoscopeEvaluator()
+    e.evaluate('def unary- (v) 0 - v')
+    e.evaluate('def binary> 10 (lhs rhs) rhs < lhs')
+    e.evaluate('def binary: 1 (x y) y')
+    e.evaluate('''
+        def binary| 5 (lhs rhs)
+            if lhs then 1 else if rhs then 1 else 0
+        ''')
+    e.evaluate('''
+        def printdensity(d)
+            if d > 8 then
+                putchard(32) # ' '
+            else if d > 4 then
+                putchard(46) # '.'
+            else if d > 2 then
+                putchard(43) # '+'
+            else
+                putchard(42) # '*'
+        ''')
+    e.evaluate('''
+        def mandelconverger(real imag iters creal cimag)
+            if iters > 255 | (real*real + imag*imag > 4) then
+                iters
+            else
+                mandelconverger(real*real - imag*imag + creal,
+                                2*real*imag + cimag,
+                                iters+1, creal, cimag)
+        ''')
+    e.evaluate('''
+        def mandelconverge(real imag)
+            mandelconverger(real, imag, 0, real, imag)
+        ''')
+    e.evaluate('''
+        def mandelhelp(xmin xmax xstep ymin ymax ystep)
+            for y = ymin, y < ymax, ystep in (
+                (for x = xmin, x < xmax, xstep in
+                    printdensity(mandelconverge(x, y)))
+                : putchard(10))
+        ''')
+    e.evaluate('''
+        def mandel(realstart imagstart realmag imagmag)
+            mandelhelp(realstart, realstart+realmag*78, realmag,
+                       imagstart, imagstart+imagmag*48, imagmag)
+        ''')
+    e.evaluate('mandel(-2.3, -1.3, 0.05, 0.07)')
+
 
 if __name__ == '__main__':
-    p = Parser()
+    generate_mandelbrot()
-    print(p.parse_toplevel('def binary% 77(a b) a + b').dump())
+    #p = Parser()
-    print(p.parse_toplevel('def fra(x t) x % t').dump())
+    #print(p.parse_toplevel('def binary% 77(a b) a + b').dump())
-    kalei = KaleidoscopeEvaluator()
+    #print(p.parse_toplevel('def fra(x t) x % t').dump())
-    kalei.evaluate('def binary% 77(a b) a + b')
+    #kalei = KaleidoscopeEvaluator()
-    print(kalei.evaluate('5 % 10', optimize=False, llvmdump=True))
+    #kalei.evaluate('def binary% 77(a b) a + b')
-
+    #print(kalei.evaluate('5 % 10', optimize=False, llvmdump=True))
-    #kalei.evaluate('def foo(a b) for x = 65, x < a, b in putchard(x)')
-    #print(kalei.evaluate('foo(79, 1)', optimize=True, llvmdump=True))