def extract_from_bin(self):
try:
with open(self.file, 'r') as f:
self.dbg_map = pyc_dbg_elf.extract_from_bin(f)
for (name, bloc) in self.dbg_map['blocs'].items():
bloc['symtbl'] = SymTable.from_map(bloc['mem_map'])
except pyc_dbg_elf.ElfError as e:
raise Exception("invalid binary '%s': %s" % (self.file, e))
def __init__(self):
self.symtbl = SymTable()
self.stack_constraints = {}
class Allocator:
def __init__(self):
self.symtbl = SymTable()
self.stack_constraints = {}
def propagate_reg_constraints(self, reg_constraints, graph, node, index):
for n in graph[node]:
if n not in reg_constraints:
continue
reg_constraints[n].add(index)
def propagate_stack_constraints(self, graph, node, index):
for intf_node in graph[node]:
if intf_node not in self.stack_constraints:
self.stack_constraints[intf_node] = set([])
self.stack_constraints[intf_node].add(index)
def dump_vars_to_stack(self, nodes, graph):
reg_amt = len(Register.registers)
i = reg_amt
for n in nodes:
if isinstance(n, Register):
continue
self.symtbl.map(n, i)
self.propagate_stack_constraints(graph, n, i)
i += 1
return i + 1 - reg_amt
def allocate_stack(self, node, graph):
if node.needs_reg:
raise Exception("%s needs a reg!" % repr(node))
i = len(Register.registers)
while node in self.stack_constraints and i in self.stack_constraints[node]:
i = i+1
self.symtbl.map(node, i)
self.propagate_stack_constraints(graph, node, i)
return i
def allocate_reg(self, node, reg_constraints, graph):
regamt = len(Register.registers)
if node not in reg_constraints:
raise Exception("node not in reg constraints!")
i = 0
for i in range(0, regamt):
if i not in reg_constraints[node]:
break
if i >= regamt:
raise Exception("no reg available!")
self.symtbl.map(node, i)
self.propagate_reg_constraints(reg_constraints, graph, node, i)
return i
def next(self, reg_constraints):
if len(reg_constraints) < 1:
raise Exception("reg_constraints is empty!")
regamt = len(Register.registers)
to_sort = []
for (node, v) in reg_constraints.items():
saturation = len(v)
#if its already out of the running on a reg
#there is no need to break ties
if saturation >= regamt:
return (node, True)
to_sort.append( (node, saturation) )
sorted_items = sorted(to_sort, key = lambda (node, v): v, reverse=True )
high_score = sorted_items[0][1]
for (node, score) in sorted_items:
if node.needs_reg == True:
return (node, False)
if score != high_score:
break
return (sorted_items[0][0], False)
def allocate(self, graph):
reg_constraints = {}
self.symtbl.boot_reg_dwellers()
reg_nodes = [Register(x) for x in Register.registers]
todo = set(graph.keys()) - set(self.symtbl.mem_map.keys()) - set(reg_nodes)
for n in todo:
reg_constraints[n] = set([])
#registers are allocated to themselves
for rnode in reg_nodes:
self.propagate_reg_constraints(reg_constraints, graph, rnode, self.symtbl[rnode])
t0 = time.time()
while len(reg_constraints) > 0:
#log("constraints: %s" % repr(self.reg_constraints) )
(node, saturated) = self.next(reg_constraints)
log("allocate reg for %s" % repr(node) )
index = \
self.allocate_stack(node, graph) if saturated \
else self.allocate_reg(node, reg_constraints, graph)
log(" allocated %d" % index)
del reg_constraints[node]
return self.symtbl