From 52f3b66575578ba118323626bb69f53af2029d8a Mon Sep 17 00:00:00 2001
From: degrigis <degrigis@ucsb.edu>
Date: Wed, 31 Jan 2024 18:14:46 -0800
Subject: [PATCH] better printing in simgrviz

Signed-off-by: degrigis <degrigis@ucsb.edu>
---
 greed/exploration_techniques/simgrviz.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/greed/exploration_techniques/simgrviz.py b/greed/exploration_techniques/simgrviz.py
index 818a9ccd..471b413a 100644
--- a/greed/exploration_techniques/simgrviz.py
+++ b/greed/exploration_techniques/simgrviz.py
@@ -63,8 +63,10 @@ def _add_node(self, state):
         self._simgGraph.add_node(state_id)
         self._simgGraph.nodes[state_id]['timestamp'] = str(self.timestamp)
         self._simgGraph.nodes[state_id]['pc'] = state.pc
-        # if state.curr_stmt.__internal_name__ == "LOG1":
-        #    self._simgGraph.nodes[state_id]['log'] = state.curr_stmt.arg2_val
+        if state.curr_stmt.__internal_name__ == "STOP" or state.curr_stmt.__internal_name__ == "RETURN":
+            self._simgGraph.nodes[state_id]['done'] = state.curr_stmt
+        if state.curr_stmt.__internal_name__ == "REVERT":
+            self._simgGraph.nodes[state_id]['fail'] = state.curr_stmt
         # self._simgGraph.nodes[state_id]['csts'] = '\n'.join([str(x.dump()) for x in state.constraints])
         self.timestamp += 1
         return state_id
@@ -88,12 +90,20 @@ def _dump_graph(self):
         for node_id in self._simgGraph.nodes:
             node = self._simgGraph.nodes[node_id]
 
-            shape = 'box'
+            shape = '<SHAPE>'
             s += '\t\"{}\" [shape={},label='.format(node_id[:10], shape)
             s += '<ts:{}<br align="left"/>'.format(node["timestamp"])
             s += '<br align="left"/>pc:{}'.format(node["pc"])
-            if node.get("log", None):
-                s += '<br align="left"/>log:{}'.format(node["log"])
+            
+            if node.get("done", None):
+                s += '<br align="left"/>{}'.format(node["done"])
+                s = s.replace("<SHAPE>", "star")
+            elif node.get("fail", None):
+                s += '<br align="left"/>{}'.format(node["fail"])
+                s = s.replace("<SHAPE>", "invtriangle")
+            else:
+                s = s.replace("<SHAPE>", "box")
+                #import ipdb; ipdb.set_trace()
             # s += '<br align="left"/>csts:{}'.format(node["csts"])
             s += '<br align="left"/>>];\n'