Collected the data to move llvm_cbuilder

llvm_cbuilder/translator.py

@@ -0,0 +1,194 @@
+# A handy translator that converts control flow into the appropriate
+# llvm_cbuilder constructs
+from numba.functions import _get_ast, fix_ast_lineno
+import inspect, functools, ast
+import logging
+
+logger = logging.getLogger(__name__)
+
+def translate(func):
+    # TODO use meta package
+    wrapper = functools.wraps(func)
+    caller_frame = inspect.currentframe().f_back
+    tree = _get_ast(func)
+    tree = ast.Module(body=tree.body)
+    tree = ExpandControlFlow().visit(tree)
+    fix_ast_lineno(tree)
+
+    # prepare locals for execution
+    local_dict = locals()
+    local_dict.update(caller_frame.f_locals)
+    local_dict.update(caller_frame.f_globals)
+
+    try:
+        compiled = compile(tree, '<string>', 'exec')
+        return eval(compiled)
+    except Exception, e:
+        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
+
+
+_if_else_template = '''
+with self.ifelse(__CONDITION__) as _ifelse_:
+    with _ifelse_.then():
+        __THEN__
+    with _ifelse_.otherwise():
+        __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):
+    '''
+    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):
+        mapping = {
+            '__CONDITION__' : node.test,
+            '__THEN__'      : node.body,
+            '__OTHERWISE__' : node.orelse,
+        }
+
+        ifelse = load_template(_if_else_template)
+        ifelse = MacroExpander(mapping).visit(ifelse)
+        newnode = self.generic_visit(ifelse)
+        return newnode
+
+    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 newnode
+
+    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 newnode
+
+    def visit_Return(self, node):
+        mapping = {'__RETURN__' : node.value}
+        ret = load_template(_return_template)
+        repl = MacroExpander(mapping).visit(ret)
+        return repl
+
+    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 = constant
+        return newnode
+
+class MacroExpander(ast.NodeTransformer):
+    def __init__(self, mapping):
+        self.mapping = mapping
+
+    def visit_With(self, node):
+        '''
+        Expand X in the following:
+            with blah:
+                X
+        Nothing should go before or after X.
+        X must be a list of nodes.
+        '''
+        if (len(node.body)==1 # the body of
+          and isinstance(node.body[0], ast.Expr)
+          and isinstance(node.body[0].value, ast.Name)):
+            try:
+                repl = self.mapping.pop(node.body[0].value.id)
+            except KeyError:
+                pass
+            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
+        '''
+
+        try:
+            repl = self.mapping.pop(node.id)
+        except KeyError:
+            pass
+        else:
+            if repl is not None and not isinstance(repl, list):
+                return repl
+        return node
+
+