diff --git a/typetapper/data.py b/typetapper/data.py index 403679d..b50c75a 100644 --- a/typetapper/data.py +++ b/typetapper/data.py @@ -201,30 +201,24 @@ class DataKind(IntEnum): @dataclass(slots=True) class Prop: self_data: Counter[DataKind] = field(default_factory=Counter) - struct_data: defaultdict[int, defaultdict[int, Counter[DataKind]]] = field(default_factory=lambda: defaultdict(lambda: defaultdict(Counter))) + struct_data: Counter[Tuple[int, int, DataKind]] = field(default_factory=Counter) unifications: Counter[Tuple[int, int]] = field(default_factory=Counter) def update(self, other: 'Prop'): self.self_data.update(other.self_data) - for offset, v1 in other.struct_data.items(): - for size, v2 in v1.items(): - self.struct_data[offset][size].update(v2) + self.struct_data.update(other.struct_data) self.unifications.update(other.unifications) def subtract(self, other: 'Prop'): self.self_data.subtract(other.self_data) - for offset, v1 in other.struct_data.items(): - for size, v2 in v1.items(): - self.struct_data[offset][size].subtract(v2) + self.struct_data.subtract(other.struct_data) self.unifications.subtract(other.unifications) def maximize(self, other: 'Prop'): for key, val in other.self_data.items(): self.self_data[key] = max(self.self_data[key], val) - for offset, v1 in other.struct_data.items(): - for size, v2 in v1.items(): - for kind, val in v2.items(): - self.struct_data[offset][size][kind] = max(self.struct_data[offset][size][kind], val) + for key, val in other.struct_data.items(): + self.struct_data[key] = max(self.struct_data[key], val) for key, val in other.unifications.items(): self.unifications[key] = max(self.unifications[key], val) @@ -238,19 +232,18 @@ class Prop: result = copy.deepcopy(self) for op in ops.ops: if isinstance(op, RefOp): - result.struct_data.clear() - result.struct_data[0][op.size] = result.self_data - result.self_data = Counter() + result.struct_data = Counter({(0, op.size, k): v for k, v in result.self_data.items()}) + result.self_data.clear() result.unifications.clear() elif isinstance(op, DerefOp): - result.self_data = result.struct_data[0][op.size] + result.self_data = Counter({k[2]: v for k, v in result.struct_data.items() if k[0] == 0 and k[1] == op.size}) result.struct_data.clear() result.unifications.clear() elif isinstance(op, ConstOffsetOp): items = list(result.struct_data.items()) result.struct_data.clear() - for k, v in items: - result.struct_data[k - op.const] = v # there is some JANK shit going on with this sign + for (offset, size, kind), v in items: + result.struct_data[(offset - op.const, size, kind)] = v # there is some JANK shit going on with this sign saved = result.self_data.get(DataKind.Pointer, None) result.self_data.clear() if saved: diff --git a/typetapper/engine.py b/typetapper/engine.py index f82bbab..c54ff78 100644 --- a/typetapper/engine.py +++ b/typetapper/engine.py @@ -109,7 +109,7 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): def load(self, addr, size, endness): prop = Prop() prop.self_data[DataKind.Pointer] += 1 - prop.struct_data[0][size][DataKind.GenericData] += 1 + prop.struct_data[(0, size, DataKind.GenericData)] += 1 addr.prop(prop, self.graph) mem_atom = MemoryAtom(self.codeloc, size, endness) @@ -265,7 +265,7 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): def store(self, addr, data, endness): prop = Prop() prop.self_data[DataKind.Pointer] += 1 - prop.struct_data[0][data.size][DataKind.GenericData] += 1 + prop.struct_data[(0, data.size, DataKind.GenericData)] += 1 addr.prop(prop, self.graph) mem_atom = MemoryAtom(self.codeloc, data.size, endness) diff --git a/typetapper/hierarchy_graph.py b/typetapper/hierarchy_graph.py index cd53ba1..6fdf4f5 100644 --- a/typetapper/hierarchy_graph.py +++ b/typetapper/hierarchy_graph.py @@ -1,4 +1,5 @@ from typing import Set, Union, TYPE_CHECKING, List, Optional, Dict, Iterable, Tuple, Callable, Any +import copy from itertools import pairwise, chain from collections import defaultdict @@ -41,6 +42,10 @@ class HierarchicalGraph(RelativeAtomGraph): else: assert isinstance(node, RelativeAtomGroup) assert node is self._root_group or node.parent is self._root_group or node.parent in self.__graph.nodes + prop = Prop() + for child in node.children: + prop.update(self.prop(child)) + assert prop == node.prop pool = set(self.__graph.edges) while pool: @@ -129,8 +134,8 @@ class HierarchicalGraph(RelativeAtomGraph): else: yield node - def _prop_propagate(self, node: RelativeAtomOrGroup, add: bool): - prop = self.prop(node) + def _prop_propagate(self, node: RelativeAtomOrGroup, add: bool, prop: Optional[Prop]=None): + prop = self.prop(node) if prop is None else prop for parent in self._ancestry(node): if add: parent.prop.update(prop) @@ -138,12 +143,17 @@ class HierarchicalGraph(RelativeAtomGraph): parent.prop.subtract(prop) def _add_node(self, relatom: RelativeAtom, path: OpSequence, has_gone_down: bool) -> bool: - res = super()._add_node(relatom, path, has_gone_down) - if relatom not in self.__graph.nodes: + if relatom in self.__graph.nodes: + old_prop = copy.deepcopy(self.prop(relatom)) + else: + old_prop = None self._atom_parents[relatom] = self._current_group self._current_group.children.add(relatom) self.__graph.add_node(relatom) + res = super()._add_node(relatom, path, has_gone_down) if res: + if old_prop is not None: + self._prop_propagate(relatom, False, old_prop) self._prop_propagate(relatom, True) self.check_invariants() return res @@ -189,6 +199,7 @@ class HierarchicalGraph(RelativeAtomGraph): # if item is a group, all edges will have two sides # if item is an atom, all edges will have one side if isinstance(item, RelativeAtom): + # if there is a self-loop, it will appear once from each side yield from ((None, (item, succ, key)) for succ, keys in self.__graph.succ[item].items() for key in keys) yield from (((pred, item, key), None) for pred, keys in self.__graph.pred[item].items() for key in keys) else: @@ -261,14 +272,13 @@ class HierarchicalGraph(RelativeAtomGraph): self.__graph.edges[new_edge][prev_next[outward]] = even_further_edge if even_further_edge: self.__graph.edges[even_further_edge][prev_next[not outward]] = new_edge - self._prop_propagate(item, False) new_parent.children.add(item) parent.children.remove(item) if isinstance(item, RelativeAtomGroup): item.parent = new_parent else: self._atom_parents[item] = new_parent - self._prop_propagate(item, True) + parent.prop.subtract(self.prop(item)) self.check_invariants() def move_node_in(self, item: RelativeAtomOrGroup, new_parent: RelativeAtomGroup): @@ -303,6 +313,10 @@ class HierarchicalGraph(RelativeAtomGraph): outward = False further_edge = self.__graph.edges[outer_edge][prev_next[outward]] + if outer == item: + # self loop! + continue + # we only need to break the outer edge self.__graph.remove_edge(*outer_edge) # is this a collapse or an expand operation? @@ -338,14 +352,13 @@ class HierarchicalGraph(RelativeAtomGraph): self.__graph.edges[new_edge][prev_next[outward]] = even_further_edge if even_further_edge: self.__graph.edges[even_further_edge][prev_next[not outward]] = new_edge - self._prop_propagate(item, False) new_parent.children.add(item) parent.children.remove(item) if isinstance(item, RelativeAtomGroup): item.parent = new_parent else: self._atom_parents[item] = new_parent - self._prop_propagate(item, True) + new_parent.prop.update(self.prop(item)) self.check_invariants() def move_node(self, node: RelativeAtomOrGroup, new_parent: RelativeAtomGroup): diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py index ff60d1f..ea4e516 100644 --- a/typetapper/hierarchy_graph_view.py +++ b/typetapper/hierarchy_graph_view.py @@ -915,7 +915,7 @@ class PropChart(QGraphicsItem): self.scene().removeItem(item) # layout vertical position mapping - offsets_set = {offset + suboffset for offset, v1 in self.prop.struct_data.items() for size, v2 in v1.items() for kind, count in v2.items() if count for suboffset in range(size)} + offsets_set = {offset + suboffset for (offset, size, kind), count in self.prop.struct_data.items() if count for suboffset in range(size)} offsets_set.update(o for offsets in self.prop.unifications for o in offsets) first_offset = min(offsets_set, default=-1) ypos_mapping = {first_offset: self.margin_y} @@ -952,7 +952,7 @@ class PropChart(QGraphicsItem): ticks.append(run_start) # layout boxes - data = [(offset, kind, size, count) for offset, v1 in self.prop.struct_data.items() for size, v2 in v1.items() for kind, count in v2.items() if count] + data = [(offset, kind, size, count) for (offset, size, kind), count in self.prop.struct_data.items() if count] data.sort(key=lambda x: (x[0], x[1], -x[2], x[3])) cols_allocated = defaultdict(int) boxes = [] diff --git a/typetapper/relative_graph.py b/typetapper/relative_graph.py index 6ee238e..0a478b6 100644 --- a/typetapper/relative_graph.py +++ b/typetapper/relative_graph.py @@ -1,3 +1,4 @@ +import copy from typing import TYPE_CHECKING, List, Optional, Tuple, Set from dataclasses import dataclass import logging @@ -25,7 +26,11 @@ class RelativeAtomAttrs: has_gone_down: bool = False def merge(self, other: 'RelativeAtomAttrs'): + changed = False + checkme = copy.deepcopy(self.prop) self.prop.maximize(other.prop) + if checkme != self.prop: + changed = True # TODO has_gone_down if self.path != other.path: # TODO unifications @@ -34,8 +39,8 @@ class RelativeAtomAttrs: score_1 = sum(0 if isinstance(op, (ConstOffsetOp, RefOp, DerefOp)) else 1 for op in other.path.ops) if score_1 < score_0: self.path = other.path - return True - return False + changed = True + return changed class RelativeAtomGraph: def __init__(self, kp: 'TypeTapperManager', baseline: List[Atom]):