From 7292619c1218597c2cc0f0c6218534a85a1ebf44 Mon Sep 17 00:00:00 2001
From: pancake <pancake@nopcode.org>
Date: Fri, 8 Mar 2019 10:07:42 +0100
Subject: [PATCH] Implement graph.trace and improve dt+ with dt++ for abt
 ##graph

---
 libr/core/cconfig.c   |  1 +
 libr/core/cmd_debug.c | 40 ++++++++++++++++++++++++++--------------
 libr/core/graph.c     | 34 +++++++++++++++++++++++++++++++---
 libr/core/visual.c    |  2 +-
 4 files changed, 59 insertions(+), 18 deletions(-)

diff --git a/libr/core/cconfig.c b/libr/core/cconfig.c
index d468ac4be5..6814827c7e 100644
--- a/libr/core/cconfig.c
+++ b/libr/core/cconfig.c
@@ -3167,6 +3167,7 @@ R_API int r_core_config_init(RCore *core) {
 	SETPREF ("tcp.islocal", "false", "Bind a loopback for tcp command server");
 
 	/* graph */
+	SETPREF ("graph.trace", "false", "Fold all non-traced basic blocks");
 	SETPREF ("graph.few", "false", "Show few basic blocks in the graph");
 	SETPREF ("graph.comments", "true", "Show disasm comments in graph");
 	SETPREF ("graph.cmtright", "false", "Show comments at right");
diff --git a/libr/core/cmd_debug.c b/libr/core/cmd_debug.c
index a9314c3595..9bd61f9a0e 100644
--- a/libr/core/cmd_debug.c
+++ b/libr/core/cmd_debug.c
@@ -4484,22 +4484,34 @@ static int cmd_debug(void *data, const char *input) {
 			core->dbg->trace = r_debug_trace_new ();
 			break;
 		case '+': // "dt+"
-			ptr = input + 3;
-			addr = r_num_math (core->num, ptr);
-			ptr = strchr (ptr, ' ');
-			if (ptr) {
-				RAnalOp *op = r_core_op_anal (core, addr);
-				if (op) {
-					RDebugTracepoint *tp = r_debug_trace_add (core->dbg, addr, op->size);
-					if (!tp) {
+			if (input[2] == '+') { // "dt++"
+				char *a, *s = r_str_trim (strdup (input + 3));
+				RList *args = r_str_split_list (s, " ");
+				RListIter *iter;
+				r_list_foreach (args, iter, a) {
+					ut64 addr = r_num_get (NULL, a);
+					(void)r_debug_trace_add (core->dbg, addr, 1);
+				}
+				r_list_free (args);
+				free (s);
+			} else {
+				ptr = input + 3;
+				addr = r_num_math (core->num, ptr);
+				ptr = strchr (ptr, ' ');
+				if (ptr) {
+					RAnalOp *op = r_core_op_anal (core, addr);
+					if (op) {
+						RDebugTracepoint *tp = r_debug_trace_add (core->dbg, addr, op->size);
+						if (!tp) {
+							r_anal_op_free (op);
+							break;
+						}
+						tp->count = r_num_math (core->num, ptr + 1);
+						r_anal_trace_bb (core->anal, addr);
 						r_anal_op_free (op);
-						break;
+					} else {
+						eprintf ("Cannot analyze opcode at 0x%08" PFMT64x "\n", addr);
 					}
-					tp->count = r_num_math (core->num, ptr + 1);
-					r_anal_trace_bb (core->anal, addr);
-					r_anal_op_free (op);
-				} else {
-					eprintf ("Cannot analyze opcode at 0x%08" PFMT64x "\n", addr);
 				}
 			}
 			break;
diff --git a/libr/core/graph.c b/libr/core/graph.c
index 07c0505be3..2cb493e01d 100644
--- a/libr/core/graph.c
+++ b/libr/core/graph.c
@@ -2190,6 +2190,26 @@ static void get_bbupdate(RAGraph *g, RCore *core, RAnalFunction *fcn) {
 	core->anal->stackptr = saved_stackptr;
 }
 
+static void fold_asm_trace(RCore *core, RAGraph *g) {
+	const RList *nodes = r_graph_get_nodes (g->graph);
+	RGraphNode *gn;
+	RListIter *it;
+	RANode *n;
+
+	graph_foreach_anode (nodes, it, gn, n) {
+		if (get_anode (g->curnode) == n) {
+			n->is_mini = false;
+			g->need_reload_nodes = true;
+			continue;
+		}
+		ut64 addr = r_num_get (NULL, n->title);
+		RDebugTracepoint *tp = r_debug_trace_get (core->dbg, addr);
+		n->is_mini = (tp == NULL);
+	}
+	g->need_update_dim = 1;
+	//agraph_refresh (r_cons_singleton ()->event_data);
+}
+
 static void delete_dup_edges (RAGraph *g) {
 	RListIter *it, *in_it, *in_it2, *in_it2_tmp;
 	RGraphNode *n, *a, *b;
@@ -3268,7 +3288,13 @@ static int check_changes(RAGraph *g, int is_interactive,
 		}
 	}
 	if (fcn) {
-		agraph_update_title (g, fcn);
+		agraph_update_title (core, g, fcn);
+	}
+	if (core && core->config) {
+		if (r_config_get_i (core->config, "graph.trace")) {
+			// fold all bbs not traced
+			fold_asm_trace (core, g);
+		}
 	}
 	if (g->need_update_dim || g->need_reload_nodes || !is_interactive) {
 		update_node_dimension (g->graph, is_mini (g), g->zoom, g->edgemode, g->is_callgraph, g->layout);
@@ -4227,8 +4253,8 @@ R_API int r_core_visual_graph(RCore *core, RAGraph *g, RAnalFunction *_fcn, int
 				" V            - toggle basicblock / call graphs\n"
 				" w            - toggle between movements speed 1 and graph.scroll\n"
 				" x/X          - jump to xref/ref\n"
-				" y            - toggle node folding/minification\n"
 				" Y            - toggle tiny graph\n"
+				" z            - toggle node folding\n"
 				" Z            - follow parent node");
 			r_cons_less ();
 			r_cons_any_key (NULL);
@@ -4316,7 +4342,9 @@ R_API int r_core_visual_graph(RCore *core, RAGraph *g, RAnalFunction *_fcn, int
 			r_core_cmd0 (core, "e!asm.hint.lea");
 			break;
 		case '$':
+			r_core_cmd (core, "dr PC=$$", 0);
 			r_core_cmd (core, "sr PC", 0);
+			g->need_reload_nodes = true;
 			break;
 		case 'R':
 			if (r_config_get_i (core->config, "scr.randpal")) {
@@ -4401,7 +4429,7 @@ R_API int r_core_visual_graph(RCore *core, RAGraph *g, RAnalFunction *_fcn, int
 			agraph_toggle_tiny (g);
 			agraph_update_seek (g, get_anode (g->curnode), true);
 			break;
-		case 'y':
+		case 'z':
 			agraph_toggle_mini (g);
 			break;
 		case 'v':
diff --git a/libr/core/visual.c b/libr/core/visual.c
index 4ceac38299..952bc4ba8d 100644
--- a/libr/core/visual.c
+++ b/libr/core/visual.c
@@ -289,7 +289,7 @@ static const char *help_msg_visual[] = {
 	"xX", "show xrefs/refs of current function from/to data/code",
 	"yY", "copy and paste selection",
 	"z", "fold/unfold comments in disassembly",
-	"Z", "toggle zoom mode",
+	"Z", "shift-tab rotate print modes", // ctoggle zoom mode",
 	"Enter", "follow address of jump/call",
 	NULL
 };