resolve RIP in some x86 fp semantics

plugins/optimization/optimization_main.ml

@@ -1,154 +1,156 @@
-open Core_kernel[@@warning "-D"]
+open Core_kernel [@@warning "-D"]
 open Bap.Std
 open Regular.Std
 open Format
-include Self()
-
+include Self ()
 module Digest = Data.Cache.Digest
 module O = Optimization_data
 
 type level = int
-
 type cls = Phys | Virt | Flag
 
 let classify is_flag var =
-  if is_flag var then Flag
-  else if Var.is_physical var then Phys
-  else Virt
+  if is_flag var then Flag else if Var.is_physical var then Phys else Virt
 
 let is_optimization_allowed is_flag level var =
   match classify is_flag var with
   | Virt -> level > 0
   | Flag -> level > 1
   | Phys -> level > 2
 
-let def_use_collector = object
-  inherit [Var.Set.t * Var.Set.t] Term.visitor
+let def_use_collector =
+  object
+    inherit [Var.Set.t * Var.Set.t] Term.visitor
 
-  method! enter_def t (defs,uses) =
-    Set.add defs (Def.lhs t), Set.union uses (Def.free_vars t)
+    method! enter_def t (defs, uses) =
+      (Set.add defs (Def.lhs t), Set.union uses (Def.free_vars t))
 
-  method! enter_phi t (defs,uses) =
-    Set.add defs (Phi.lhs t), Set.union uses (Phi.free_vars t)
+    method! enter_phi t (defs, uses) =
+      (Set.add defs (Phi.lhs t), Set.union uses (Phi.free_vars t))
 
-  method! enter_jmp t (defs,uses) =
-    defs, Set.union uses (Jmp.free_vars t)
-end
+    method! enter_jmp t (defs, uses) = (defs, Set.union uses (Jmp.free_vars t))
+  end
 
 let computed_def_use sub =
-  def_use_collector#visit_sub sub (Var.Set.empty,Var.Set.empty)
+  def_use_collector#visit_sub sub (Var.Set.empty, Var.Set.empty)
 
 let compute_dead can_touch protected sub =
-  let defs,uses = computed_def_use sub in
+  let defs, uses = computed_def_use sub in
   let dead = Set.diff defs uses in
   let live v = not (Set.mem dead v) in
-  Term.enum blk_t sub |>
-  Seq.fold ~init:Tid.Set.empty ~f:(fun dead blk ->
-      Term.enum ~rev:true def_t blk |>
-      Seq.fold ~init:(protected,dead) ~f:(fun (protected,dead) def ->
+  Term.enum blk_t sub
+  |> Seq.fold ~init:Tid.Set.empty ~f:(fun dead blk ->
+      Term.enum ~rev:true def_t blk
+      |> Seq.fold ~init:(protected, dead) ~f:(fun (protected, dead) def ->
           let v = Def.lhs def in
-          if not (can_touch v) || live v
-          then (protected,dead)
-          else if Set.mem protected (Var.base v)
-          then Set.remove protected (Var.base v), dead
-          else protected, Set.add dead (Term.tid def)) |>
-      snd)
+          if (not (can_touch v)) || live v then (protected, dead)
+          else if Set.mem protected (Var.base v) then
+            (Set.remove protected (Var.base v), dead)
+          else (protected, Set.add dead (Term.tid def)))
+      |> snd)
 
 let is_alive dead t = not (Set.mem dead (Term.tid t))
 let live_phi dead blk = Term.filter phi_t ~f:(is_alive dead) blk
 
 let live_def can_touch dead blk =
   Term.filter def_t blk ~f:(fun d ->
-      not (can_touch (Def.lhs d)) || is_alive dead d)
+      (not (can_touch (Def.lhs d))) || is_alive dead d)
 
 let rec substitute vars exp =
