diff --git a/typetapper/data.py b/typetapper/data.py index 243baa1..403679d 100644 --- a/typetapper/data.py +++ b/typetapper/data.py @@ -134,6 +134,20 @@ class OpSequence: def invert(self) -> 'OpSequence': return OpSequence(tuple(x.invert() for x in reversed(self.ops))) + def compute_unifications(self) -> List[Tuple[int, int]]: + base_offset = 0 + strides = [] + for op in self.ops: + if isinstance(op, ConstOffsetOp): + base_offset += op.const + elif isinstance(op, StrideOffsetOp): + strides.append(op.stride) + else: + base_offset = 0 + strides = [] + + return [(base_offset, base_offset + stride) for stride in strides] + def simplify_op_sequence(seq: List[Op]): i = 0 while i < len(seq): @@ -227,11 +241,11 @@ class Prop: result.struct_data.clear() result.struct_data[0][op.size] = result.self_data result.self_data = Counter() - self.unifications.clear() + result.unifications.clear() elif isinstance(op, DerefOp): result.self_data = result.struct_data[0][op.size] result.struct_data.clear() - self.unifications.clear() + result.unifications.clear() elif isinstance(op, ConstOffsetOp): items = list(result.struct_data.items()) result.struct_data.clear() @@ -300,21 +314,29 @@ class LiveData: return LiveData(loc, self.sources + other.sources, size, strides) def commit(self, target: Atom, graph: networkx.DiGraph): + prop = Prop() for src, seq in self.sources: + for start, end in seq.compute_unifications(): + prop.unifications[(start, end)] += 1 graph.add_edge(src, target, ops=seq, cf=[]) + self.atom_prop(target, prop, graph) def prop(self, prop: Prop, graph: networkx.DiGraph): for atom, ops in self.sources: tprop = prop.transform(ops.invert()) - try: - eprop: Prop = graph.nodes[atom].get('prop') - except KeyError: - graph.add_node(atom, prop=tprop) + self.atom_prop(atom, tprop, graph) + + @staticmethod + def atom_prop(atom, tprop: Prop, graph: networkx.DiGraph): + try: + eprop: Prop = graph.nodes[atom].get('prop') + except KeyError: + graph.add_node(atom, prop=tprop) + else: + if eprop: + eprop.update(tprop) else: - if eprop: - eprop.update(tprop) - else: - graph.nodes[atom]['prop'] = tprop + graph.nodes[atom]['prop'] = tprop def prop_self(self, kind: DataKind, graph: networkx.DiGraph): prop = Prop() diff --git a/typetapper/engine.py b/typetapper/engine.py index e745e02..f82bbab 100644 --- a/typetapper/engine.py +++ b/typetapper/engine.py @@ -173,7 +173,6 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): else: strideC[args[1]] += sign - neg1 = args[1] if sign == -1: neg1 = neg1.appended(self.codeloc, NegOp(), neg1.size, tuple((k, -n) for k, n in neg1.strides)) diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py index 48c28c1..d39cc5c 100644 --- a/typetapper/hierarchy_graph_view.py +++ b/typetapper/hierarchy_graph_view.py @@ -1,4 +1,4 @@ -from itertools import chain +from itertools import chain, pairwise from typing import Tuple, Optional, Union, Dict, List, Set from collections import defaultdict from math import sqrt, sin, cos, pi @@ -10,7 +10,8 @@ import logging from PySide6.QtCore import QRectF, QPointF, QTimeLine, QEasingCurve, QMarginsF, QTimer, Qt, QEvent from PySide6.QtGui import QColor, QFont, QBrush, QPainterPath, QPen, QAction, QShortcut from PySide6.QtWidgets import QGraphicsItem, QGraphicsRectItem, QGraphicsSimpleTextItem, QHBoxLayout, QGraphicsScene, \ - QGraphicsEllipseItem, QGraphicsSceneMouseEvent, QMenu, QDialog, QVBoxLayout, QLineEdit, QInputDialog, QMessageBox + QGraphicsEllipseItem, QGraphicsSceneMouseEvent, QMenu, QDialog, QVBoxLayout, QLineEdit, QInputDialog, QMessageBox, \ + QGraphicsLineItem, QGraphicsTextItem import networkx import pygraphviz @@ -40,6 +41,7 @@ PLUS_ICON_COLOR = QColor(0xff, 0xff, 0xff) PLUS_HOVERED_COLOR = QColor(0x40, 0xc0, 0xc0) PLUS_BORDER_COLOR = QColor(0x30, 0xc0, 0x50) PLUS_ICON_FONT = QFont("default", 12) +ARC_COLOR = QColor(0xff, 0xff, 0xff) SCENE_MARGIN = 200 DPI = 72 @@ -733,8 +735,10 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): else: l.error("How'd you select more than one item during expansion") elif not self._selection_event_occurring: - self.selected_nodes.am_obj = {model for model, qnode in zip(self.ug_lookup, self.qnodes) if qnode.isSelected()} - self.selected_nodes.am_event(src='qt') + if self.items(): + # don't do this if we are closing the pane (there will be a desync between items() and qnodes) + self.selected_nodes.am_obj = {model for model, qnode in zip(self.ug_lookup, self.qnodes) if qnode.isSelected()} + self.selected_nodes.am_event(src='qt') for qedge in self.qedges: qedge.update() @@ -849,7 +853,7 @@ class HGNode(AnimatableItem): class PropChart(QGraphicsItem): - def __init__(self, parent, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5.): + def __init__(self, parent, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5., gap_size=15., tick_granularity=4): self.max_width = max_width self.default_unit = default_unit self.unit = default_unit @@ -857,6 +861,8 @@ class PropChart(QGraphicsItem): self.margin_x = margin_x self.margin_y = margin_y self.padding_left = padding_left + self.gap_size = gap_size + self.tick_granularity = tick_granularity self.width = 0. self.height = 0. self.objects = [] @@ -876,67 +882,114 @@ class PropChart(QGraphicsItem): def paint(self, painter, option, widget=...): pass + def layout(self): + raise NotImplementedError + def _layout_marks(self): self.objects.clear() for item in list(self.childItems()): 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.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} + ypos_max = self.margin_y + ticks = [] + gaps = [] + run_start = first_offset + if first_offset % self.tick_granularity == 0: + ticks.append(first_offset) + + for last_offset, offset in pairwise(sorted(offsets_set)): + if offset - last_offset > self.tick_granularity: + if len(ticks) == 0 or ticks[-1] < run_start: + ticks.append(run_start) + + run_start = offset + gaps.append(offset) + ypos_mapping[offset] = ypos_mapping[last_offset] + self.byte_height + self.gap_size + + if offset % self.tick_granularity == 0: + ticks.append(offset) + else: + ypos_mapping[offset] = ypos_mapping[last_offset] + (offset - last_offset) * self.byte_height + + for suboffset in range(offset, last_offset, -1): + if suboffset % self.tick_granularity == 0: + ticks.append(suboffset) + if suboffset != offset: + ypos_mapping[suboffset] = ypos_mapping[offset] - (offset - suboffset) * self.byte_height + + ypos_max = ypos_mapping[offset] + self.byte_height + + if offsets_set and (len(ticks) == 0 or ticks[-1] < run_start): + 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.sort(key=lambda x: (x[0], x[1], -x[2], x[3])) cols_allocated = defaultdict(int) - row_allocated = 0 - offset_allocated = None - marks = [] - ticks = [] + boxes = [] for offset, kind, size, count in data: xpos = max(cols_allocated[suboffset] for suboffset in range(offset, offset + size)) for suboffset in range(offset, offset + size): cols_allocated[suboffset] = xpos + count - if offset_allocated is None or offset >= offset_allocated: - ypos = row_allocated - row_allocated += size - offset_allocated = offset + size - ticks.append((offset, ypos)) - else: - ypos = row_allocated - (offset_allocated - offset) - stretch = max(0, offset + size - offset_allocated) - row_allocated += stretch - offset_allocated += stretch - - marks.append((offset, kind, size, count, xpos, ypos)) + boxes.append((offset, kind, size, count, xpos)) total_width = max(cols_allocated.values()) if cols_allocated else 0 if total_width * self.unit > self.max_width: self.unit = self.max_width / total_width + # create base object self.width = self.max_width + 2 * self.margin_x + self.padding_left - self.height = row_allocated * self.byte_height + 2 * self.margin_y + self.height = ypos_max + self.margin_y box = QGraphicsRectItem(QRectF(0, 0, self.width, self.height), self) box.setBrush(BOX_COLOR) box.setPen(QPen(BOX_BORDER_COLOR, BOX_BORDER_WIDTH)) self.objects.append(box) self.object_bg = box - for offset, kind, size, count, xpos, ypos in marks: - rect = QGraphicsRectItem(QRectF(self.margin_x + self.padding_left + xpos * self.unit, self.margin_y + ypos * self.byte_height, count * self.unit, size * self.byte_height), self) + # create box objects + for offset, kind, size, count, xpos in boxes: + rect = QGraphicsRectItem(QRectF(self.margin_x + self.padding_left + xpos * self.unit, ypos_mapping[offset], count * self.unit, size * self.byte_height), self) rect.setBrush(kind_to_color(kind)) rect.setPen(QColor(0x10, 0x10, 0x10)) self.objects.append(rect) + # create tick objects tick_width = 0 - for offset, ypos in ticks: - tick = QGraphicsSimpleTextItem(str(offset), self) + max_tick_ypos = 0 + for offset in ticks: + tick = QGraphicsSimpleTextItem(format(offset, 'x'), self) tick.setBrush(TEXT_COLOR) tick.setFont(TEXT_FONT) - width = tick.boundingRect().width() + rect = tick.boundingRect() + width = rect.width() tick_width = max(width, tick_width) - tick.setPos(QPointF(self.margin_x - width, self.margin_y + ypos * self.byte_height)) + tick.setPos(QPointF(self.margin_x - width, ypos_mapping[offset] - rect.height() / 2 + self.byte_height / 2)) + max_tick_ypos = max(rect.height() + tick.y(), max_tick_ypos) self.objects.append(tick) - if tick_width > self.margin_x: - shift = tick_width - self.margin_x + ticktick = QGraphicsLineItem(self.margin_x + self.padding_left / 2, ypos_mapping[offset], self.margin_x + self.padding_left, ypos_mapping[offset], self) + self.objects.append(ticktick) + + # create squiggle objects + for gap in gaps: + squiggle = SquiggleMark(self.margin_x + self.padding_left - 1.5, ypos_mapping[gap] - 12.5, self) + self.objects.append(squiggle) + + # create arc objects + for offset1, offset2 in self.prop.unifications: + arc = UnificationArc(self.margin_x + self.padding_left, ypos_mapping[offset1], ypos_mapping[offset2], self) + self.objects.append(arc) + + # if during rendering ticks we ran out of space horizontally, expand (and shift everything to the right of it) + if tick_width > self.margin_x - self.padding_left: + shift = tick_width - (self.margin_x - self.padding_left) for obj in self.objects: if obj is self.object_bg: rect = obj.rect() @@ -946,6 +999,14 @@ class PropChart(QGraphicsItem): obj.setX(obj.x() + shift) self.width += shift + # if during rendering ticks we ran out of space vertically, expand + if max_tick_ypos + self.margin_y > self.height: + shift = max_tick_ypos + self.margin_y - self.height + rect = self.object_bg.rect() + rect.setHeight(rect.height() + shift) + self.object_bg.setRect(rect) + self.height += shift + class PropChartHG(HGNode, PropChart): def __init__(self, parent, **kwargs): self.object_label = None @@ -973,8 +1034,17 @@ class PropChartHG(HGNode, PropChart): self.setY(v.y() - self.height / 2) def layout(self): - lbl_height = 20 self._layout_marks() + self.object_label = QGraphicsTextItem(self.label, self) + self.object_label.setTextWidth(self.width - self.margin_x * 2) + self.object_label.setPos(QPointF(self.margin_x, self.margin_y)) + self.object_label.setFont(TEXT_FONT) + self.object_label.setDefaultTextColor(TEXT_COLOR) + self.objects.append(self.object_label) + self.object_btn = PlusButton(self, self.begin_expand, 16) + self.object_btn.setPos(self.width - 18, 2) + lbl_height = self.object_label.boundingRect().height() + self.height += lbl_height for obj in self.objects: if obj is self.object_bg: @@ -982,15 +1052,13 @@ class PropChartHG(HGNode, PropChart): rect = obj.rect() rect.setHeight(obj.rect().height() + lbl_height) obj.setRect(rect) + elif obj is self.object_label: + pass + elif obj is self.object_btn: + pass else: obj.setY(obj.y() + lbl_height) - self.object_label = QGraphicsSimpleTextItem(self.label, self) - self.object_label.setPos(QPointF(self.margin_x, self.margin_y)) - self.object_label.setFont(TEXT_FONT) - self.objects.append(self.object_label) - self.object_btn = PlusButton(self, self.begin_expand, 16) - self.object_btn.setPos(self.width - 18, 2) self.setTransformOriginPoint(self.center - self.pos()) @@ -1274,6 +1342,58 @@ class PlusButton(QGraphicsItem): self.object_bg.setBrush(PLUS_COLOR) super().hoverLeaveEvent(event) +class UnificationArc(QGraphicsItem): + def __init__(self, xpos, ypos1, ypos2, parent, bent=10., thickness=3.): + super().__init__(parent) + self.setPos(xpos, ypos1) + self.height = ypos2 - ypos1 + self.bent = bent + self.thickness = thickness + self.setOpacity(0.5) + self.setAcceptHoverEvents(True) + + def shape(self): + result = QPainterPath() + result.moveTo(0, 0) + result.lineTo(-self.thickness, 0) + result.cubicTo(-self.bent, self.height / 2, -self.bent, self.height / 2, -self.thickness, self.height) + result.lineTo(0, self.height) + result.cubicTo(-self.bent + self.thickness, self.height / 2, -self.bent + self.thickness, self.height / 2, 0, 0) + result.closeSubpath() + return result + + def paint(self, painter, option, widget=...): + painter.fillPath(self.shape(), ARC_COLOR) + + def boundingRect(self): + return QRectF(-self.bent, 0, self.bent, self.height) + + def hoverEnterEvent(self, event): + self.setOpacity(1) + + def hoverLeaveEvent(self, event): + self.setOpacity(0.5) + +class SquiggleMark(QGraphicsItem): + def __init__(self, xpos, ypos, parent, nlines=5, dx=3., dy=2., thickness=1.): + super().__init__(parent) + + self.setPos(xpos, ypos) + self.nlines = nlines + self.dx = dx + self.dy = dy + self.thickness = thickness + + def paint(self, painter, option, widget=...): + line = QPainterPath() + line.moveTo(0, 0) + for i in range(self.nlines): + line.lineTo(self.dx if i % 2 == 0 else 0, self.dy * (i + 1)) + painter.strokePath(line, QPen(TEXT_COLOR, self.thickness)) + + def boundingRect(self): + return QRectF(-self.thickness, -self.thickness, self.dx + self.thickness*2, self.dy * self.nlines + self.thickness*2) + def arrowhead_vectors(start: QPointF, end: QPointF, stem_width: float, flare_width: float) -> Tuple[QPointF, QPointF, QPointF, QPointF, QPointF, QPointF]: direction = end - start norm = normalized(direction)