Insert final restoration code before scheduling heuristic again

This is the restoration code that will be part of the superblock even
after "if-lifting" (moving compensation/restoration code out of the
superblock)

+ Minor formatting

1 files changed, 34 insertions, 7 deletions

diff --git a/scheduling/MyRTLpathScheduleraux.ml b/scheduling/MyRTLpathScheduleraux.ml
index 8661a11b..2fb72d78 100644
--- a/scheduling/MyRTLpathScheduleraux.ml
+++ b/scheduling/MyRTLpathScheduleraux.ml
@@ -373,12 +373,12 @@ let restoration_instructions live_renames =
   to_insert
 
 let restoration_instructions' live_renames ~next_free_reg =
-let next_free_reg =
-  let next_free_reg = ref next_free_reg in
-  (fun () ->
-    let r = !next_free_reg in
-    next_free_reg := Camlcoq.P.succ !next_free_reg;
-    r )
+  let next_free_reg =
+    let next_free_reg = ref next_free_reg in
+    (fun () ->
+      let r = !next_free_reg in
+      next_free_reg := Camlcoq.P.succ !next_free_reg;
+      r )
   in
 
   let to_rename = PTree.map1
@@ -1283,7 +1283,34 @@ let scheduler f =
   debug "\n";
   flush_all ();
 
-  (* TODO: Insert final restoration code before downschedule_compensation_code *)
+
+  let (sb_renamings, code, pm, next_free_pc) = ListLabels.fold_left sb_renamings
+    ~init:([], code, pm, next_free_pc)
+    ~f:(fun (sbs, code, pm, next_free_pc) (sb, live_renames) ->
+      let nr_instr = Array.length sb.instructions in
+      let last_pc = sb.instructions.(nr_instr - 1) in
+      let final_renames = PTree.map (fun pc renames -> if Camlcoq.P.eq pc last_pc then renames else []) live_renames in
+      let final_restoration = restoration_instructions final_renames in
+      let (sb, code, pm, next_free_pc, fwmap) = insert_code sb code pm final_restoration ~next_free_pc in
+      let live_renames =
+        PTree.fold
+          (fun acc pc insts ->
+            let pc' = apply_map' fwmap pc in
+            let insts = if Camlcoq.P.eq pc last_pc then [] else insts in
+            PTree.set pc' insts acc)
+          live_renames
+          PTree.empty
+      in
+      ((sb, live_renames)::sbs, code, pm, next_free_pc)
+    )
+  in
+
+  debug "After inserting the final restoration code:\n";
+  (* print_code code; *)
+  (* print_path_map pm; *)
+  print_superblocks (fst @@ List.split sb_renamings) code;
+  debug "\n";
+  flush_all ();
 
   (* WARNING: mutation *)
   let sb_tocompensatepcs_liverenames =