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]