class HeapifyVisitor(CopyVisitor):
def __init__(self):
super(HeapifyVisitor,self).__init__()
# Initialize helper visitors
self.local_visitor = LocalVarsVisitor()
self.free_visitor = FreeVarsVisitor()
self.nested_visitor = NestedFreeVarsVisitor()
def preorder(self, tree, *args):
self.local_visitor.preorder(tree)
self.free_visitor.preorder(tree)
# Global set of variables needing heapification
self.needs_heapification = set([])
return super(HeapifyVisitor, self).preorder(tree, *args)
def visitModule(self, n):
self.needs_heapification = self.needs_heapification | self.nested_visitor.preorder(n.node)
lvs = n.local_vars
new_stmts = []
for local in lvs:
if local in self.needs_heapification:
new_stmts.append(VarAssign(local, List([ZERO])))
body = self.dispatch(n.node)
body.nodes = new_stmts + body.nodes
return Module(body)
def visitSLambda(self, n):
lvs = n.local_vars
self.needs_heapification = self.needs_heapification | self.nested_visitor.preorder(n.code)
code = self.dispatch(n.code)
new_params = []
new_stmts = []
for param in n.params:
if param in self.needs_heapification:
new_param = generate_name('heaped_at_')+param
new_stmts.append(VarAssign(param, List([ZERO])))
new_stmts.append(SubscriptAssign(Name(param), [ZERO], Name(new_param)))
new_params.append(new_param)
else:
new_params.append(param)
for local in lvs:
if local in self.needs_heapification:
new_stmts.append(VarAssign(local, List([ZERO])))
ret = SLambda(new_params, StmtList(new_stmts + code.nodes))
ret.free_vars = n.free_vars
return ret
def visitVarAssign(self, n):
if n.target in self.needs_heapification:
return SubscriptAssign(Name(n.target), [ZERO], self.dispatch(n.value))
else: return VarAssign(n.target, self.dispatch(n.value))
def visitName(self, n):
if n.name in self.needs_heapification:
return Subscript(Name(n.name), [ZERO])
else: return n
def visitLambda(self, n):
free_in_fun = FreeVarsVisitor().preorder(n.code)
for x in free_in_fun - set(n.argnames):
needs_heapify[x] = True
def visitFunction(self, n):
local_vars = FindLocalsVisitor().preorder(n.code)
free_in_fun = FreeVarsVisitor().preorder(n.code)
for x in (free_in_fun - local_vars) - set(n.argnames):
needs_heapify[x] = True
def __init__(self):
super(HeapifyVisitor,self).__init__()
# Initialize helper visitors
self.local_visitor = LocalVarsVisitor()
self.free_visitor = FreeVarsVisitor()
self.nested_visitor = NestedFreeVarsVisitor()
