preserves brancher information in the BIL code of an instruction

lib/bap_disasm/bap_disasm_rec.ml

@@ -299,6 +299,43 @@ let create_indexes (dests : dests Addr.Table.t) =
 
 let filter_valid s = {s with inits = Set.inter s.inits s.valid}
 
+let join_destinations ?default dests =
+  let jmp x = [ Bil.(jmp (int x)) ] in
+  let undecided x =
+    Bil.if_ (Bil.unknown "destination" (Type.Imm 1)) (jmp x) [] in
+  let init = match default with
+    | None -> []
+    | Some x -> jmp x in
+  Set.fold dests ~init ~f:(fun bil x -> undecided x :: bil)
+
+let make_switch x dests =
+  let case addr = Bil.(if_ (x = int addr) [jmp (int addr)] []) in
+  let default = Bil.jmp x in
+  Set.fold dests ~init:[default] ~f:(fun ds a -> case a :: ds)
+
+let dests_of_bil bil =
+  (object inherit [Addr.Set.t] Stmt.visitor
+    method! visit_jmp e dests = match e with
+      | Int w -> Set.add dests w
+      | _ -> dests
+  end)#run bil Addr.Set.empty
+
+let add_destinations bil = function
+  | [] -> bil
+  | dests ->
+    let d = dests_of_bil bil in
+    let d' = Addr.Set.of_list dests in
+    let n = Set.diff d' d in
+    if Set.is_empty n then bil
+    else
+    if has_jump bil then
+      (object inherit Stmt.mapper
+        method! map_jmp = function
+          | Int addr -> join_destinations ~default:addr n
+          | indirect -> make_switch indirect n
+      end)#run bil
+    else bil @ join_destinations n
+
 let stage2 dis stage1 =
   let stage1 = filter_valid stage1 in
   let leads, terms, kinds = create_indexes stage1.dests in
@@ -353,9 +390,17 @@ let stage2 dis stage1 =
             Dis.stop s (Dis.insns s)) |>
       List.map ~f:(function
           | mem, None -> mem,(None,None)
-          | mem, (Some ins as insn) -> match stage1.lift mem ins with
-            | Ok bil -> mem,(insn,Some bil)
-            | _ -> mem, (insn, None)) in
+          | mem, (Some ins as insn) ->
+            let dests =
+              Addrs.find stage1.dests (Memory.max_addr mem) |>
+              Option.value ~default:[] |>
+              List.filter_map ~f:(function
+                  | a, (`Cond | `Jump) -> a
+                  | _ -> None) in
+            let bil = match stage1.lift mem ins with
+              | Ok bil -> bil
+              | _ -> [] in
+            mem, (insn, Some (add_destinations bil dests))) in
     return {stage1; addrs; succs; preds; disasm}
 
 let stage3 s2 =