Added ability to trigger on addresses and line numbers

This feature adds triggers of the format `addr:0xfffff` or `line:file.ml:99:15`.
All triggers can be prefixed with either `symbol:`, `addr:`, or `line:`, but if
none is given, it will default to a symbol if possible and otherwise parse it as
a line number or address.

There is no feature to allow our fuzzy finding to work with the new options, but
it could be possible to implement fuzzy finding on files for example.

Signed-off-by: Aaron Lamoreaux <alamoreaux@janestreet.com>

core/elf.ml

@@ -136,6 +136,25 @@ let matching_functions t symbol_re =
   !res
 ;;
 
+let traverse_debug_line ~f t =
+  Option.iter t.debug ~f:(fun body ->
+      let cursor = Owee_buf.cursor body in
+      let pointers_to_other_sections =
+        Owee_elf.debug_line_pointers t.all_elf t.sections
+      in
+      let rec load_table_next () =
+        match Owee_debug_line.read_chunk cursor ~pointers_to_other_sections with
+        | None -> ()
+        | Some (header, chunk) ->
+          let process header (state : Owee_debug_line.state) () =
+            if not state.end_sequence then f header state
+          in
+          Owee_debug_line.fold_rows (header, chunk) process ();
+          load_table_next ()
+      in
+      load_table_next ())
+;;
+
 let find_symbol t name =
   let some_name = Some name in
   with_return (fun return ->
@@ -148,6 +167,72 @@ let find_symbol t name =
       None)
 ;;
 
+let find_selection t name : Selection.t option =
+  let find_line_selection name =
+    let split_name = String.split name ~on:':' in
+    let desired_filename = List.nth split_name 0 in
+    let desired_line =
+      Option.try_with (fun () -> List.nth_exn split_name 1 |> Int.of_string)
+    in
+    let desired_col =
+      Option.try_with (fun () -> List.nth_exn split_name 2 |> Int.of_string)
+    in
+    let cols = ref [] in
+    traverse_debug_line
+      ~f:(fun header state ->
+        let filename = Owee_debug_line.get_filename header state in
+        let line = state.line in
+        let col = state.col in
+        match desired_filename, desired_line, filename with
+        | Some desired_filename, Some desired_line, Some filename ->
+          if String.(desired_filename = filename) && desired_line = line
+          then cols := (col, state.address) :: !cols
+        | None, _, _ | _, None, _ | _, _, None -> ())
+      t;
+    let cols =
+      List.sort !cols ~compare:(fun (col1, _) (col2, _) -> Int.compare col1 col2)
+    in
+    match cols with
+    | [] -> None
+    | (col, address) :: _ ->
+      (match desired_col with
+      | None ->
+        if List.length cols > 1
+        then
+          Core.eprint_s
+            [%message
+              "Multiple snapshot symbols with same column. Selecting one."
+                ~columns:(List.map ~f:Tuple2.get1 cols : int list)
+                ~selected:(col : int)];
+        Some (Selection.Address { address; name })
+      | Some desired_col ->
+        (match
+           List.find_map cols ~f:(fun (col, _) ->
+               if col = desired_col then Some address else None)
+         with
+        | None -> None
+        | Some address -> Some (Selection.Address { address; name })))
+  in
+  let find_addr_selection name =
+    Option.bind
+      (Option.try_with (fun () -> Int.Hex.of_string name))
+      ~f:(fun address -> Some (Selection.Address { address; name }))
+  in
+  let find_symbol_selection name =
+    Option.map (find_symbol t name) ~f:(fun symbol -> Selection.Symbol symbol)
+  in
+  if String.is_prefix name ~prefix:"symbol:"
+  then find_symbol_selection (String.drop_prefix name 7)
+  else if String.is_prefix name ~prefix:"line:"
+  then find_line_selection (String.drop_prefix name 5)
+  else if String.is_prefix name ~prefix:"addr:"
+  then find_addr_selection (String.drop_prefix name 5)
+  else
+    Option.first_some
+      (Option.first_some (find_symbol_selection name) (find_line_selection name))
+      (find_addr_selection name)
+;;
+
 let all_symbols t =
   let res = String.Table.create () in
   Owee_elf.Symbol_table.iter t.symbol ~f:(fun symbol ->
@@ -163,12 +248,9 @@ let all_symbols t =
   String.Table.to_alist res
 ;;
 
-let symbol_stop_info t pid symbol =
-  let name = Owee_elf.Symbol_table.Symbol.name symbol t.string in
-  let name = Option.value_exn ~message:"stop_info symbols must have a name" name in
+let selection_stop_info t pid selection =
   let filename = Filename_unix.realpath t.filename in
-  let addr = Owee_elf.Symbol_table.Symbol.value symbol in
-  let addr =
+  let compute_addr addr =
     if t.statically_mappable
     then addr
     else
@@ -181,49 +263,46 @@ let symbol_stop_info t pid symbol =
              else None)
       |> List.hd_exn
   in
