def minmax(typesystem, args, op):
if len(args) < 2:
return
res = args[0]
for arg in args[1:]:
lhs_type = get_type(res)
rhs_type = get_type(arg)
res_type = typesystem.promote(lhs_type, rhs_type)
if lhs_type != res_type:
res = nodes.CoercionNode(res, res_type)
if rhs_type != res_type:
arg = nodes.CoercionNode(arg, res_type)
lhs_temp = nodes.TempNode(res_type)
rhs_temp = nodes.TempNode(res_type)
res_temp = nodes.TempNode(res_type)
lhs = lhs_temp.load(invariant=True)
rhs = rhs_temp.load(invariant=True)
expr = ast.IfExp(ast.Compare(lhs, [op], [rhs]), lhs, rhs)
body = [
ast.Assign([lhs_temp.store()], res),
ast.Assign([rhs_temp.store()], arg),
ast.Assign([res_temp.store()], expr),
]
res = nodes.ExpressionNode(body, res_temp.load(invariant=True))
return res
def rewrite_enumerate(self, node):
"""
Rewrite a loop like
for i, x in enumerate(array[, start]):
...
into
_arr = array
[_s = start]
for _i in range(len(_arr)):
i = _i [+ _s]
x = _arr[_i]
...
"""
call = node.iter
if (len(call.args) not in (1, 2) or call.keywords or call.starargs
or call.kwargs):
self.error(call, 'expected 1 or 2 arguments to enumerate()')
target = node.target
if (not isinstance(target, (ast.Tuple, ast.List))
or len(target.elts) != 2):
self.error(call, 'expected 2 iteration variables')
array = call.args[0]
start = call.args[1] if len(call.args) > 1 else None
idx = target.elts[0]
var = target.elts[1]
array_temp = untypedTemp()
if start:
start_temp = untypedTemp() # TODO: only allow integer start
idx_temp = nodes.TempNode(typesystem.Py_ssize_t)
# for _i in range(len(_arr)):
node.target = idx_temp.store()
node.iter = ast.Call(ast.Name('range', ast.Load()), [
ast.Call(ast.Name('len', ast.Load()), [array_temp.load(True)], [],
None, None)
], [], None, None)
# i = _i [+ _s]
new_idx = idx_temp.load()
if start:
new_idx = ast.BinOp(new_idx, ast.Add(), start_temp.load(True))
node.body.insert(0, ast.Assign([idx], new_idx))
# x = _arr[_i]
value = ast.Subscript(array_temp.load(True),
ast.Index(idx_temp.load()), ast.Load())
node.body.insert(1, ast.Assign([var], value))
# _arr = array; [_s = start]; ...
body = [ast.Assign([array_temp.store()], array), node]
if start:
body.insert(1, ast.Assign([start_temp.store()], start))
return map(self.visit, body)
def allocate_struct_on_stack(self, assmnt_node, target):
# Allocate struct on stack
temp = nodes.TempNode(target.type)
assmnt_node.targets[0] = temp.store()
assmnt_node.value = self.visit(assmnt_node.value)
# Expose LLVM value through SSA (patch the Variable or the
# LHS). We need to store the pointer to the struct (the alloca)
ssa_assmnt = ast.Assign(targets=[target], value=temp.store())
return ast.Suite(body=[assmnt_node, ssa_assmnt])
def rewrite_zip(self, node):
"""
Rewrite a loop like
for x, y... in zip(xs, ys...):
...
into
_xs = xs; _ys = ys...
for _i in range(min(len(_xs), len(_ys)...)):
x = _xs[_i]; y = _ys[_i]...
...
"""
call = node.iter
if not call.args or call.keywords or call.starargs or call.kwargs:
self.error(call, 'expected at least 1 argument to zip()')
target = node.target
if (not isinstance(target, (ast.Tuple, ast.List))
or len(target.elts) != len(call.args)):
self.error(call,
'expected %d iteration variables' % len(call.args))
temps = [untypedTemp() for _ in xrange(len(call.args))]
idx_temp = nodes.TempNode(typesystem.Py_ssize_t)
# min(len(_xs), len(_ys)...)
len_call = ast.Call(ast.Name('min', ast.Load()), [
ast.Call(ast.Name('len', ast.Load()), [tmp.load(True)], [], None,
None) for tmp in temps
], [], None, None)
# for _i in range(...):
node.target = idx_temp.store()
node.iter = ast.Call(ast.Name('range', ast.Load()), [len_call], [],
None, None)
# x = _xs[_i]; y = _ys[_i]...
node.body = [ast.Assign([tgt],
ast.Subscript(tmp.load(True),
ast.Index(idx_temp.load()),
ast.Load()))
for tgt, tmp in zip(target.elts, temps)] + \
node.body
# _xs = xs; _ys = ys...
body = [ast.Assign([tmp.store()], arg)
for tmp, arg in zip(temps, call.args)] + \
[node]
return map(self.visit, body)
def getiter(self, context, for_node, llvm_module):
self.index = nodes.TempNode(Py_ssize_t, "iterator_index")
return assign(self.index, nodes.const(0, Py_ssize_t))
def rewrite_range_iteration(self, node):
"""
Handle range iteration:
for i in range(start, stop, step):
...
becomes
nsteps = compute_nsteps(start, stop, step)
temp = 0
while temp < nsteps:
target = start + temp * step
...
temp += 1
"""
self.generic_visit(node)
temp = nodes.TempNode(node.target.type, 'target_temp')
nsteps = nodes.TempNode(Py_ssize_t, 'nsteps')
start, stop, step = unpack_range_args(node.iter)
if isinstance(step, nodes.ConstNode):
have_step = step.pyval != 1
else:
have_step = True
start, stop, step = [
nodes.CloneableNode(n) for n in (start, stop, step)
]
if have_step:
compute_nsteps = """
$length = {{stop}} - {{start}}
{{nsteps}} = $length / {{step}}
if {{nsteps_load}} * {{step}} != $length: #$length % {{step}}:
# Test for truncation
{{nsteps}} = {{nsteps_load}} + 1
# print "nsteps", {{nsteps_load}}
"""
else:
compute_nsteps = "{{nsteps}} = {{stop}} - {{start}}"
if node.orelse:
else_clause = "else: {{else_body}}"
else:
else_clause = ""
templ = textwrap.dedent("""
%s
{{temp}} = 0
while {{temp_load}} < {{nsteps_load}}:
{{target}} = {{start}} + {{temp_load}} * {{step}}
{{body}}
{{temp}} = {{temp_load}} + 1
%s
""") % (textwrap.dedent(compute_nsteps), else_clause)
# Leave the bodies empty, they are already analyzed
body = ast.Suite(body=[])
else_body = ast.Suite(body=[])
#--------------------------------------------------------------------
# Substitute template and infer types
#--------------------------------------------------------------------
result = self.run_template(templ,
vars=dict(length=Py_ssize_t),
start=start,
stop=stop,
step=step,
nsteps=nsteps.store(),
nsteps_load=nsteps.load(),
temp=temp.store(),
temp_load=temp.load(),
target=node.target,
body=body,
else_body=else_body)
#--------------------------------------------------------------------
# Patch the body and else clause
#--------------------------------------------------------------------
body.body.extend(node.body)
else_body.body.extend(node.orelse)
while_node = result.body[-1]
assert isinstance(while_node, ast.While)
target_increment = while_node.body[-1]
assert isinstance(target_increment, ast.Assign)
# Add target variable increment basic block
node.incr_block.body = [target_increment]
while_node.body[-1] = node.incr_block
#--------------------------------------------------------------------
# Create a While with the ForNode's cfg blocks merged in
#--------------------------------------------------------------------
while_node = make_while_loop(while_node)
copy_basic_blocks(node, while_node)
while_node = nodes.build_while(**vars(while_node))
# Create the place to jump to for 'continue'
while_node.continue_block = node.incr_block
# Set the new while loop in the templated Suite
result.body[-1] = while_node
return result
def rewrite_array_iteration(self, node):
"""
Convert 1D array iteration to for-range and indexing:
for value in my_array:
...
becomes
for i in my_array.shape[0]:
value = my_array[i]
...
"""
logger.debug(ast.dump(node))
orig_target = node.target
orig_iter = node.iter
#--------------------------------------------------------------------
# Replace node.target with a temporary
#--------------------------------------------------------------------
target_name = orig_target.id + '.idx'
target_temp = nodes.TempNode(Py_ssize_t)
node.target = target_temp.store()
#--------------------------------------------------------------------
# Create range(A.shape[0])
#--------------------------------------------------------------------
call_func = ast.Name(id='range', ctx=ast.Load())
nodes.typednode(call_func, typesystem.range_)
shape_index = ast.Index(nodes.ConstNode(0, typesystem.Py_ssize_t))
shape_index.type = typesystem.npy_intp
stop = ast.Subscript(value=nodes.ShapeAttributeNode(orig_iter),
slice=shape_index,
ctx=ast.Load())
nodes.typednode(stop, npy_intp)
#--------------------------------------------------------------------
# Create range iterator and replace node.iter
#--------------------------------------------------------------------
call_args = [
nodes.ConstNode(0, typesystem.Py_ssize_t),
nodes.CoercionNode(stop, typesystem.Py_ssize_t),
nodes.ConstNode(1, typesystem.Py_ssize_t),
]
node.iter = ast.Call(func=call_func, args=call_args)
nodes.typednode(node.iter, call_func.type)
node.index = target_temp.load(invariant=True)
#--------------------------------------------------------------------
# Add assignment to new target variable at the start of the body
#--------------------------------------------------------------------
index = ast.Index(value=node.index)
index.type = target_temp.type
subscript = ast.Subscript(value=orig_iter, slice=index, ctx=ast.Load())
nodes.typednode(subscript, get_type(orig_iter).dtype)
#--------------------------------------------------------------------
# Add assignment to new target variable at the start of the body
#--------------------------------------------------------------------
assign = ast.Assign(targets=[orig_target], value=subscript)
node.body = [assign] + node.body
#--------------------------------------------------------------------
# Specialize new for loop through range iteration
#--------------------------------------------------------------------
return self.visit(node)
def rewrite_range_iteration(self, node):
"""
Handle range iteration:
for i in range(start, stop, step):
...
becomes
nsteps = compute_nsteps(start, stop, step)
temp = 0
while temp < nsteps:
target = start + temp * step
...
temp += 1
"""
self.generic_visit(node)
temp = nodes.TempNode(node.target.type, 'target_temp')
nsteps = nodes.TempNode(Py_ssize_t, 'nsteps')
start, stop, step = unpack_range_args(node.iter)
if isinstance(step, nodes.ConstNode):
have_step = step.pyval != 1
else:
have_step = True
start, stop, step = map(nodes.CloneableNode, (start, stop, step))
if have_step:
templ = textwrap.dedent("""
{{temp}} = 0
{{nsteps}} = ({{stop}} - {{start}} + {{step}} -
(1 if {{step}} >= 0 else -1)) / {{step}}
while {{temp_load}} < {{nsteps_load}}:
{{target}} = {{start}} + {{temp_load}} * {{step}}
{{temp}} = {{temp_load}} + 1
{{body}}
""")
else:
templ = textwrap.dedent("""
{{temp}} = {{start}}
{{nsteps}} = {{stop}}
while {{temp_load}} < {{nsteps_load}}:
{{target}} = {{temp_load}}
{{temp}} = {{temp_load}} + 1
{{body}}
""")
if node.orelse:
templ += "\nelse: {{else_body}}"
# Leave the bodies empty, they are already analyzed
body = ast.Suite(body=[])
else_body = ast.Suite(body=[])
#--------------------------------------------------------------------
# Substitute template and infer types
#--------------------------------------------------------------------
result = self.run_template(templ,
vars=dict(length=Py_ssize_t),
start=start,
stop=stop,
step=step,
nsteps=nsteps.store(),
nsteps_load=nsteps.load(),
temp=temp.store(),
temp_load=temp.load(),
target=node.target,
body=body,
else_body=else_body)
ast.copy_location(result, node)
if hasattr(node, 'lineno'):
visitor = missing.FixMissingLocations(node.lineno,
node.col_offset,
override=True)
visitor.visit(result)
#--------------------------------------------------------------------
# Patch the body and else clause
#--------------------------------------------------------------------
body.body.extend(node.body)
else_body.body.extend(node.orelse)
while_node = result.body[-1]
assert isinstance(while_node, ast.While)
#--------------------------------------------------------------------
# Create a While with the ForNode's cfg blocks merged in
#--------------------------------------------------------------------
while_node = make_while_loop(while_node)
copy_basic_blocks(node, while_node)
while_node = nodes.build_while(**vars(while_node))
# Create the place to jump to for 'continue'
while_node.continue_block = node.cond_block
# Set the new while loop in the templated Suite
result.body[-1] = while_node
return result
def untypedTemp():
"Temp node with a yet unknown type"
type = typesystem.DeferredType(None)
temp = nodes.TempNode(type)
type.variable = temp.variable
return temp