WIP

Author: hhugo

compiler/lib/code.ml

@@ -480,6 +480,17 @@ let block pc p = Addr.Map.find pc p.blocks
 
 let add_block pc block p = { p with blocks = Addr.Map.add pc block p.blocks }
 
+let update_block pc p ~f =
+  { p with
+    blocks =
+      Addr.Map.update
+        pc
+        (function
+          | None -> raise Not_found
+          | Some b -> Some (f b))
+        p.blocks
+  }
+
 let remove_block pc p = { p with blocks = Addr.Map.remove pc p.blocks }
 
 let free_pc p =

compiler/lib/code.mli

@@ -231,6 +231,8 @@ val add_block : Addr.t -> block -> program -> program
 
 val remove_block : Addr.t -> program -> program
 
+val update_block : Addr.t -> program -> f:(block -> block) -> program
+
 val program : Addr.t -> block Addr.Map.t -> program
 
 val map_blocks : f:(block -> block) -> program -> program

compiler/lib/eval.ml

@@ -778,20 +778,13 @@ let drop_exception_handler drop_count p =
               incr drop_count;
               let b = { b with branch = Branch cont1 } in
               let p = Code.add_block pc b p in
-              let blocks =
-                List.fold_left
-                  ~f:(fun blocks pc2 ->
-                    Addr.Map.update
-                      pc2
-                      (function
-                        | Some ({ branch = Poptrap cont; _ } as b) ->
-                            Some { b with branch = Branch cont }
-                        | None | Some _ -> assert false)
-                      blocks)
-                  rewrite
-                  ~init:(Code.blocks p)
-              in
-              Code.program (Code.start p) blocks)
+              List.fold_left
+                ~f:(fun p pc2 ->
+                  Code.update_block pc2 p ~f:(function
+                    | { branch = Poptrap cont; _ } as b -> { b with branch = Branch cont }
+                    | _ -> assert false))
+                rewrite
+                ~init:p)
       | _ -> p)
     (Code.blocks p)
     p

compiler/lib/subst.ml

@@ -66,9 +66,7 @@ module Excluding_Binders = struct
   let block s block =
     { params = block.params; body = instrs s block.body; branch = last s block.branch }
 
-  let program s p =
-    let blocks = Addr.Map.map (fun b -> block s b) (Code.blocks p) in
-    Code.program (Code.start p) blocks
+  let program s p = Code.map_blocks p ~f:(fun b -> block s b)
 
   let rec cont' s pc p visited =
     if Addr.Set.mem pc visited