From 152a97786568c09917641ef7a631714a059f46ef Mon Sep 17 00:00:00 2001
From: Felipe Pena <felipensp@gmail.com>
Date: Sat, 17 Feb 2024 05:51:44 -0300
Subject: [PATCH] v.debug: implement tracing hook feature (#20818)

---
 vlib/v/checker/fn.v                   |  5 +-
 vlib/v/debug/tests/trace/trace_test.v | 65 +++++++++++++++++++++++++
 vlib/v/debug/trace.v                  | 68 +++++++++++++++++++++++++++
 vlib/v/debug/tracing_test.v           | 61 ++++++++++++++++++++++++
 vlib/v/gen/c/fn.v                     | 41 +++++++++++++---
 vlib/v/pref/pref.c.v                  |  4 ++
 6 files changed, 235 insertions(+), 9 deletions(-)
 create mode 100644 vlib/v/debug/tests/trace/trace_test.v
 create mode 100644 vlib/v/debug/trace.v
 create mode 100644 vlib/v/debug/tracing_test.v

diff --git a/vlib/v/checker/fn.v b/vlib/v/checker/fn.v
index 1a9db9f62..10fa1eb59 100644
--- a/vlib/v/checker/fn.v
+++ b/vlib/v/checker/fn.v
@@ -1524,8 +1524,9 @@ fn (mut c Checker) fn_call(mut node ast.CallExpr, mut continue_check &bool) ast.
 
 // register_trace_call registers the wrapper funcs for calling funcs for callstack feature
 fn (mut c Checker) register_trace_call(node ast.CallExpr, func ast.Fn) {
-	is_traceable := c.pref.is_callstack && c.table.cur_fn != unsafe { nil } && c.pref.is_callstack
-		&& c.file.imports.any(it.mod == 'v.debug') && node.name != 'v.debug.callstack'
+	is_traceable := (c.pref.is_callstack || c.pref.is_trace) && c.table.cur_fn != unsafe { nil }
+		&& node.language == .v && c.file.imports.any(it.mod == 'v.debug')
+		&& node.name !in ['v.debug.callstack', 'v.debug.add_after_call', 'v.debug.add_before_call', 'v.debug.remove_after_call', 'v.debug.remove_before_call']
 	if is_traceable {
 		hash_fn, fn_name := c.table.get_trace_fn_name(c.table.cur_fn, node)
 		calling_fn := if func.is_method {
diff --git a/vlib/v/debug/tests/trace/trace_test.v b/vlib/v/debug/tests/trace/trace_test.v
new file mode 100644
index 000000000..879de9e0c
--- /dev/null
+++ b/vlib/v/debug/tests/trace/trace_test.v
@@ -0,0 +1,65 @@
+@[has_globals]
+module main
+
+import time
+import v.debug
+
+__global counter = HookCounter{}
+
+struct HookCounter {
+pub mut:
+	after  int
+	before int
+}
+
+@[heap]
+struct Profiler {
+mut:
+	fn_map map[string]time.StopWatch
+}
+
+fn (mut p Profiler) trace_after(fn_name string) {
+	counter.after++
+	println('>>> called after ${fn_name} - ${p.fn_map[fn_name].elapsed()}')
+}
+
+fn (mut p Profiler) trace_before(fn_name string) {
+	p.fn_map[fn_name] = time.new_stopwatch()
+	println('>>> called before ${fn_name}')
+	counter.before++
+}
+
+fn test_call() {
+	println('${@FN} got called')
+	time.sleep(100)
+}
+
+fn hook_before(fn_name string) {
+	counter.before++
+}
+
+fn hook_after(fn_name string) {
+	counter.after++
+}
+
+fn test_main() {
+	counter = HookCounter{}
+	mut myprofiler := Profiler{}
+	hook1 := debug.add_after_call(myprofiler.trace_after)
+	hook2 := debug.add_before_call(myprofiler.trace_before)
+	hook3 := debug.add_before_call(hook_before)
+	hook4 := debug.add_after_call(hook_after)
+	test_call()
+	debug.remove_after_call(hook1)
+	debug.remove_before_call(hook2)
+	debug.remove_before_call(hook3)
+	debug.remove_after_call(hook4)
+	test_call()
+	$if trace ? {
+		assert counter.before == 4
+		assert counter.after == 4
+	} $else {
+		assert counter.before == 0
+		assert counter.after == 0
+	}
+}
diff --git a/vlib/v/debug/trace.v b/vlib/v/debug/trace.v
new file mode 100644
index 000000000..d323280c7
--- /dev/null
+++ b/vlib/v/debug/trace.v
@@ -0,0 +1,68 @@
+@[has_globals]
+module debug
+
+@[markused]
+__global g_trace = TraceHook{}
+
+type HookFnCall = fn (string)
+
+@[noinit]
+pub struct TraceHook {
+mut:
+	in_hook bool
+pub mut:
+	trace_after_call  []HookFnCall
+	trace_before_call []HookFnCall
+}
+
+// after_call_hook calls the registered hook fns
+@[markused]
+fn after_call_hook(fn_name string) {
+	g_trace.in_hook = true
+	for func in g_trace.trace_after_call {
+		func(fn_name)
+	}
+	g_trace.in_hook = false
+}
+
+// before_call_hook calls the registered hook fns
+@[markused]
+fn before_call_hook(fn_name string) {
+	g_trace.in_hook = true
+	for func in g_trace.trace_before_call {
+		func(fn_name)
+	}
+	g_trace.in_hook = false
+}
+
+// add_after_call adds a fn hook to after hook list and returns its id
+@[inline; markused]
+pub fn add_after_call(func HookFnCall) HookFnCall {
+	g_trace.trace_after_call << func
+	return func
+}
+
+// add_before_call adds a fn hook to before hook list and return its id
+@[inline; markused]
+pub fn add_before_call(func HookFnCall) HookFnCall {
+	g_trace.trace_before_call << func
+	return func
+}
+
+// remove_after_call removes a fn hook from after hook list by its idx
+@[inline; markused]
+pub fn remove_after_call(func HookFnCall) {
+	idx := g_trace.trace_after_call.index(func)
+	if idx != -1 {
+		g_trace.trace_after_call.delete(idx)
+	}
+}
+
+// remove_before_call removes a fn hook from before hook list by its idx
+@[inline; markused]
+pub fn remove_before_call(func HookFnCall) {
+	idx := g_trace.trace_before_call.index(func)
+	if idx != -1 {
+		g_trace.trace_before_call.delete(idx)
+	}
+}
diff --git a/vlib/v/debug/tracing_test.v b/vlib/v/debug/tracing_test.v
new file mode 100644
index 000000000..96e574cdd
--- /dev/null
+++ b/vlib/v/debug/tracing_test.v
@@ -0,0 +1,61 @@
+import os
+import term
+import time
+
+const vexe = @VEXE
+const trace_tests_path = os.join_path(@VEXEROOT, 'vlib', 'v', 'debug', 'tests', 'trace')
+const bar = term.yellow('-'.repeat(105))
+const be_verbose = os.getenv('GITHUB_JOB') != '' || os.getenv('VERBOSE') != ''
+
+const expect_exe = os.quoted_path(os.find_abs_path_of_executable('expect') or {
+	eprintln('skipping test, since expect is missing')
+	exit(0)
+})
+
+fn testsuite_begin() {
+	os.chdir(@VEXEROOT) or {}
+}
+
+fn gprintln(msg string) {
+	println(term.green(msg))
+	flush_stdout()
+}
+
+fn gprint(msg string) {
+	print(term.green(msg))
+	flush_stdout()
+}
+
+fn test_trace() {
+	all_expect_files := os.walk_ext(trace_tests_path, '.v')
+	assert all_expect_files.len > 0, 'no .v files found in ${trace_tests_path}'
+	mut oks := 0
+	for eidx, efile in all_expect_files.sorted() {
+		// if !efile.contains('sumtype') { gprintln('skipping $efile') continue }
+		vfile := efile
+
+		if be_verbose {
+			println(bar)
+		}
+		gprint('>>>> Running [${eidx + 1}/${all_expect_files.len}] ${term.magenta(efile):-68} ... ')
+		if be_verbose {
+			println('')
+		}
+
+		compile_sw := time.new_stopwatch()
+		comp_res := os.system('${os.quoted_path(vexe)} -d trace test ${os.quoted_path(vfile)}')
+		cdur_ms := compile_sw.elapsed().milliseconds()
+		if be_verbose {
+			gprintln('>>>>>>>>>>> compilation took ${cdur_ms} ms, comp_res: ${comp_res}')
+		}
+		if comp_res != 0 {
+			assert false, term.red('failed test cmd: ${vfile}')
+		}
+		assert true
+		oks++
+	}
+	os.chdir(@VEXEROOT) or {}
+
+	println(bar)
+	gprintln('Passed debugger tests: ${term.bold(oks.str())} of ${all_expect_files.len} total.')
+}
diff --git a/vlib/v/gen/c/fn.v b/vlib/v/gen/c/fn.v
index debf1bb21..c82a3083c 100644
--- a/vlib/v/gen/c/fn.v
+++ b/vlib/v/gen/c/fn.v
@@ -297,18 +297,45 @@ fn (mut g Gen) gen_fn_decl(node &ast.FnDecl, skip bool) {
 			g.definitions.write_string(');\n')
 
 			orig_fn_args := call_fn.func.params.map(it.name).join(', ')
-
-			if g.cur_fn.is_method || g.cur_fn.is_static_type_method {
-				g.writeln('\tarray_push((array*)&g_callstack, _MOV((v__debug__FnTrace[]){ ((v__debug__FnTrace){.name = _SLIT("${g.table.type_to_str(g.cur_fn.receiver.typ)}.${g.cur_fn.name.all_after_last('__static__')}"),.file = _SLIT("${call_fn.file}"),.line = ${call_fn.line},}) }));')
-			} else {
-				g.writeln('\tarray_push((array*)&g_callstack, _MOV((v__debug__FnTrace[]){ ((v__debug__FnTrace){.name = _SLIT("${g.cur_fn.name}"),.file = _SLIT("${call_fn.file}"),.line = ${call_fn.line},}) }));')
+			add_trace_hook := g.pref.is_trace
+				&& call_fn.name !in ['v.debug.add_after_call', 'v.debug.add_before_call', 'v.debug.remove_after_call', 'v.debug.remove_before_call']
+			if g.pref.is_callstack {
+				if g.cur_fn.is_method || g.cur_fn.is_static_type_method {
+					g.writeln('\tarray_push((array*)&g_callstack, _MOV((v__debug__FnTrace[]){ ((v__debug__FnTrace){.name = _SLIT("${g.table.type_to_str(g.cur_fn.receiver.typ)}.${g.cur_fn.name.all_after_last('__static__')}"),.file = _SLIT("${call_fn.file}"),.line = ${call_fn.line},}) }));')
+				} else {
+					g.writeln('\tarray_push((array*)&g_callstack, _MOV((v__debug__FnTrace[]){ ((v__debug__FnTrace){.name = _SLIT("${g.cur_fn.name}"),.file = _SLIT("${call_fn.file}"),.line = ${call_fn.line},}) }));')
+				}
 			}
 			if call_fn.return_type == 0 || call_fn.return_type == ast.void_type {
+				if add_trace_hook {
+					g.writeln('\tif (!g_trace.in_hook) {')
+					g.writeln('\t\tv__debug__before_call_hook(_SLIT("${call_fn.name}"));')
+					g.writeln('\t}')
+				}
 				g.writeln('\t${c_name(call_fn.name)}(${orig_fn_args});')
-				g.writeln('\tarray_pop((array*)&g_callstack);')
+				if add_trace_hook {
+					g.writeln('\tif (!g_trace.in_hook) {')
+					g.writeln('\t\tv__debug__after_call_hook(_SLIT("${call_fn.name}"));')
+					g.writeln('\t}')
+				}
+				if g.pref.is_callstack {
+					g.writeln('\tarray_pop((array*)&g_callstack);')
+				}
 			} else {
+				if add_trace_hook {
+					g.writeln('\tif (!g_trace.in_hook) {')
+					g.writeln('\t\tv__debug__before_call_hook(_SLIT("${call_fn.name}"));')
+					g.writeln('\t}')
+				}
 				g.writeln('\t${g.typ(call_fn.return_type)} ret = ${c_name(call_fn.name)}(${orig_fn_args});')
-				g.writeln('\tarray_pop((array*)&g_callstack);')
+				if g.pref.is_callstack {
+					g.writeln('\tarray_pop((array*)&g_callstack);')
+				}
+				if add_trace_hook {
+					g.writeln('\tif (!g_trace.in_hook) {')
+					g.writeln('\t\tv__debug__after_call_hook(_SLIT("${call_fn.name}"));')
+					g.writeln('\t}')
+				}
 				g.writeln('\treturn ret;')
 			}
 			g.writeln('}')
diff --git a/vlib/v/pref/pref.c.v b/vlib/v/pref/pref.c.v
index 980376569..dffdd48eb 100644
--- a/vlib/v/pref/pref.c.v
+++ b/vlib/v/pref/pref.c.v
@@ -134,6 +134,7 @@ pub mut:
 	is_quiet           bool     // do not show the repetitive explanatory messages like the one for `v -prod run file.v` .
 	is_cstrict         bool     // turn on more C warnings; slightly slower
 	is_callstack       bool     // turn on callstack registers on each call when v.debug is imported
+	is_trace           bool     // turn on possibility to trace fn call where v.debug is imported
 	eval_argument      string   // `println(2+2)` on `v -e "println(2+2)"`. Note that this source code, will be evaluated in vsh mode, so 'v -e 'println(ls(".")!)' is valid.
 	test_runner        string   // can be 'simple' (fastest, but much less detailed), 'tap', 'normal'
 	profile_file       string   // the profile results will be stored inside profile_file
@@ -1051,6 +1052,9 @@ pub fn parse_args_and_show_errors(known_external_commands []string, args []strin
 	if 'callstack' in res.compile_defines_all {
 		res.is_callstack = true
 	}
+	if 'trace' in res.compile_defines_all {
+		res.is_trace = true
+	}
 	// keep only the unique res.build_options:
 	mut m := map[string]string{}
 	for x in res.build_options {
-- 
2.39.5