From a81352a846144e528e01f8f65ed3d5e871720c12 Mon Sep 17 00:00:00 2001
From: Clyve Gassant <cgassant@sas.upenn.edu>
Date: Tue, 19 Jul 2022 15:17:28 +0000
Subject: [PATCH] demangle ocaml symbols function

Signed-off-by: Clyve Gassant <cgassant@sas.upenn.edu>
---
 core/demangle_ocaml_symbols.ml  | 48 +++++++++++++++++++++++++++++++++
 core/demangle_ocaml_symbols.mli | 11 ++++++++
 core/dune                       |  2 +-
 core/symbol_selection.ml        |  8 +++++-
 test/demangle_ocaml_symbols.ml  | 27 +++++++++++++++++++
 5 files changed, 94 insertions(+), 2 deletions(-)
 create mode 100644 core/demangle_ocaml_symbols.ml
 create mode 100644 core/demangle_ocaml_symbols.mli
 create mode 100644 test/demangle_ocaml_symbols.ml

diff --git a/core/demangle_ocaml_symbols.ml b/core/demangle_ocaml_symbols.ml
new file mode 100644
index 000000000..ab36aa55e
--- /dev/null
+++ b/core/demangle_ocaml_symbols.ml
@@ -0,0 +1,48 @@
+open! Core
+
+let decode_two_digit_hexadecimal_number first_character second_character =
+  let%bind.Option hex_digit_of_first_character = Char.get_hex_digit first_character in
+  let%bind.Option hex_digit_of_second_character = Char.get_hex_digit second_character in
+  let leftshift_first_character = hex_digit_of_first_character lsl 4 in
+  let bit_or_on_the_hexadecimals =
+    leftshift_first_character lor hex_digit_of_second_character
+  in
+  Char.of_int bit_or_on_the_hexadecimals
+;;
+
+let parser =
+  let open Angstrom in
+  let strip_numeric_suffix =
+    char '_' *> skip_while Char.is_digit *> end_of_input *> return None
+  in
+  let double_underscores = string "__" >>| fun _ -> Some '.' in
+  let two_digit_hexadecimal_number =
+    let hex_character = satisfy Char.is_hex_digit in
+    let hexcode =
+      decode_two_digit_hexadecimal_number <$> char '$' *> hex_character <*> hex_character
+    in
+    hexcode
+    >>= fun integer ->
+    match integer with
+    | None -> fail "invalid integer"
+    | Some character -> return (Some character)
+  in
+  let normal_character = any_char >>| fun character -> Some character in
+  let token =
+    choice
+      ~failure_msg:"unrecognized token"
+      [ strip_numeric_suffix
+      ; double_underscores
+      ; two_digit_hexadecimal_number
+      ; normal_character
+      ]
+  in
+  string "caml" *> many1 token
+;;
+
+let demangle mangled_symbol =
+  let mangled_string = Angstrom.parse_string ~consume:All parser mangled_symbol in
+  match mangled_string with
+  | Ok list -> Some (String.of_char_list (List.filter_map list ~f:Fn.id))
+  | Error _ -> None
+;;
diff --git a/core/demangle_ocaml_symbols.mli b/core/demangle_ocaml_symbols.mli
new file mode 100644
index 000000000..27a189389
--- /dev/null
+++ b/core/demangle_ocaml_symbols.mli
@@ -0,0 +1,11 @@
+open! Core
+
+(** The logic for this function is derived from perf [0].
+    This function is used to change the symbols in the application 
+    executable from a mangled form to a demangled form in ocaml. Now
+    when running [magic-trace run -trigger] the symbols will appear in
+    their demangled form. Will return None if the symbol is not
+    recognized as an OCaml symbol.
+   
+   [0]: https://github.com/torvalds/linux/blob/5bfc75d92efd494db37f5c4c173d3639d4772966/tools/perf/util/demangle-ocaml.c *)
+val demangle : string -> string option
diff --git a/core/dune b/core/dune
index 35855af51..5543ef5bb 100644
--- a/core/dune
+++ b/core/dune
@@ -2,7 +2,7 @@
  (name magic_trace_core)
  (public_name magic-trace.magic_trace_core)
  (libraries core async core_unix.filename_unix expect_test_helpers_core fzf
-   magic_trace owee re)
+   magic_trace owee re angstrom)
  (inline_tests)
  (preprocess
   (pps ppx_jane)))
diff --git a/core/symbol_selection.ml b/core/symbol_selection.ml
index 7dd1affba..45eb9bd30 100644
--- a/core/symbol_selection.ml
+++ b/core/symbol_selection.ml
@@ -18,7 +18,13 @@ let select_owee_symbol ~elf ~header select =
   let open Deferred.Or_error.Let_syntax in
   let all_symbols = Elf.all_symbols ~select elf in
   let all_symbol_names = List.map all_symbols ~f:Tuple2.get1 in
-  match%bind Fzf.pick_one ~header (Inputs all_symbol_names) with
+  let demangled_symbols =
+    List.map all_symbol_names ~f:(fun mangled_symbol ->
+      match Demangle_ocaml_symbols.demangle mangled_symbol with
+      | None -> mangled_symbol, mangled_symbol
+      | Some demangled_symbol -> demangled_symbol, mangled_symbol)
+  in
+  match%bind Fzf.pick_one ~header (Assoc demangled_symbols) with
   | None -> Deferred.Or_error.error_string "No symbol selected"
   | Some chosen_name ->
     let chosen_symbol =
diff --git a/test/demangle_ocaml_symbols.ml b/test/demangle_ocaml_symbols.ml
new file mode 100644
index 000000000..f343b4b70
--- /dev/null
+++ b/test/demangle_ocaml_symbols.ml
@@ -0,0 +1,27 @@
+open! Core
+open Magic_trace_core
+
+let demangle_symbol_test symbol =
+  let demangle_symbol = Demangle_ocaml_symbols.demangle symbol in
+  print_s [%sexp (demangle_symbol : string option)]
+;;
+
+let%expect_test "real mangled symbol" =
+  demangle_symbol_test "camlAsync_unix__Unix_syscalls__to_string_57255";
+  [%expect {| (Async_unix.Unix_syscalls.to_string) |}]
+;;
+
+let%expect_test "proper hexcode" =
+  demangle_symbol_test "caml$3f";
+  [%expect {| (?) |}]
+;;
+
+let%expect_test "improper hexcode" =
+  demangle_symbol_test "caml$7l";
+  [%expect {| ($7l) |}]
+;;
+
+let%expect_test "when the symbol is not a demangled ocaml symbol" =
+  demangle_symbol_test "dr__$3e$21_358";
+  [%expect {| () |}]
+;;