Fix next_instr_graph out-of-subroutine pointers

The disasm-based next_instr_graph would introduce next instructions out
of the current subroutine for eg. calls, jmp to plts, etc.

DwarfSynth/Simplest.ml

@@ -120,7 +120,9 @@ let map_option f = function
   | None -> None
   | Some x -> Some (f x)
 
-let build_next_instr (disasm: BStd.disasm): AddrSet.t AddrMap.t =
+exception Block_not_in_subroutine
+
+let build_next_instr sub_ranges (disasm: BStd.disasm): AddrSet.t AddrMap.t =
   (** Build a map of memory_address -> AddrSet.t holding, for each address, the
       set of instructions coming right after the instruction at given address.
       There might be multiple such addresses, if the current instruction is at
@@ -138,8 +140,9 @@ let build_next_instr (disasm: BStd.disasm): AddrSet.t AddrMap.t =
          with _ -> cur_map)
         in
         build_of_instr_list new_map (elt2 :: tl)
-    | (cur_mem, _) :: [] ->
+    | (cur_mem, cur_insn) :: [] ->
-      let last_addr = (try Some (to_int64_addr @@ BStd.Memory.min_addr cur_mem)
+      let last_addr =
+        (try Some (to_int64_addr @@ BStd.Memory.min_addr cur_mem)
          with _ -> None) in
       cur_map, last_addr
 
@@ -165,8 +168,24 @@ let build_next_instr (disasm: BStd.disasm): AddrSet.t AddrMap.t =
   in
 
   let build_of_block cur_map block =
+    (try
+       (* First, check that this block belongs to a subroutine *)
+       let block_first_address = (
+         try
+           to_int64_addr @@ BStd.Block.addr block
+         with _ -> raise Block_not_in_subroutine) in
+       let sub_first_addr, sub_last_addr = (
+         try AddrMap.find_last
+               (fun start_addr -> start_addr <= block_first_address) sub_ranges
+         with Not_found ->
+           raise Block_not_in_subroutine
+       ) in
+
+       (* Add the sequence of instuctions inside the block itself *)
        let cur_map, last_addr =
          build_of_instr_list cur_map (BStd.Block.insns block) in
+
+       (* Then the set of possible destinations for the block terminator *)
        (match last_addr with
         | Some last_addr ->
           let following_set = BStd.Graphs.Cfg.Node.outputs block cfg
@@ -176,17 +195,24 @@ let build_next_instr (disasm: BStd.disasm): AddrSet.t AddrMap.t =
                                        (block_addresses
                                           (BStd.Graphs.Cfg.Edge.dst edge))
                                        set)
+                              |> AddrSet.filter (fun addr ->
+                                  sub_first_addr <= addr
+                                  && addr <= sub_last_addr)
+                                (* ^ We must ensure the landing address belongs
+                                   to the current subroutine for our purpose *)
           in
           AddrMap.add last_addr following_set cur_map
         | None -> cur_map
        )
+     with Block_not_in_subroutine ->
+       cur_map
+    )
   in
 
   BStd.Seq.fold (BStd.Graphs.Cfg.nodes cfg)
     ~init:AddrMap.empty
     ~f:build_of_block
 
-
 let find_rbp_pop_set cfg entry =
   (** Returns a BStd.Tid.Set.t of the terms actually "popping" %rbp, that is,
       the terms that should trigger a change to RbpUndef of the %rbp register.
@@ -701,11 +727,29 @@ let of_prog prog next_instr_graph : subroutine_cfa_map =
     ~init:StrMap.empty
     ~f:fold_step
 
+let build_sub_ranges prog: (memory_address) AddrMap.t =
+  (** Builds a map mapping the first address of each subroutine to its last
+      address. This map can be interpreted as a list of address ranges with
+      easy fast access to a member (cf Map.S.find_first) *)
+
+  let fold_subroutine accu sub =
+    let first_addr = int64_addr_of sub in
+    let last_addr = find_last_addr sub in
+    AddrMap.add first_addr (last_addr) accu
+  in
+
+  let subroutines = BStd.Term.enum BStd.sub_t prog in
+  BStd.Seq.fold subroutines
+    ~init:AddrMap.empty
+    ~f:fold_subroutine
+
 let of_proj no_rbp_undef proj : subroutine_cfa_map =
   (** Extracts the `cfa_changes` of a project *)
   __settings.no_rbp_undef <- no_rbp_undef ;
-  let next_instr_graph = build_next_instr (BStd.Project.disasm proj) in
   let prog = BStd.Project.program proj in
+  let sub_ranges = build_sub_ranges prog in
+  let next_instr_graph =
+    build_next_instr sub_ranges (BStd.Project.disasm proj) in
   of_prog prog next_instr_graph
 
 let clean_lost_track_subs pre_dwarf : subroutine_cfa_map =