BAP CFG to AST patch

cfg.diff

Index: ocaml/cfg_ast.ml
===================================================================
--- ocaml/cfg_ast.ml	(revision 7629)
+++ ocaml/cfg_ast.ml	(working copy)
@@ -13,7 +13,7 @@
 module C = Cfg.AST
 
 
-module D = Debug.Make(struct let name = "CFG_AST" and default=`NoDebug end)
+module D = Debug.Make(struct let name = "Cfg_ast" and default=`NoDebug end)
 open D
 
 let v2s v = bbid_to_string(C.G.V.label v)
@@ -176,6 +176,7 @@
 *)
 let to_prog c =
   let size = C.G.nb_vertex c in
+  let c = C.remove_vertex c (C.G.V.create BB_Indirect) in
   let module BH = Hashtbl.Make(C.G.V) in
   let tails = BH.create size (* maps head vertex to the tail of the trace *)
     (* maps vertex to succ it was joined with, forming a trace *)
@@ -184,7 +185,11 @@
   let get_revstmts b =
     try BH.find hrevstmts b
     with Not_found ->
-      let s = List.rev (C.get_stmts c b) in
+      let s = match C.G.V.label b with
+        | BB _ -> List.rev (C.get_stmts c b)
+        (* Don't include special node contents *)
+        | _ -> []
+      in
       BH.add hrevstmts b s;
       s
   in
@@ -193,14 +198,14 @@
   let rec grow_trace cond head =
       match bh_find_option tails head with
       | None ->
-	  () (* must have already been joined previously*)
+	  () (* must have already been joined previously *)
       | Some tail ->
 	  assert(not(BH.mem joined tail));
 	  let rec find_succ = function
 	    | [] -> ()
 	    | suc::rest ->
 		match bh_find_option tails suc with
-		| Some succtail when cond tail suc &&
+		| Some succtail when cond head tail suc &&
                                      suc <> head ->
                   assert (succtail <> head);
                   assert (suc <> head);
@@ -262,22 +267,18 @@
 		     ^ v2s src ^ " points to"^dests)
   in
   (* join traces without jumps *)
-  grow_traces (fun b suc -> normal b && joinable suc && not(has_jump b));
+  grow_traces (fun _ b suc -> normal b && normal suc && not(has_jump b));
   (* join other traces (if we cared, we could remove some jumps later) *)
-  grow_traces (fun b suc -> normal b && joinable suc);
-  (* join the entry node, NOT with the trace containing the exit *)
+  grow_traces (fun _ b suc -> normal b && normal suc);
+  (* join the entry node *)
   grow_trace
-    (fun b suc ->
-       let suctail = BH.find tails suc in
-       C.G.V.label suctail <> BB_Exit
-    )
+    (fun h t suc -> (normal t || C.G.V.label t = BB_Entry) && normal suc)
     (C.G.V.create BB_Entry);
-  (* add jumps for edges that need them *)
-  C.G.iter_vertex 
-    (fun b -> 
-       C.G.iter_succ (fun s -> if not(BH.mem joined b) then ensure_jump b s) c b
-    )
-    c;
+  (* now join traces with exit, but do NOT join with the entry trace.
+     the entry trace must be printed first, but the exit trace must be
+     printed last. *)
+  grow_traces (fun h t suc -> normal h && normal t && joinable suc);
+  (* Make sure the exit trace is last *)
   let revordered_heads, exittrace =
     BH.fold
       (fun h t (rh,et) ->
@@ -302,6 +303,25 @@
     in
     List.fold_right head_to_revnodes revordered_heads []
   in
+  (* Because the entry trace must go first, and the exit trace must go
+     last, they are not joined, since it's not always possible for
+     these to happen at the same time.  Sometimes it is possible,
+     though, and here we identify the predecessor trace of the exit,
+     to avoid a jump to the exit if one is not needed. *)
+  let exit_edge = 
+    let special v =
+      match C.G.V.label v with | BB _ -> true | BB_Entry -> true | BB_Exit -> true | _ -> false
+    in
+    match List.filter special revnodes with
+    | x::y::_ when C.G.V.label x = BB_Exit -> dprintf "ok %s,%s" (C.v2s y) (C.v2s x); Some (y, x)
+    | _ -> None
+  in
+  (* add jumps for edges that need them *)
+  C.G.iter_vertex 
+    (fun b -> 
+       C.G.iter_succ (fun s -> if not(BH.mem joined b) && Some(b, s) <> exit_edge then ensure_jump b s) c b
+    )
+    c;
   let add_stmts stmts b =
     dprintf "to_prog: Adding statements for %s" (v2s b);
     let stmts = List.rev_append (get_revstmts b) stmts in