prunes unreachable code in the optimization pass

oasis/optimization

@@ -8,6 +8,6 @@ Library optimization_plugin
   Path: plugins/optimization
   FindlibName: bap-plugin-optimization
   CompiledObject: best
-  BuildDepends: bap, core_kernel, ppx_jane, regular
+  BuildDepends: bap, core_kernel, graphlib, ppx_jane, regular
   InternalModules: Optimization_main, Optimization_data
   XMETAExtraLines: tags="pass,analysis,optimization"

plugins/optimization/optimization_data.ml

@@ -1,6 +1,7 @@
 open Core_kernel
 open Bap.Std
 open Regular.Std
+open Graphlib.Std
 
 type jmp_update = {
   cond : exp;
@@ -43,24 +44,86 @@ let updates_of_sub sub =
 
 let create ~deads sub = {deads; updates = updates_of_sub sub}
 
-let apply sub {deads; updates} =
-  let apply_to_def d =
-    if Set.mem deads (Term.tid d) then None
-    else
-      match Map.find updates (Term.tid d) with
-      | None -> Some d
-      | Some (Rhs e) -> Some (Def.with_rhs d e)
-      | _ -> assert false in
-  let apply_to_jmp j =
-    match Map.find updates (Term.tid j) with
-    | None -> j
-    | Some (Jmp {cond; kind}) ->
-      let j = Jmp.with_cond j cond in
-      Jmp.with_kind j kind
-    | _ -> assert false in
+let (++) = Set.union
+
+let dead_jmps_of_blk b =
+  Term.to_sequence jmp_t b |>
+  Seq.fold ~init:(Set.empty (module Tid), false)
+    ~f:(fun (deads, is_unreachable) jmp ->
+        if is_unreachable
+        then Set.add deads (Term.tid jmp), is_unreachable
+        else match Jmp.cond jmp with
+          | Bil.Int x when x = Word.b1 -> deads, true
+          | Bil.Int x when x = Word.b0 ->
+            Set.add deads (Term.tid jmp), is_unreachable
+          | _ -> deads, is_unreachable) |> fst
+
+let dead_jmps sub =
+  Term.to_sequence blk_t sub |>
+  Seq.fold ~init:(Set.empty (module Tid))
+    ~f:(fun tids b -> tids ++ dead_jmps_of_blk b)
+
+let remove_dead_edges g dead_jmps =
+  let module G = Graphs.Tid in
+  Seq.fold (G.edges g)
+    ~init:g ~f:(fun g edge ->
+        if Set.mem dead_jmps (Graphs.Tid.Edge.label edge)
+        then G.Edge.remove edge g
+        else g)
+
+let dead_blks g =
+  let module G = Graphs.Tid in
+  fst @@
+  Graphlib.depth_first_search (module G) g
+    ~init:(Set.empty (module Tid), false)
+    ~start_tree:(fun node (deads, _) ->
+        deads, Tid.equal node G.start)
+    ~enter_node:(fun _ node (deads, is_reachable) ->
+        if is_reachable then deads, is_reachable
+        else Set.add deads node, is_reachable)
+
+let find_unreachable sub t =
+  let dead_jmps = dead_jmps sub in
+  let dead_blks =
+    remove_dead_edges (Sub.to_graph sub) dead_jmps |>
+    dead_blks in
+  {t with deads = t.deads ++ dead_jmps ++ dead_blks }
+
+let update_def updates d =
+  match Map.find updates (Term.tid d) with
+  | None -> d
+  | Some (Rhs e) -> Def.with_rhs d e
+  | _ -> assert false
+
+let update_jmp updates j =
+  match Map.find updates (Term.tid j) with
+  | None -> j
+  | Some (Jmp {cond; kind}) ->
+    let j = Jmp.with_cond j cond in
+    Jmp.with_kind j kind
+  | _ -> assert false
+
+let update sub {updates} =
   Term.map blk_t sub ~f:(fun b ->
-      Term.filter_map def_t b ~f:apply_to_def |>
-      Term.map jmp_t ~f:apply_to_jmp)
+      Term.map def_t b ~f:(update_def updates) |>
+      Term.map jmp_t ~f:(update_jmp updates))
+
+let filter_map_alive deads cls ?(f=ident) x =
+  Term.filter_map cls x ~f:(fun t ->
+      if Set.mem deads (Term.tid t) then None
+      else Some (f t))
+
+let remove_dead_code sub {deads} =
+  let update_blk b =
+    filter_map_alive deads def_t b |>
+    filter_map_alive deads jmp_t in
+  filter_map_alive deads blk_t sub ~f:update_blk
+
+let apply sub {deads; updates} =
+  let update_blk b =
+    filter_map_alive deads def_t b ~f:(update_def updates) |>
+    filter_map_alive deads jmp_t ~f:(update_jmp updates) in
+  filter_map_alive deads blk_t sub ~f:update_blk
 
 include Data.Make(struct
     type nonrec t = t

plugins/optimization/optimization_data.mli

@@ -10,4 +10,12 @@ val create : deads:Tid.Set.t -> sub term -> t
 
 val apply : sub term -> t -> sub term
 
+
+val update : sub term -> t -> sub term
+
+val find_unreachable : sub term -> t -> t
+
+val remove_dead_code : sub term -> t -> sub term
+
+
 include Data.S with type t := t

plugins/optimization/optimization_main.ml

@@ -134,16 +134,23 @@ let process_sub free can_touch sub =
   O.create dead sub'
 
 let digest_of_sub sub level =
-  let digest =
-    (object
-      inherit [Digest.t] Term.visitor
-      method! enter_arg t dst = Digest.add dst "%a" Arg.pp t
-      method! enter_def t dst = Digest.add dst "%a" Def.pp t
-      method! enter_jmp t dst = Digest.add dst "%a" Jmp.pp t
-    end)#visit_sub sub
-      (Digest.create ~namespace:"optimization") in
+  let add addrs t = match Term.get_attr t address with
+    | None -> addrs
+    | Some a -> Set.add addrs a in
+  let addrs =
+    Term.to_sequence blk_t sub |>
+    Seq.fold ~init:(Set.empty (module Addr))
+      ~f:(fun addrs b ->
+          Seq.fold (Blk.elts b) ~init:(add addrs b)
+            ~f:(fun addrs -> function
+                | `Def d -> add addrs d
+                | `Jmp j -> add addrs j
+                | `Phi p -> add addrs p)) in
+  let digest = Digest.create ~namespace:"optimization" in
+  let digest = Set.fold addrs ~init:digest ~f:(fun d a ->
+      Digest.add d "%a" Addr.pp a) in
   let digest = Digest.add digest "%s" (Sub.name sub) in
-  Digest.add digest "%s" (string_of_int level)
+  Digest.add digest "%d" level
 
 let run level proj =
   let arch = Project.arch proj in
@@ -153,13 +160,14 @@ let run level proj =
   Project.with_program proj @@
   Term.map sub_t prog ~f:(fun sub ->
       let digest = digest_of_sub sub level in
-      let data = match O.Cache.load digest with
-        | Some data -> data
-        | None ->
-          let data = process_sub free can_touch sub in
-          O.Cache.save digest data;
-          data in
-      O.apply sub data)
+      match O.Cache.load digest with
+      | Some data -> O.apply sub data
+      | None ->
+        let data = process_sub free can_touch sub in
+        let sub = O.update sub data in
+        let data = O.find_unreachable sub data in
+        O.Cache.save digest data;
+        O.remove_dead_code sub data)
 
 let () =
   Config.manpage [