diff --git a/src/optimize_machine.ml b/src/optimize_machine.ml
index c5056c3e0dc480f96fd74d02948d4caf7d673701..cde625bc4a469f8ce49cc05c30ff405353d58929 100644
--- a/src/optimize_machine.ml
+++ b/src/optimize_machine.ml
@@ -907,37 +907,56 @@ and instrs_replace_var m fvar instrs =
   let instrs = List.fold_left (instr_replace_var m fvar) [] instrs in
   List.rev instrs
 
-let add_ghost_assigns fvar instrs =
-  (* let is_reused_def v = Hashtbl.find_opt reuse v.var_id in *)
-  let add_reused vars i =
+let rec add_ghost_assign fvar instr =
+  let aux spec i =
     let v = fvar i in
-    if v.var_id = i.var_id then vars else VMap.add i v vars
+    if (fvar i).var_id = i.var_id then spec
+    else add_ghost_assign_spec i (vdecl_to_val v) spec
   in
-  let rec is_defining_reused vars instr =
-    match instr.instr_desc with
-    | MLocalAssign (i, _) ->
-      add_reused vars i
-    | MStep (il, _, _) ->
-      List.fold_left add_reused vars il
-    | MBranch (_, hl, _) ->
-      List.fold_left
-        (fun vars (_, instrs) -> List.fold_left is_defining_reused vars instrs)
-        vars
-        hl
-    | _ ->
-      vars
-  in
-  List.map
-    (fun instr ->
-      let vars = is_defining_reused VMap.empty instr in
-      let instr_spec =
-        VMap.fold
-          (fun i v -> add_ghost_assign_spec i (vdecl_to_val v))
-          vars
-          instr.instr_spec
-      in
-      { instr with instr_spec })
-    instrs
+  match instr.instr_desc with
+  | MLocalAssign (i, _) ->
+    let instr_spec = aux instr.instr_spec i in
+    { instr with instr_spec }
+  | MStep (il, _, _) ->
+    let instr_spec = List.fold_left aux instr.instr_spec il in
+    { instr with instr_spec }
+  | MBranch (e, hl, b) ->
+    let hl = List.map (fun (h, l) -> h, add_ghost_assigns fvar l) hl in
+    { instr with instr_desc = MBranch (e, hl, b) }
+  | _ ->
+    instr
+
+and add_ghost_assigns fvar = List.map (add_ghost_assign fvar)
+(* let is_reused_def v = Hashtbl.find_opt reuse v.var_id in *)
+(* let add_reused vars i = *)
+(*   let v = fvar i in *)
+(*   if v.var_id = i.var_id then vars else VMap.add i v vars *)
+(* in *)
+(* let rec is_defining_reused vars instr = *)
+(*   match instr.instr_desc with *)
+(*   | MLocalAssign (i, _) -> *)
+(*     add_reused vars i *)
+(*   | MStep (il, _, _) -> *)
+(*     List.fold_left add_reused vars il *)
+(*   | MBranch (_, hl, _) -> *)
+(*     List.fold_left *)
+(* (fun vars (_, instrs) -> List.fold_left is_defining_reused vars instrs) *)
+(*       vars *)
+(*       hl *)
+(*   | _ -> *)
+(*     vars *)
+(* in *)
+(* List.map *)
+(*   (fun instr -> *)
+(*     let vars = is_defining_reused VMap.empty instr in *)
+(*     let instr_spec = *)
+(*       VMap.fold *)
+(*         (fun i v -> add_ghost_assign_spec i (vdecl_to_val v)) *)
+(*         vars *)
+(*         instr.instr_spec *)
+(*     in *)
+(*     { instr with instr_spec }) *)
+(*   instrs *)
 
 let ghost_vd v = { v with var_id = "__" ^ v.var_id }