diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py index 79b7220..513c711 100644 --- a/typetapper/hierarchy_graph_view.py +++ b/typetapper/hierarchy_graph_view.py @@ -1,12 +1,15 @@ -from typing import Tuple, Optional, Union, Dict +from typing import Tuple, Optional, Union, Dict, List from collections import defaultdict from math import sqrt, sin, cos, pi +from dataclasses import dataclass from PySide6.QtCore import QRectF, QPointF from PySide6.QtGui import QColor, QFont, QBrush, QPainterPath, QPen -from PySide6.QtWidgets import QGraphicsItem, QGraphicsRectItem, QGraphicsSimpleTextItem, QHBoxLayout, QGraphicsScene +from PySide6.QtWidgets import QGraphicsItem, QGraphicsRectItem, QGraphicsSimpleTextItem, QHBoxLayout, QGraphicsScene, \ + QGraphicsEllipseItem import networkx +import pygraphviz from angrmanagement.config import Conf from angrmanagement.data.instance import Instance @@ -47,6 +50,11 @@ class HierarchicalGraphView(BaseView): layout.addWidget(self.graph) self.setLayout(layout) +@dataclass(frozen=True, eq=False) +class ExternNode: + goes_out: bool + goes_in: bool + class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): def __init__(self, parent, instance: Instance, hg: HierarchicalGraph, current_group: Optional[RelativeAtomGroup]=None): self.hg = hg @@ -75,12 +83,16 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): assert False def select(self, node: Optional[RelativeAtomOrGroup]): - if not self.current_node.am_none: + try: self.nodes[self.current_node.am_obj].update() + except KeyError: + pass self.current_node.am_obj = node self.current_node.am_event() - if not self.current_node.am_none: + try: self.nodes[self.current_node.am_obj].update() + except KeyError: + pass def _layout(self): for item in list(self.scene().items()): @@ -90,14 +102,25 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): word_height = 40 byte_height = word_height / self.hg.kp.kb._project.arch.bytes - ng = self.hg.local_graph(self.current_group.am_obj) + current_group = self.current_group.am_obj + ng = self.hg.local_graph(current_group) ug = networkx.Graph(ng) for edge in ug.edges: del ug.edges[edge]['prev'] del ug.edges[edge]['next'] + enodes = [] + if current_group in ug.nodes: + for adj in list(ug.adj[current_group]): + enode = ExternNode(goes_out=current_group in ng.succ[adj], goes_in=current_group in ng.pred[adj]) + enodes.append(enode) + ug.remove_edge(current_group, adj) + ug.add_edge(adj, enode) + assert len(ug.adj[current_group]) == 0 + ug.remove_node(current_group) + lookup = list(ug.nodes()) reverse = {x: i for i, x in enumerate(lookup)} - charts = [PropChartHG( + qnodes: List[PropNode] = [PropExtern(None) if isinstance(node, ExternNode) else PropChartHG( None, self.hg.prop(node), byte_height=byte_height, @@ -109,29 +132,35 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): ag = networkx.drawing.nx_agraph.to_agraph(ug) for node in ag.nodes_iter(): idx = int(node) - node.attr['width'] = charts[idx].width / 72 - node.attr['height'] = charts[idx].height / 72 - node.attr['shape'] = 'box' - for edge in ag.edges_iter(): - edge.attr['len'] = 50 / 72 + 200 / 72 + qnodes[idx].set_agraph_node_properties(node) + ag.edge_attr['len'] = 50 / 72 + 200 / 72 ag.graph_attr['overlap'] = 'false' ag._layout() #ag.draw('/home/audrey/render.dot', prog='nop') for node in ag.nodes_iter(): idx = int(node) - chart = charts[idx] + chart = qnodes[idx] posx, posy = node.attr['pos'].split(',') - chart.setPos(QPointF(float(posx) - chart.width / 2, -float(posy) - chart.height / 2)) + chart.center = QPointF(float(posx), -float(posy)) self.scene().addItem(chart) for u, v in ug.edges: - towards_end = lookup[v] in ng.succ[lookup[u]] - towards_start = lookup[u] in ng.succ[lookup[v]] - arrow = PropArrow(None, charts[u], charts[v], towards_start, towards_end) + lu = lookup[u] + lv = lookup[v] + if isinstance(lu, ExternNode): + towards_start = lu.goes_out + towards_end = lu.goes_in + elif isinstance(lv, ExternNode): + towards_start = lv.goes_in + towards_end = lv.goes_out + else: + towards_end = lv in ng.succ[lu] + towards_start = lu in ng.succ[lv] + arrow = PropArrow(None, qnodes[u], qnodes[v], towards_start, towards_end) self.scene().addItem(arrow) - self.nodes = {chart.nav: chart for chart in charts} + self.nodes = {chart.nav: chart for chart in qnodes if isinstance(chart, PropChartHG)} rect = self.scene().itemsBoundingRect() self.resetTransform() @@ -166,8 +195,23 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): event.accept() self.select(None) +class PropNode(QGraphicsItem): + @property + def center(self): + raise NotImplementedError -class PropChart(QGraphicsItem): + @center.setter + def center(self, v): + raise NotImplementedError + + def clip_edge(self, extent: QPointF) -> QPointF: + raise NotImplementedError + + def set_agraph_node_properties(self, node: pygraphviz.Node): + raise NotImplementedError + + +class PropChart(PropNode): def __init__(self, parent, prop: Prop, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5.): self.prop = prop self.max_width = max_width @@ -186,6 +230,38 @@ class PropChart(QGraphicsItem): self._layout_marks() + def set_agraph_node_properties(self, node: pygraphviz.Node): + node.attr['width'] = self.width / 72 + node.attr['height'] = self.height / 72 + node.attr['shape'] = 'box' + + def clip_edge(self, extent: QPointF) -> QPointF: + center_start = self.center + center_end = extent + delta_x = center_end.x() - center_start.x() + delta_y = center_end.y() - center_start.y() + slope = delta_y / delta_x if delta_x else float('inf') + + start_corner_slope = self.height / self.width + start_horiz = start_corner_slope > slope > -start_corner_slope + + if start_horiz: + start_right = center_end.x() > center_start.x() + if start_right: + start_ex = self.width / 2 + else: + start_ex = -self.width / 2 + start_ey = start_ex * slope + else: + start_bottom = center_end.y() > center_start.y() + if start_bottom: + start_ey = self.height / 2 + else: + start_ey = -self.height / 2 + start_ex = start_ey / slope + + return QPointF(center_start.x() + start_ex, center_start.y() + start_ey) + def boundingRect(self): return QRectF(0, 0, self.width, self.height) @@ -199,6 +275,11 @@ class PropChart(QGraphicsItem): pos.setY(pos.y() + self.height / 2) return pos + @center.setter + def center(self, v: QPointF): + self.setX(v.x() - self.width / 2) + self.setY(v.y() - self.height / 2) + def _layout_marks(self): 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()] data.sort() @@ -293,9 +374,55 @@ class PropChartHG(PropChart): super().paint(painter, option, widget) +class PropExtern(PropNode): + def __init__(self, parent, radius=28.): + self.radius = radius + self.objects = [] + self.object_bg: Optional[QGraphicsEllipseItem] = None + self.object_text = None + super().__init__(parent) + + self._layout() + + def _layout(self): + self.object_bg = QGraphicsEllipseItem(QRectF(-self.radius, -self.radius, self.radius*2, self.radius*2), self) + self.object_bg.setBrush(BOX_COLOR) + self.object_bg.setPen(QPen(BOX_BORDER_COLOR, BOX_BORDER_WIDTH)) + self.objects.append(self.object_bg) + + self.object_text = QGraphicsSimpleTextItem("Out", self) + self.object_text.setPos(-self.object_text.boundingRect().center()) + self.object_text.setFont(TEXT_FONT) + self.objects.append(self.object_text) + + @property + def center(self): + return self.pos() + + @center.setter + def center(self, v: QPointF): + self.setPos(v) + + def clip_edge(self, extent: QPointF) -> QPointF: + vector = extent - self.center + edge_vector = normalized(vector) * self.radius + return self.center + edge_vector + + def set_agraph_node_properties(self, node: pygraphviz.Node): + node.attr['width'] = self.radius * 2 / 72 + node.attr['shape'] = 'circle' + + def boundingRect(self) -> QRectF: + return QRectF(-self.radius, -self.radius, self.radius*2, self.radius*2) + + def shape(self) -> QPainterPath: + return self.object_bg.shape().translated(-self.radius, -self.radius) + + def paint(self, painter, option, widget=...): + pass class PropArrow(QGraphicsItem): - def __init__(self, parent, start: PropChart, end: PropChart, toward_start: bool, toward_end: bool, stroke=3., arrow_size=8.): + def __init__(self, parent, start: PropNode, end: PropNode, toward_start: bool, toward_end: bool, stroke=3., arrow_size=8.): self.start = start self.end = end self.toward_start = toward_start @@ -317,50 +444,8 @@ class PropArrow(QGraphicsItem): self.layout() def layout(self): - center_start = self.start.center - center_end = self.end.center - delta_x = center_end.x() - center_start.x() - delta_y = center_end.y() - center_start.y() - slope = delta_y / delta_x if delta_x else float('inf') - - start_corner_slope = self.start.height / self.start.width - end_corner_slope = self.end.height / self.end.width - - start_horiz = start_corner_slope > slope > -start_corner_slope - end_horiz = end_corner_slope > slope > -end_corner_slope - - if start_horiz: - start_right = center_end.x() > center_start.x() - if start_right: - start_ex = self.start.width / 2 - else: - start_ex = -self.start.width / 2 - start_ey = start_ex * slope - else: - start_bottom = center_end.y() > center_start.y() - if start_bottom: - start_ey = self.start.height / 2 - else: - start_ey = -self.start.height / 2 - start_ex = start_ey / slope - - if end_horiz: - end_right = center_start.x() > center_end.x() - if end_right: - end_ex = self.end.width / 2 - else: - end_ex = -self.end.width / 2 - end_ey = end_ex * slope - else: - end_bottom = center_start.y() > center_end.y() - if end_bottom: - end_ey = self.end.height / 2 - else: - end_ey = -self.end.height / 2 - end_ex = end_ey / slope - - start_pt = QPointF(center_start.x() + start_ex, center_start.y() + start_ey) - end_pt = QPointF(center_end.x() + end_ex, center_end.y() + end_ey) + start_pt = self.start.clip_edge(self.end.center) + end_pt = self.end.clip_edge(self.start.center) if self.percentage != 1.: end_pt = start_pt + (end_pt - start_pt) * self.percentage