-  let size = Owee_elf.Symbol_table.Symbol.size_in_bytes symbol in
-  let offset = Int64.( - ) addr (Int64.of_int t.base_offset) in
-  let filter = [%string {|stop %{offset#Int64}/%{size#Int64}@%{filename}|}] in
-  { Stop_info.name; addr; filter }
+  let compute_filter ~name ~addr ~size =
+    let offset = Int64.( - ) addr (Int64.of_int t.base_offset) in
+    let filter = [%string {|stop %{offset#Int64}/%{size#Int64}@%{filename}|}] in
+    { Stop_info.name; addr; filter }
+  in
+  match selection with
+  | Selection.Symbol symbol ->
+    let name = Owee_elf.Symbol_table.Symbol.name symbol t.string in
+    let name = Option.value_exn ~message:"stop_info symbols must have a name" name in
+    let addr = Owee_elf.Symbol_table.Symbol.value symbol in
+    let addr = compute_addr addr in
+    let size = Owee_elf.Symbol_table.Symbol.size_in_bytes symbol in
+    compute_filter ~name ~addr ~size
+  | Address { address; name } ->
+    let addr = compute_addr (Int64.of_int address) in
+    compute_filter ~name ~addr ~size:1L
 ;;
 
 let addr_table t =
   let table = Int.Table.create () in
-  Option.iter t.debug ~f:(fun body ->
-      (* We only want to include line info from the start address of symbols in the table,
-       lest it grow too large on big executables. We don't need to include mappings for
-       lines in the middle of functions. *)
-      let symbol_starts = Int.Hash_set.create () in
-      Owee_elf.Symbol_table.iter t.symbol ~f:(fun symbol ->
-          if is_func symbol
-          then
-            Hash_set.add
-              symbol_starts
-              (Owee_elf.Symbol_table.Symbol.value symbol |> Int64.to_int_exn));
-      let cursor = Owee_buf.cursor body in
-      let pointers_to_other_sections =
-        Owee_elf.debug_line_pointers t.all_elf t.sections
-      in
-      let rec load_table_next () =
-        match Owee_debug_line.read_chunk cursor ~pointers_to_other_sections with
-        | None -> ()
-        | Some (header, chunk) ->
-          let process header (state : Owee_debug_line.state) () =
-            if (not state.end_sequence) && Hash_set.mem symbol_starts state.address
-            then
-              Hashtbl.set
-                table
-                ~key:state.address
-                ~data:
-                  { Location.filename = Owee_debug_line.get_filename header state
-                  ; line = state.line
-                  ; col = state.col
-                  }
-          in
-          Owee_debug_line.fold_rows (header, chunk) process ();
-          load_table_next ()
-      in
-      load_table_next ());
+  let symbol_starts = Int.Hash_set.create () in
+  Owee_elf.Symbol_table.iter t.symbol ~f:(fun symbol ->
+      if is_func symbol
+      then
+        Hash_set.add
+          symbol_starts
+          (Owee_elf.Symbol_table.Symbol.value symbol |> Int64.to_int_exn));
+  traverse_debug_line
+    ~f:(fun header state ->
+      if Hash_set.mem symbol_starts state.address
+      then
+        Hashtbl.set
+          table
+          ~key:state.address
+          ~data:
+            { Location.filename = Owee_debug_line.get_filename header state
+            ; line = state.line
+            ; col = state.col
+            })
+    t;
   table
 ;;
 

core/elf.mli

@@ -11,7 +11,7 @@ val create : Filename.t -> t option
     to stop when executing code from that symbol. The filter string syntax is in ftrace
     format and is accepted both by the perf command line tool and the
     [PERF_EVENT_IOC_SET_FILTER] ioctl for the [perf_event_open] syscall. *)
-val symbol_stop_info : t -> Pid.t -> Owee_elf.Symbol_table.Symbol.t -> Stop_info.t
+val selection_stop_info : t -> Pid.t -> Selection.t -> Stop_info.t
 
 val addr_table : t -> Addr_table.t
 val ocaml_exception_info : t -> Ocaml_exception_info.t option
@@ -22,6 +22,7 @@ val matching_functions : t -> Re.re -> Owee_elf.Symbol_table.Symbol.t String.Map
 
 val all_symbols : t -> (string * Owee_elf.Symbol_table.Symbol.t) list
 val find_symbol : t -> string -> Owee_elf.Symbol_table.Symbol.t option
+val find_selection : t -> string -> Selection.t option
 
 module Symbol_resolver : sig
   type nonrec t =

core/elf_intf.ml

@@ -20,3 +20,12 @@ module Stop_info = struct
     ; filter : string
     }
 end
+
+module Selection = struct
+  type t =
+    | Symbol of Owee_elf.Symbol_table.Symbol.t
+    | Address of
+        { address : int
+        ; name : string
+        }
+end

core/when_to_snapshot.ml

@@ -14,14 +14,18 @@ let param =
        from:\n\
        (1) If you do not provide [-trigger], magic-trace takes a snapshot when you \
        Ctrl+C it.\n\
-       (2) If you provide [-trigger ?], magic-trace will open up a fuzzy-find dialog to \
-       help you choose a symbol to trace. Fails if you don't have the \"fzf\" binary in \
-       your PATH.\n\
-       (3) If you provide [-trigger <FUNCTION NAME>], magic-trace will snapshot when the \
+       (2) If you provide [-trigger ?], [-trigger line:?], or [-trigger symbol:?], \
+       magic-trace will open up a fuzzy-find dialog to help you choose a symbol to \
+       trace. Fails if you don't have the \"fzf\" binary in your PATH. The prefixes \
+       [symbol:] and [line:] will restrict to only seeing either functions or file \
+       locations respectively.\n\
+       (3) If you provide [-trigger <SELECTION>], magic-trace will snapshot when the \
        application being traced calls the given function. This takes the fully-mangled \
        name, so if you're using anything except C, fuzzy-find mode will probably be \
        easier to use. [-trigger .] is a shorthand for [-trigger \
-       magic_trace_stop_indicator].\n\
+       magic_trace_stop_indicator]. [SELECTION] can be [<FUNCTION NAME>], [<ADDRESS>], \
+       or [<FILE>:<LINE>:<COL>]. Prefixes [addr:], [symbol:], and [line:] restrict how \
+       magic-trace should parse the selection.\n\
        (*) Regardless of trigger mode, magic-trace will always snapshot when the \
        application terminates if it has not yet triggered for any other reason."
   |> map ~f:(function

src/trace.ml

@@ -328,11 +328,11 @@ module Make_commands (Backend : Backend_intf.S) = struct
           let%bind snap_sym =
             Deferred.return
               (Result.of_option
-                 (Elf.find_symbol elf symbol_name)
+                 (Elf.find_selection elf symbol_name)
                  ~error:
                    (Error.of_string [%string "Snapshot symbol not found: %{symbol_name}"]))
           in
-          let snap_loc = Elf.symbol_stop_info elf head_pid snap_sym in
+          let snap_loc = Elf.selection_stop_info elf head_pid snap_sym in
           return (Some snap_loc))
     in
     let%map.Deferred.Or_error recording, recording_data =