-  let substituter = object
-    inherit Exp.mapper as super
-    method! map_let var ~exp ~body =
-      let exp = super#map_exp exp in
-      let body = substitute (Map.remove vars var) body in
-      Bil.let_ var exp body
-    method! map_var v = match Map.find vars v with
-      | None -> Bil.var v
-      | Some e -> e
-  end in
+  let substituter =
+    object
+      inherit Exp.mapper as super
+
+      method! map_let var ~exp ~body =
+        let exp = super#map_exp exp in
+        let body = substitute (Map.remove vars var) body in
+        Bil.let_ var exp body
+
+      method! map_var v =
+        match Map.find vars v with None -> Bil.var v | Some e -> e
+    end
+  in
   substituter#map_exp exp |> Exp.fold_consts
 
 let equal_kinds j j' = compare_jmp_kind (Jmp.kind j) (Jmp.kind j') = 0
 
 let substitute sub vars =
   let def d =
     let rhs = substitute vars (Def.rhs d) in
-    if Exp.(Def.rhs d <> rhs) then
-      O.mark_updated (Def.with_rhs d rhs)
-    else d in
+    if Exp.(Def.rhs d <> rhs) then O.mark_updated (Def.with_rhs d rhs) else d
+  in
   let jmp j =
     let j' = Jmp.map_exp j ~f:(substitute vars) in
-    if Exp.(Jmp.cond j <> Jmp.cond j') || not (equal_kinds j j')
-    then O.mark_updated j'
-    else j in
+    if Exp.(Jmp.cond j <> Jmp.cond j') || not (equal_kinds j j') then
+      O.mark_updated j'
+    else j
+  in
   Term.map blk_t sub ~f:(Blk.map_elts ~def ~jmp)
 
 (* A simple constant propagation. Note, the input is required to be in SSA. *)
 let propagate_consts can_touch sub =
-  Seq.fold (Term.enum blk_t sub) ~init:Var.Map.empty
-    ~f:(fun vars b ->
-        Seq.fold (Term.enum def_t b) ~init:vars
-          ~f:(fun vars d ->
-              let v = Def.lhs d in
-              if can_touch v then
-                match Def.rhs d with
-                | Bil.Unknown _ | Bil.Int _ as exp ->
-                  Map.set vars ~key:v ~data:exp
-                | _ -> vars
-              else vars)) |>
-  substitute sub
+  Seq.fold (Term.enum blk_t sub) ~init:Var.Map.empty ~f:(fun vars b ->
+      Seq.fold (Term.enum def_t b) ~init:vars ~f:(fun vars d ->
+          let v = Def.lhs d in
+          if can_touch v then
+            match Def.rhs d with
+            | (Bil.Unknown _ | Bil.Int _) as exp ->
+                Map.set vars ~key:v ~data:exp
+            | _ -> vars
+          else vars))
+  |> substitute sub
 
 let clean can_touch dead sub =
   Term.map blk_t sub ~f:(fun b -> live_def can_touch dead b |> live_phi dead)
 
 let is_flag arch =
-  let module T = (val (target_of_arch arch)) in
+  let module T = (val target_of_arch arch) in
   T.CPU.is_flag
 
 let free_vars prog =
   report_progress ~note:"free-vars" ();
-  let (++) = Set.union in
+  let ( ++ ) = Set.union in
   let collect cls t ~f =
-    Seq.fold (Term.enum cls t) ~init:Var.Set.empty
-      ~f:(fun acc x -> acc ++ f x) in
+    Seq.fold (Term.enum cls t) ~init:Var.Set.empty ~f:(fun acc x -> acc ++ f x)
+  in
   let sub_free sub = Sub.free_vars sub |> Set.filter ~f:Var.is_physical in
   let sub_args sub =
     collect arg_t sub ~f:(fun arg ->
-        Set.add (Exp.free_vars (Arg.rhs arg)) (Arg.lhs arg)) in
+        Set.add (Exp.free_vars (Arg.rhs arg)) (Arg.lhs arg))
+  in
   collect sub_t prog ~f:(fun sub -> sub_args sub ++ sub_free sub)
 
 let process_sub free can_touch sub =
   let rec loop dead s =
     let s = propagate_consts can_touch s in
     let dead' = compute_dead can_touch free s in
     let dead = Set.union dead dead' in
