def visit_Return(self, node):
''' Substitute return with return variable assignment + break
'''
self.generic_visit(node)
if self._return_var is None:
self._return_var = new_var_name(self)
return [ast.Assign(
targets=[ast.Name(id=self._return_var, ctx=ast.Store())],
value=node.value),
ast.Break()]
def _new_binding_node(self, value):
''' Generate unique variable name, add it to constants with given value,
and return the node that loads generated variable.
'''
literal_node = self._get_literal_node(value)
if literal_node is not None:
return literal_node
else:
var_name = new_var_name(self)
self._constants[var_name] = value
return ast.Name(id=var_name, ctx=ast.Load())
def visit_Name(self, node):
''' Replacing known variables with literal values
'''
self.generic_visit(node)
if node.id in self._locals:
if node.id in self._mangled:
mangled_id = self._mangled[node.id]
else:
mangled_id = new_var_name(self)
self._mangled[node.id] = mangled_id
return ast.Name(id=mangled_id, ctx=node.ctx)
else:
return node
def _inlined_fn(self, node):
''' Return a list of nodes, representing inlined function call,
and a node, repesenting the variable that stores result.
'''
is_known, fn = self._get_node_value_if_known(node.func)
assert is_known
fn_ast = fn_to_ast(fn).body[0]
inliner = Inliner(self._var_count, get_locals(fn_ast))
fn_ast = inliner.visit(fn_ast)
self._var_count = inliner.get_var_count()
inlined_body = []
assert not node.kwargs and not node.starargs # TODO
for callee_arg, fn_arg in zip(node.args, fn_ast.args.args):
# setup mangled values before call
# TODO - if callee_arg is "simple" - literal or name,
# and is never assigned in inlined_body
# then do not make an assignment, just use it in inlined_body
inlined_body.append(ast.Assign(
targets=[ast.Name(id=fn_arg.id, ctx=ast.Store())],
value=callee_arg))
is_known, value = self._get_node_value_if_known(callee_arg)
if is_known:
# TODO - check that mutations are detected
self._constants[fn_arg.id] = value
inlined_code = self._visit(fn_ast.body) # optimize inlined code
if isinstance(inlined_code[-1], ast.Break): # single return
inlined_body.extend(inlined_code[:-1])
else: # multiple returns - wrap in "while"
while_var = new_var_name(self)
inlined_body.extend([
ast.Assign(
targets=[ast.Name(id=while_var, ctx=ast.Store())],
value=self._get_literal_node(True)),
ast.While(
test=ast.Name(id=while_var, ctx=ast.Load()),
body=[
ast.Assign(
targets=[ast.Name(id=while_var, ctx=ast.Store())],
value=self._get_literal_node(False))] +
inlined_code,
orelse=[])
])
all_nodes = inlined_body + \
[ast.Name(id=inliner.get_return_var(), ctx=ast.Load())]
remove_assignments(all_nodes)
return all_nodes[:-1], all_nodes[-1]
