Add llvm_cbuilder translate features:

* handle while loop
 * handle for loop
 * handle immediate values

Todo: handle "break" and "continue"

translator.py

@@ -1,7 +1,10 @@
 # A handy translator that converts control flow into the appropriate
 # llvm_cbuilder constructs
 import inspect, functools, ast
-import pprint
+import logging
+
+logger = logging.getLogger(__name__)
+
 def translate(func):
     # TODO use meta package
     wrapper = functools.wraps(func)
@@ -29,8 +32,15 @@ def translate(func):
     local_dict.update(caller_frame.f_locals)
     local_dict.update(caller_frame.f_globals)
 
+    try:
+        return eval(compile(tree, '<string>', 'exec'))
+    except:
+        logger.debug(ast.dump(tree))
+        from ArminRonacher import codegen # uses Armin Ronacher's codegen to debug
+        # http://dev.pocoo.org/hg/sandbox/file/852a1248c8eb/ast/codegen.py
+        logger.debug(codegen.to_source(tree))
+        raise
 
-    return eval(compile(tree, '<string>', 'exec'))
 
 _if_else_template = '''
 with self.ifelse(__CONDITION__) as _ifelse_:
@@ -40,14 +50,32 @@ with self.ifelse(__CONDITION__) as _ifelse_:
         __OTHERWISE__
 '''
 
+_while_template = '''
+with self.loop() as _loop_:
+    with _loop_.condition() as _setcond_:
+        _setcond_(__CONDITION__)
+    with _loop_.body():
+        __BODY__
+'''
+
+_for_range_template = '''
+with self.for_range(*__ARGS__) as (_loop_, __ITER__):
+    __BODY__
+'''
+
 _return_template = 'self.ret(__RETURN__)'
 
+_const_int_template = 'self.constant(C.int, __VALUE__)'
+_const_long_template = 'self.constant(C.long, __VALUE__)'
+_const_float_template = 'self.constant(C.double, __VALUE__)'
+
 def load_template(string):
     '''
     Since ast.parse() returns a ast.Module node,
     it is more useful to trim the Module and get to the first item of body
     '''
     tree = ast.parse(string)  # return a Module
+    assert isinstance(tree, ast.Module)
     return tree.body[0]       # get the first item of body
 
 class ExpandControlFlow(ast.NodeTransformer):
@@ -55,8 +83,23 @@ class ExpandControlFlow(ast.NodeTransformer):
     Expand control flow contructs.
     These are the most tedious thing to do in llvm_cbuilder.
     '''
+
+    ## Use breadcumb to track parent nodes
+    #    def __init__(self):
+    #        self.breadcumb = []
+    #
+    #    def visit(self, node):
+    #        self.breadcumb.append(node)
+    #        try:
+    #            return super(ExpandControlFlow, self).visit(node)
+    #        finally:
+    #            self.breadcumb.pop()
+    #
+    #    @property
+    #    def parent(self):
+    #        return self.breadcumb[-2]
+
     def visit_If(self, node):
-        condition = node.test
         mapping = {
             '__CONDITION__' : node.test,
             '__THEN__'      : node.body,
@@ -65,8 +108,36 @@ class ExpandControlFlow(ast.NodeTransformer):
 
         ifelse = load_template(_if_else_template)
         ifelse = MacroExpander(mapping).visit(ifelse)
-        newnode = ast.copy_location(ifelse, node)
+        newnode = self.generic_visit(ifelse)
-        return self.generic_visit(newnode)
+        return ast.copy_location(newnode, node)
+
+    def visit_While(self, node):
+        mapping = {
+            '__CONDITION__' : node.test,
+            '__BODY__'      : node.body,
+        }
+        whileloop = load_template(_while_template)
+        whileloop = MacroExpander(mapping).visit(whileloop)
+        newnode = self.generic_visit(whileloop)
+        return ast.copy_location(newnode, node)
+
+    def visit_For(self, node):
+        try:
+            if node.iter.func.id not in ['range', 'xrange']:
+                return node
+        except AttributeError:
+            return node
+
+        mapping = {
+            '__ITER__' : node.target,
+            '__BODY__' : node.body,
+            '__ARGS__' : ast.Tuple(elts=node.iter.args, ctx=ast.Load()),
+        }
+
+        forloop = load_template(_for_range_template)
+        forloop = MacroExpander(mapping).visit(forloop)
+        newnode = self.generic_visit(forloop)
+        return ast.copy_location(newnode, node)
 
     def visit_Return(self, node):
         mapping = {'__RETURN__' : node.value}
@@ -74,6 +145,24 @@ class ExpandControlFlow(ast.NodeTransformer):
         repl = MacroExpander(mapping).visit(ret)
         return ast.copy_location(repl, node)
 
+    def visit_Num(self, node):
+        '''convert immediate values
+        '''
+        typemap = {
+            int   : _const_int_template,
+            long  : _const_long_template,  # TODO: disable long for py3
+            float : _const_float_template,
+        }
+
+        template = load_template(typemap[type(node.n)])
+
+        mapping = {
+            '__VALUE__' : node,
+        }
+        constant = MacroExpander(mapping).visit(template).value
+        newnode = ast.copy_location(constant, node)
+        return newnode
+
 class MacroExpander(ast.NodeTransformer):
     def __init__(self, mapping):
         self.mapping = mapping
@@ -96,19 +185,20 @@ class MacroExpander(ast.NodeTransformer):
             else:
                 old = node.body[0]
                 node.body = repl
+
         return self.generic_visit(node) # recursively apply expand all macros
 
     def visit_Name(self, node):
         '''
         Expand all Name node to simple value
         '''
-        if type(node.ctx) is ast.Load:
+
-            try:
+        try:
-                repl = self.mapping.pop(node.id)
+            repl = self.mapping.pop(node.id)
-            except KeyError:
+        except KeyError:
-                pass
+            pass
-            else:
+        else:
-                if repl is not None and not isinstance(repl, list):
+            if repl is not None and not isinstance(repl, list):
-                    return ast.copy_location(repl, node)
+                return ast.copy_location(repl, node)
         return node