-    if Set.is_empty dead' then s, dead
-    else loop dead (clean can_touch dead' s) in
+    if Set.is_empty dead' then (s, dead)
+    else loop dead (clean can_touch dead' s)
+  in
   let sub', dead = loop Tid.Set.empty (Sub.ssa sub) in
   O.create dead sub'
 
 let digest_of_sub sub level =
-  let add addrs t = match Term.get_attr t address with
+  let add addrs t =
+    match Term.get_attr t address with
     | None -> addrs
-    | Some a -> Set.add addrs a in
+    | 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
+    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 =
+    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 "%d" level
 
@@ -158,73 +160,64 @@ let run level proj =
   Project.map_program proj ~f:(fun prog ->
       let free = free_vars prog in
       Term.map sub_t prog ~f:(fun sub ->
-          let digest = digest_of_sub sub level in
-          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 data = process_sub free can_touch sub in
+          let sub = O.update sub data in
+          let data = O.find_unreachable sub data in
+          O.remove_dead_code sub data))
 
 let () =
-  Config.manpage [
-    `S "SYNOPSIS";
-    `Pre "
-     $(b,--)$(mname)
-";
-    `S "DESCRIPTION";
-
-    `P "An autorun pass that conservatively removes dead code. The
-  removed dead code is usually produced by a lifter, though it might
-  be possible that a binary indeed contains a dead code. The algorithm
-  doesn't remove variables that are stored in memory, only registers
-  are considered";
-
-    `S "ALGORITHM";
-
-    `P "Applies constant folding, dead code elimination, and constant
-  propagation in a loop until the fixed point is reached. The
-  algorithm is interprocedural, however it is not call graph
-  sensitive, as instead of considering the call graph we use an over
-  approximation that any function can call any other, thus any
-  variables that occurs free in any function is considered to be
-  non-constant. The algorithm is, however, flow sensitive on the
-  control flow graph level and uses the SSA form to encode data
-  facts. Since, it is not always safe to rely on the control flow
-  integrity and CFG precision, by default we optimize only flags and
-  virtual variables under a presumption that those two kind of data
-  points rarely used non-locally.
-
-  See the $(b,--optimization-level) parameter for the list of
-  available optimization levels and their consequences.";
-
-    `S "DEPENDENCIES";
-    `P "$(b,bap-plugin-api)(1)";
-
-    `S "SEE ALSO";
-    `P "$(b,bap-plugin-api)(1), $(b,bap-plugin-ssa)(1)";
-  ];
+  Config.manpage
+    [
+      `S "SYNOPSIS";
+      `Pre "\n     $(b,--)$(mname)\n";
+      `S "DESCRIPTION";
+      `P
+        "An autorun pass that conservatively removes dead code. The\n\
+        \  removed dead code is usually produced by a lifter, though it might\n\
+        \  be possible that a binary indeed contains a dead code. The algorithm\n\
+        \  doesn't remove variables that are stored in memory, only registers\n\
+        \  are considered";
+      `S "ALGORITHM";
+      `P
+        "Applies constant folding, dead code elimination, and constant\n\
+        \  propagation in a loop until the fixed point is reached. The\n\
+        \  algorithm is interprocedural, however it is not call graph\n\
+        \  sensitive, as instead of considering the call graph we use an over\n\
+        \  approximation that any function can call any other, thus any\n\
+        \  variables that occurs free in any function is considered to be\n\
+        \  non-constant. The algorithm is, however, flow sensitive on the\n\
+        \  control flow graph level and uses the SSA form to encode data\n\
+        \  facts. Since, it is not always safe to rely on the control flow\n\
+        \  integrity and CFG precision, by default we optimize only flags and\n\
+        \  virtual variables under a presumption that those two kind of data\n\
+        \  points rarely used non-locally.\n\n\
+        \  See the $(b,--optimization-level) parameter for the list of\n\
+        \  available optimization levels and their consequences.";
+      `S "DEPENDENCIES";
+      `P "$(b,bap-plugin-api)(1)";
+      `S "SEE ALSO";
+      `P "$(b,bap-plugin-api)(1), $(b,bap-plugin-ssa)(1)";
+    ];
   let level =
     let doc =
-      "Specifies the optimization level. The higher the value the more
-       aggressive (and less safe) optimizations are applied. On level
-       0 we touch nothing, only some constant folding may occur. On
-       level 1 we optimize only the synthetic code that was generated
-       by the lifter. Since such code can't leave a scope of instruction
-       it is not affected by the imprecision of a control flow graph.
-       On level 2, we also move and optimize processor flags. This
-       removes a significant amount of code and simplifies the program
-       and is a fair compromise between safety and performance.
-       (Since flags are rarely used non-locally).
-       Finally, on level 3 we extend our analysis to all variables." in
-
-    Config.(param int ~default:0 ~doc "level") in
-
-  Config.declare_extension
-    ~doc:"removes dead code and propagates constants"
-    ~provides:["pass"; "optimization"; "analysis"; "simplification"]
-    (fun {Config.get=(!)} ->
-       if !level > 0
-       then Project.register_pass ~deps:["api"] ~autorun:true (run !level))
+      "Specifies the optimization level. The higher the value the more\n\
+      \       aggressive (and less safe) optimizations are applied. On level\n\
+      \       0 we touch nothing, only some constant folding may occur. On\n\
+      \       level 1 we optimize only the synthetic code that was generated\n\
+      \       by the lifter. Since such code can't leave a scope of instruction\n\
+      \       it is not affected by the imprecision of a control flow graph.\n\
+      \       On level 2, we also move and optimize processor flags. This\n\
+      \       removes a significant amount of code and simplifies the program\n\
+      \       and is a fair compromise between safety and performance.\n\
+      \       (Since flags are rarely used non-locally).\n\
+      \       Finally, on level 3 we extend our analysis to all variables."
+    in
+
+    Config.(param int ~default:0 ~doc "level")
+  in
+
+  Config.declare_extension ~doc:"removes dead code and propagates constants"
+    ~provides:[ "pass"; "optimization"; "analysis"; "simplification" ]
+    (fun { Config.get = ( ! ) } ->
+      if !level > 0 then
+        Project.register_pass ~deps:[ "api" ] ~autorun:true (run !level))

plugins/primus_lisp/primus_lisp_semantic_primitives.ml

@@ -256,6 +256,10 @@ let export = Primus.Lisp.Type.Spec.[
     "get-current-program-counter", unit @-> int,
     "(get-current-program-counter) is an alias to (get-program-counter)";
 
+    "get-instruction-length", unit @-> int,
+    "(get-instruction-length) returns the length of the current instruction \
+     in bytes.";
+
     "set-symbol-value", tuple [any; a] @-> a,
     "(set-symbol-value S X) sets the value of the symbol S to X.
          Returns X";
@@ -839,6 +843,12 @@ module Primitives(CT : Theory.Core)(T : Target) = struct
     | None -> !!(empty s)
     | Some addr -> forget@@const_int s addr
 
+  let get_instruction_length s lbl =
+    let open Bap.Std in
+    KB.collect Memory.slot lbl >>= function
+    | None -> !!(empty s)
+    | Some mem -> forget@@int s (Memory.length mem)
+
   let require_symbol v k =
     match symbol v with
     | Some name -> k name
@@ -1479,6 +1489,7 @@ module Primitives(CT : Theory.Core)(T : Target) = struct
     | "store-word",_-> data@@store_word t args
     | "get-program-counter",[]
     | "get-current-program-counter",[] -> pure@@get_pc s lbl
+    | "get-instruction-length",[] -> pure@@get_instruction_length s lbl
     | "set-symbol-value",[sym;x] -> data@@set_symbol t sym x
     | "symbol-concat",syms -> pure@@symbol_concat s syms
     | "symbol",[x] -> pure@@mksymbol s x