From c859dc3f5c432823524adfc1049eab4bdcfea205 Mon Sep 17 00:00:00 2001
From: Hugo Heuzard <hugo.heuzard@gmail.com>
Date: Mon, 26 May 2025 16:40:56 +0200
Subject: [PATCH] WIP

---
 compiler/lib-wasm/closure_conversion.ml |  12 +-
 compiler/lib-wasm/generate.ml           |  44 ++---
 compiler/lib-wasm/globalize.ml          |   4 +-
 compiler/lib/code.ml                    | 118 +++++++-----
 compiler/lib/code.mli                   |  33 ++--
 compiler/lib/deadcode.ml                | 229 ++++++++++++------------
 compiler/lib/duplicate.ml               |  32 ++--
 compiler/lib/effects.ml                 | 134 +++++++-------
 compiler/lib/eval.ml                    |  72 +++-----
 compiler/lib/flow.ml                    |   9 +-
 compiler/lib/freevars.ml                |  24 +--
 compiler/lib/generate.ml                |  33 ++--
 compiler/lib/generate_closure.ml        | 108 +++++------
 compiler/lib/global_deadcode.ml         |  15 +-
 compiler/lib/global_flow.ml             |  34 ++--
 compiler/lib/inline.ml                  | 124 ++++++-------
 compiler/lib/lambda_lifting.ml          |  28 ++-
 compiler/lib/lambda_lifting_simple.ml   |  39 ++--
 compiler/lib/parse_bytecode.ml          |   7 +-
 compiler/lib/partial_cps_analysis.ml    |   4 +-
 compiler/lib/phisimpl.ml                |   5 +-
 compiler/lib/pure_fun.ml                |  20 +--
 compiler/lib/specialize.ml              | 190 +++++++++-----------
 compiler/lib/specialize_js.ml           |  39 ++--
 compiler/lib/structure.ml               |  34 ++--
 compiler/lib/structure.mli              |   2 +-
 compiler/lib/subst.ml                   |  30 ++--
 compiler/lib/subst.mli                  |   7 -
 compiler/lib/tailcall.ml                |  63 +++----
 29 files changed, 695 insertions(+), 798 deletions(-)

diff --git a/compiler/lib-wasm/closure_conversion.ml b/compiler/lib-wasm/closure_conversion.ml
index 8f30e2fc90..9be79d39c5 100644
--- a/compiler/lib-wasm/closure_conversion.ml
+++ b/compiler/lib-wasm/closure_conversion.ml
@@ -60,7 +60,7 @@ let collect_free_vars program var_depth depth pc closures =
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc () ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       Freevars.iter_block_free_vars add_if_free_variable block;
       List.iter block.body ~f:(fun i ->
           match i with
@@ -71,7 +71,7 @@ let collect_free_vars program var_depth depth pc closures =
               | Some _ | None -> ())
           | _ -> ()))
     pc
-    program.blocks
+    program
     ();
   !vars
 
@@ -87,7 +87,7 @@ let rec traverse var_depth closures program pc depth =
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc (program : Code.program) ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       mark_bound_variables var_depth block depth;
       let program =
         List.fold_left
@@ -151,9 +151,9 @@ let rec traverse var_depth closures program pc depth =
             in
             List.concat (List.rev (Array.to_list l)))
       in
-      { program with blocks = Code.Addr.Map.add pc { block with body } program.blocks })
+      Code.add_block pc { block with body } program)
     pc
-    program.blocks
+    program
     program
 
 let f p =
@@ -161,6 +161,6 @@ let f p =
   let nv = Var.count () in
   let var_depth = Array.make nv (-1) in
   let closures = ref Var.Map.empty in
-  let p = traverse var_depth closures p p.start 0 in
+  let p = traverse var_depth closures p (Code.start p) 0 in
   if Debug.find "times" () then Format.eprintf "  closure conversion: %a@." Timer.print t;
   p, !closures
diff --git a/compiler/lib-wasm/generate.ml b/compiler/lib-wasm/generate.ml
index 1ca2e82346..f9f32066e3 100644
--- a/compiler/lib-wasm/generate.ml
+++ b/compiler/lib-wasm/generate.ml
@@ -36,7 +36,7 @@ module Generate (Target : Target_sig.S) = struct
     { live : int array
     ; in_cps : Effects.in_cps
     ; deadcode_sentinal : Var.t
-    ; blocks : block Addr.Map.t
+    ; p : program
     ; closures : Closure_conversion.closure Var.Map.t
     ; global_context : Code_generation.context
     }
@@ -830,7 +830,7 @@ module Generate (Target : Target_sig.S) = struct
     Code.traverse
       { fold = fold_children_skip_try_body }
       (fun pc n ->
-        let block = Addr.Map.find pc p.blocks in
+        let block = Code.block pc p in
         List.fold_left
           ~f:(fun n i ->
             match i with
@@ -863,7 +863,7 @@ module Generate (Target : Target_sig.S) = struct
           ~init:n
           block.body)
       pc
-      p.blocks
+      p
       (false, false)
 
   let wrap_with_handler needed pc handler ~result_typ ~fall_through ~context body =
@@ -914,10 +914,10 @@ module Generate (Target : Target_sig.S) = struct
       ((pc, _) as cont)
       cloc
       acc =
-    let g = Structure.build_graph ctx.blocks pc in
+    let g = Structure.build_graph ctx.p pc in
     let dom = Structure.dominator_tree g in
     let rec translate_tree result_typ fall_through pc context =
-      let block = Addr.Map.find pc ctx.blocks in
+      let block = Code.block pc ctx.p in
       let keep_ouside pc' =
         match block.branch with
         | Switch _ -> true
@@ -925,7 +925,7 @@ module Generate (Target : Target_sig.S) = struct
         | _ -> Structure.is_merge_node g pc'
       in
       let code ~context =
-        let block = Addr.Map.find pc ctx.blocks in
+        let block = Code.block pc ctx.p in
         let* () = translate_instrs ctx context block.body in
         translate_node_within
           ~result_typ
@@ -960,7 +960,7 @@ module Generate (Target : Target_sig.S) = struct
             if
               (not (List.is_empty rem))
               ||
-              let block = Addr.Map.find pc ctx.blocks in
+              let block = Code.block pc ctx.p in
               match block.branch with
               | Cond _ | Pushtrap _ -> false (*ZZZ also some Switch*)
               | _ -> true
@@ -970,7 +970,7 @@ module Generate (Target : Target_sig.S) = struct
           in
           translate_tree result_typ fall_through pc' context
       | [] -> (
-          let block = Addr.Map.find pc ctx.blocks in
+          let block = Code.block pc ctx.p in
           let branch = block.branch in
           match branch with
           | Branch cont -> translate_branch result_typ fall_through pc cont context
@@ -1028,7 +1028,7 @@ module Generate (Target : Target_sig.S) = struct
         if List.is_empty args
         then return ()
         else
-          let block = Addr.Map.find dst ctx.blocks in
+          let block = Code.block dst ctx.p in
           parallel_renaming block.params args
       in
       match fall_through with
@@ -1077,7 +1077,7 @@ module Generate (Target : Target_sig.S) = struct
         ~param_names
         ~body:
           (let* () =
-             let block = Addr.Map.find pc ctx.blocks in
+             let block = Code.block pc ctx.p in
              match block.body with
              | Event start_loc :: _ -> event start_loc
              | _ -> no_event
@@ -1190,13 +1190,7 @@ module Generate (Target : Target_sig.S) = struct
   Code.Print.program (fun _ _ -> "") p;
 *)
     let ctx =
-      { live = live_vars
-      ; in_cps
-      ; deadcode_sentinal
-      ; blocks = p.blocks
-      ; closures
-      ; global_context
-      }
+      { live = live_vars; in_cps; deadcode_sentinal; p; closures; global_context }
     in
     let toplevel_name = Var.fresh_n "toplevel" in
     let functions =
@@ -1275,16 +1269,12 @@ let fix_switch_branches p =
                with
                | Some x -> x
                | None ->
-                   let pc' = !p'.free_pc in
+                   let pc' = Code.free_pc !p' in
                    p' :=
-                     { !p' with
-                       blocks =
-                         Addr.Map.add
-                           pc'
-                           { params = []; body = []; branch = Branch cont }
-                           !p'.blocks
-                     ; free_pc = pc' + 1
-                     };
+                     Code.add_block
+                       pc'
+                       { params = []; body = []; branch = Branch cont }
+                       !p';
                    updates := Addr.Map.add pc ((args, pc') :: l) !updates;
                    pc')
             , [] ))
@@ -1295,7 +1285,7 @@ let fix_switch_branches p =
       match block.branch with
       | Switch (_, l) -> fix_branches l
       | _ -> ())
-    p.blocks;
+    (Code.blocks p);
   !p'
 
 module G = Generate (Gc_target)
diff --git a/compiler/lib-wasm/globalize.ml b/compiler/lib-wasm/globalize.ml
index c4bd6c8ca6..d1fbbe715e 100644
--- a/compiler/lib-wasm/globalize.ml
+++ b/compiler/lib-wasm/globalize.ml
@@ -100,13 +100,13 @@ let traverse_instruction st i =
   | Event _ -> st
 
 let traverse_block p st pc =
-  let b = Code.Addr.Map.find pc p.Code.blocks in
+  let b = Code.block pc p in
   let st = List.fold_left ~f:(fun st x -> declare x st) ~init:st b.Code.params in
   List.fold_left ~f:(fun st i -> traverse_instruction st i) ~init:st b.Code.body
 
 let f p g closures =
   let l = Structure.blocks_in_reverse_post_order g in
-  let in_loop = Freevars.find_loops_in_closure p p.Code.start in
+  let in_loop = Freevars.find_loops_in_closure p (Code.start p) in
   let st =
     List.fold_left
       ~f:(fun st pc ->
diff --git a/compiler/lib/code.ml b/compiler/lib/code.ml
index eceb923f11..f42eee8219 100644
--- a/compiler/lib/code.ml
+++ b/compiler/lib/code.ml
@@ -469,9 +469,38 @@ type block =
 type program =
   { start : Addr.t
   ; blocks : block Addr.Map.t
-  ; free_pc : Addr.t
   }
 
+let start p = p.start
+
+let blocks p = p.blocks
+
+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 =
+  match Addr.Map.max_binding_opt p.blocks with
+  | None -> p.start + 1
+  | Some (pc, _) -> pc + 1
+
+let program start blocks = { start; blocks }
+
+let map_blocks ~f p = { p with blocks = Addr.Map.map f p.blocks }
+
 let noloc = No
 
 let location_of_pc pc = Before pc
@@ -657,7 +686,7 @@ let fold_closures p f accu =
 
 (****)
 
-let prepend ({ start; blocks; free_pc } as p) body =
+let prepend ({ start; blocks } as p) body =
   match body with
   | [] -> p
   | _ -> (
@@ -665,21 +694,13 @@ let prepend ({ start; blocks; free_pc } as p) body =
       | block ->
           { p with
             blocks = Addr.Map.add start { block with body = body @ block.body } blocks
-          }
-      | exception Not_found ->
-          let new_start = free_pc in
-          let blocks =
-            Addr.Map.add new_start { params = []; body; branch = Stop } blocks
-          in
-          let free_pc = free_pc + 1 in
-          { start = new_start; blocks; free_pc })
+          })
 
 let empty_block = { params = []; body = []; branch = Stop }
 
 let empty =
   let start = 0 in
-  let blocks = Addr.Map.singleton start empty_block in
-  { start; blocks; free_pc = start + 1 }
+  program start (Addr.Map.singleton start empty_block)
 
 let is_empty p =
   match Addr.Map.cardinal p.blocks with
@@ -694,41 +715,41 @@ let is_empty p =
       | _ -> false)
   | _ -> false
 
-let poptraps blocks pc =
-  let rec loop blocks pc visited depth acc =
+let poptraps p pc =
+  let rec loop p pc visited depth acc =
     if Addr.Set.mem pc visited
     then acc, visited
     else
       let visited = Addr.Set.add pc visited in
-      let block = Addr.Map.find pc blocks in
+      let block = block pc p in
       match block.branch with
       | Return _ | Raise _ | Stop -> acc, visited
-      | Branch (pc', _) -> loop blocks pc' visited depth acc
+      | Branch (pc', _) -> loop p pc' visited depth acc
       | Poptrap (pc', _) ->
           if depth = 0
           then Addr.Set.add pc' acc, visited
-          else loop blocks pc' visited (depth - 1) acc
+          else loop p pc' visited (depth - 1) acc
       | Pushtrap ((pc', _), _, (pc_h, _)) ->
-          let acc, visited = loop blocks pc' visited (depth + 1) acc in
-          let acc, visited = loop blocks pc_h visited depth acc in
+          let acc, visited = loop p pc' visited (depth + 1) acc in
+          let acc, visited = loop p pc_h visited depth acc in
           acc, visited
       | Cond (_, (pc1, _), (pc2, _)) ->
-          let acc, visited = loop blocks pc1 visited depth acc in
-          let acc, visited = loop blocks pc2 visited depth acc in
+          let acc, visited = loop p pc1 visited depth acc in
+          let acc, visited = loop p pc2 visited depth acc in
           acc, visited
       | Switch (_, a) ->
           let acc, visited =
             Array.fold_right
               ~init:(acc, visited)
-              ~f:(fun (pc, _) (acc, visited) -> loop blocks pc visited depth acc)
+              ~f:(fun (pc, _) (acc, visited) -> loop p pc visited depth acc)
               a
           in
           acc, visited
   in
-  loop blocks pc Addr.Set.empty 0 Addr.Set.empty |> fst
+  loop p pc Addr.Set.empty 0 Addr.Set.empty |> fst
 
-let fold_children blocks pc f accu =
-  let block = Addr.Map.find pc blocks in
+let fold_children p pc f accu =
+  let block = block pc p in
   match block.branch with
   | Return _ | Raise _ | Stop -> accu
   | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
@@ -744,13 +765,13 @@ let fold_children blocks pc f accu =
       let accu = Array.fold_right ~init:accu ~f:(fun (pc, _) accu -> f pc accu) a1 in
       accu
 
-let fold_children_skip_try_body blocks pc f accu =
-  let block = Addr.Map.find pc blocks in
+let fold_children_skip_try_body p pc f accu =
+  let block = block pc p in
   match block.branch with
   | Return _ | Raise _ | Stop -> accu
   | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
   | Pushtrap ((pc', _), _, (pc_h, _)) ->
-      let accu = Addr.Set.fold f (poptraps blocks pc') accu in
+      let accu = Addr.Set.fold f (poptraps p pc') accu in
       let accu = f pc_h accu in
       accu
   | Cond (_, (pc1, _), (pc2, _)) ->
@@ -761,7 +782,7 @@ let fold_children_skip_try_body blocks pc f accu =
       let accu = Array.fold_right ~init:accu ~f:(fun (pc, _) accu -> f pc accu) a1 in
       accu
 
-type 'c fold_blocs = block Addr.Map.t -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c -> 'c
+type 'c fold_blocs = program -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c -> 'c
 
 type fold_blocs_poly = { fold : 'a. 'a fold_blocs } [@@unboxed]
 
@@ -801,43 +822,43 @@ let rec preorder_traverse' { fold } f pc visited blocks acc =
 let preorder_traverse fold f pc blocks acc =
   snd (preorder_traverse' fold f pc Addr.Set.empty blocks acc)
 
-let fold_closures_innermost_first { start; blocks; _ } f accu =
-  let rec visit blocks pc f accu =
+let fold_closures_innermost_first p f accu =
+  let rec visit p pc f accu =
     traverse
       { fold = fold_children }
       (fun pc accu ->
-        let block = Addr.Map.find pc blocks in
+        let block = block pc p in
         List.fold_left block.body ~init:accu ~f:(fun accu i ->
             match i with
             | Let (x, Closure (params, cont, cloc)) ->
-                let accu = visit blocks (fst cont) f accu in
+                let accu = visit p (fst cont) f accu in
                 f (Some x) params cont cloc accu
             | _ -> accu))
       pc
-      blocks
+      p
       accu
   in
-  let accu = visit blocks start f accu in
-  f None [] (start, []) None accu
+  let accu = visit p p.start f accu in
+  f None [] (p.start, []) None accu
 
-let fold_closures_outermost_first { start; blocks; _ } f accu =
-  let rec visit blocks pc f accu =
+let fold_closures_outermost_first p f accu =
+  let rec visit p pc f accu =
     traverse
       { fold = fold_children }
       (fun pc accu ->
-        let block = Addr.Map.find pc blocks in
+        let block = block pc p in
         List.fold_left block.body ~init:accu ~f:(fun accu i ->
             match i with
             | Let (x, Closure (params, cont, cloc)) ->
                 let accu = f (Some x) params cont cloc accu in
-                visit blocks (fst cont) f accu
+                visit p (fst cont) f accu
             | _ -> accu))
       pc
-      blocks
+      p
       accu
   in
-  let accu = f None [] (start, []) None accu in
-  visit blocks start f accu
+  let accu = f None [] (p.start, []) None accu in
+  visit p p.start f accu
 
 let rec last_instr l =
   match l with
@@ -897,7 +918,7 @@ let cont_compare (pc, args) (pc', args') =
 
 let with_invariant = Debug.find "invariant"
 
-let do_compact { blocks; start; free_pc = _ } =
+let do_compact { blocks; start } =
   let remap =
     let max = fst (Addr.Map.max_binding blocks) in
     let a = Array.make (max + 1) 0 in
@@ -934,9 +955,8 @@ let do_compact { blocks; start; free_pc = _ } =
       blocks
       Addr.Map.empty
   in
-  let free_pc = (Addr.Map.max_binding blocks |> fst) + 1 in
   let start = remap.(start) in
-  { blocks; start; free_pc }
+  program start blocks
 
 let compact p =
   let t = Timer.make () in
@@ -957,19 +977,19 @@ let compact p =
   p
 
 let used_blocks p =
-  let visited = BitSet.create' p.free_pc in
+  let visited = BitSet.create' (free_pc p) in
   let rec mark_used pc =
     if not (BitSet.mem visited pc)
     then (
       BitSet.set visited pc;
-      let block = Addr.Map.find pc p.blocks in
+      let block = block pc p in
       List.iter
         ~f:(fun i ->
           match i with
           | Let (_, Closure (_, (pc', _), _)) -> mark_used pc'
           | _ -> ())
         block.body;
-      fold_children p.blocks pc (fun pc' () -> mark_used pc') ())
+      fold_children p pc (fun pc' () -> mark_used pc') ())
   in
   mark_used p.start;
   visited
diff --git a/compiler/lib/code.mli b/compiler/lib/code.mli
index c7b3696673..5b408cfc46 100644
--- a/compiler/lib/code.mli
+++ b/compiler/lib/code.mli
@@ -219,11 +219,25 @@ type block =
   ; branch : last
   }
 
-type program =
-  { start : Addr.t
-  ; blocks : block Addr.Map.t
-  ; free_pc : Addr.t
-  }
+type program
+
+val start : program -> Addr.t
+
+val blocks : program -> block Addr.Map.t
+
+val block : Addr.t -> program -> block
+
+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
+
+val free_pc : program -> Addr.t
 
 module Print : sig
   type xinstr =
@@ -247,7 +261,7 @@ module Print : sig
   val cont : Format.formatter -> cont -> unit
 end
 
-type 'c fold_blocs = block Addr.Map.t -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c -> 'c
+type 'c fold_blocs = program -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c -> 'c
 
 type fold_blocs_poly = { fold : 'a. 'a fold_blocs } [@@unboxed]
 
@@ -288,13 +302,12 @@ val fold_children : 'c fold_blocs
 
 val fold_children_skip_try_body : 'c fold_blocs
 
-val poptraps : block Addr.Map.t -> Addr.t -> Addr.Set.t
+val poptraps : program -> Addr.t -> Addr.Set.t
 
-val traverse :
-  fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> block Addr.Map.t -> 'c -> 'c
+val traverse : fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> program -> 'c -> 'c
 
 val preorder_traverse :
-  fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> block Addr.Map.t -> 'c -> 'c
+  fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> program -> 'c -> 'c
 
 val last_instr : instr list -> instr option
 (** Last instruction of a block body, ignoring events *)
diff --git a/compiler/lib/deadcode.ml b/compiler/lib/deadcode.ml
index fe6dec7bc6..2c11331111 100644
--- a/compiler/lib/deadcode.ml
+++ b/compiler/lib/deadcode.ml
@@ -41,7 +41,7 @@ let add_def defs x i =
 type variable_uses = int array
 
 type t =
-  { blocks : block Addr.Map.t
+  { p : program
   ; live : variable_uses
   ; defs : def list array
   ; reachable_blocks : BitSet.t
@@ -94,7 +94,7 @@ and mark_reachable st pc =
   if not (BitSet.mem st.reachable_blocks pc)
   then (
     BitSet.set st.reachable_blocks pc;
-    let block = Addr.Map.find pc st.blocks in
+    let block = Code.block pc st.p in
     List.iter block.body ~f:(fun i ->
         match i with
         | Let (_, Prim (Extern "caml_update_dummy", [ Pv x; Pv y ])) ->
@@ -155,26 +155,24 @@ let rec filter_args st pl al =
   | [], [] -> []
   | _ -> assert false
 
-let filter_cont blocks st (pc, args) =
-  let params = (Addr.Map.find pc blocks).params in
+let filter_cont p st (pc, args) =
+  let params = (Code.block pc p).params in
   pc, filter_args st params args
 
-let filter_closure blocks st i =
+let filter_closure p st i =
   match i with
-  | Let (x, Closure (l, cont, gloc)) ->
-      Let (x, Closure (l, filter_cont blocks st cont, gloc))
+  | Let (x, Closure (l, cont, gloc)) -> Let (x, Closure (l, filter_cont p st cont, gloc))
   | _ -> i
 
-let filter_live_last blocks st l =
+let filter_live_last p st l =
   match l with
   | Return _ | Raise _ | Stop -> l
-  | Branch cont -> Branch (filter_cont blocks st cont)
-  | Cond (x, cont1, cont2) ->
-      Cond (x, filter_cont blocks st cont1, filter_cont blocks st cont2)
-  | Switch (x, a1) -> Switch (x, Array.map a1 ~f:(fun cont -> filter_cont blocks st cont))
+  | Branch cont -> Branch (filter_cont p st cont)
+  | Cond (x, cont1, cont2) -> Cond (x, filter_cont p st cont1, filter_cont p st cont2)
+  | Switch (x, a1) -> Switch (x, Array.map a1 ~f:(fun cont -> filter_cont p st cont))
   | Pushtrap (cont1, x, cont2) ->
-      Pushtrap (filter_cont blocks st cont1, x, filter_cont blocks st cont2)
-  | Poptrap cont -> Poptrap (filter_cont blocks st cont)
+      Pushtrap (filter_cont p st cont1, x, filter_cont p st cont2)
+  | Poptrap cont -> Poptrap (filter_cont p st cont)
 
 (****)
 
@@ -204,9 +202,9 @@ let remove_unused_blocks' p =
         let b = BitSet.mem used pc in
         if not b then incr count;
         b)
-      p.blocks
+      (Code.blocks p)
   in
-  { p with blocks }, !count
+  Code.program (Code.start p) blocks, !count
 
 let remove_unused_blocks p =
   let previous_p = p in
@@ -227,8 +225,8 @@ let rec add_arg_dep defs params args =
   | [], [] -> ()
   | _ -> assert false
 
-let add_cont_dep blocks defs (pc, args) =
-  let block = Addr.Map.find pc blocks in
+let add_cont_dep p defs (pc, args) =
+  let block = Code.block pc p in
   add_arg_dep defs block.params args
 
 let empty_body b =
@@ -239,7 +237,7 @@ let empty_body b =
 let merge_blocks p =
   let previous_p = p in
   let t = Timer.make () in
-  let preds = Array.make p.free_pc 0 in
+  let preds = Array.make (Code.free_pc p) 0 in
   let assigned = ref Var.Set.empty in
   let merged = ref 0 in
   let subst =
@@ -265,20 +263,20 @@ let merge_blocks p =
             mark_cont cont2
         | Poptrap cont -> mark_cont cont
         | Return _ | Raise _ | Stop -> ())
-      p.blocks
+      (Code.blocks p)
   in
   let p =
-    let visited = BitSet.create' p.free_pc in
-    let rec process_branch pc blocks =
-      let block = Addr.Map.find pc blocks in
+    let visited = BitSet.create' (Code.free_pc p) in
+    let rec process_branch pc p =
+      let block = Code.block pc p in
       match block.branch with
       | Branch (pc_, args) when preds.(pc_) = 1 ->
-          let to_inline = Addr.Map.find pc_ blocks in
+          let to_inline = Code.block pc_ p in
           if List.exists to_inline.params ~f:(fun x -> Var.Set.mem x !assigned)
-          then block, blocks
+          then block, p
           else (
             incr merged;
-            let to_inline, blocks = process_branch pc_ blocks in
+            let to_inline, p = process_branch pc_ p in
             List.iter2 args to_inline.params ~f:(fun arg param ->
                 Code.Var.propagate_name param arg;
                 subst.(Code.Var.idx param) <- arg);
@@ -297,23 +295,20 @@ let merge_blocks p =
                    aux block.body)
               }
             in
-            let blocks = Addr.Map.remove pc_ blocks in
-            let blocks = Addr.Map.add pc block blocks in
-            block, blocks)
-      | _ -> block, blocks
+            let p = Code.remove_block pc_ p in
+            let p = Code.add_block pc block p in
+            block, p)
+      | _ -> block, p
     in
-    let rec traverse pc blocks =
+    let rec traverse pc p =
       if BitSet.mem visited pc
-      then blocks
+      then p
       else
         let () = BitSet.set visited pc in
-        let _block, blocks = process_branch pc blocks in
-        Code.fold_children blocks pc traverse blocks
+        let _block, p = process_branch pc p in
+        Code.fold_children p pc traverse p
     in
-    let blocks =
-      Code.fold_closures p (fun _ _ (pc, _) _ blocks -> traverse pc blocks) p.blocks
-    in
-    { p with blocks }
+    Code.fold_closures p (fun _ _ (pc, _) _ p -> traverse pc p) p
   in
   let p =
     if !merged = 0
@@ -364,72 +359,69 @@ let remove_empty_blocks st (p : Code.program) : Code.program =
         then Addr.Hashtbl.add shortcuts pc (params, cont)
     | _ -> ()
   in
-  Addr.Map.iter register_block_if_empty p.blocks;
-  let blocks =
-    (* We are relying on the fact that forward branches target blocks
-       with higher addresses in the code generated by the OCaml
-       compiler. By processing the blocks in descending address order,
-       simplifying a branch can make it possible to simplify earlier
-       branches. *)
-    Seq.fold_left
-      (fun blocks (pc, block) ->
-        if
-          match block.branch with
-          | Branch (pc, _) | Poptrap (pc, _) -> not (Addr.Hashtbl.mem shortcuts pc)
-          | Cond (_, (pc1, _), (pc2, _)) | Pushtrap ((pc1, _), _, (pc2, _)) ->
-              not (Addr.Hashtbl.mem shortcuts pc1 || Addr.Hashtbl.mem shortcuts pc2)
-          | Switch (_, a) ->
-              not (Array.exists ~f:(fun (pc, _) -> Addr.Hashtbl.mem shortcuts pc) a)
-          | Return _ | Raise _ | Stop -> true
-        then blocks
-        else
-          Addr.Map.add
-            pc
-            (match block with
-            | { body; branch = Cond (x, cont1, cont2); _ } ->
-                let cont1' = resolve cont1 in
-                let cont2' = resolve cont2 in
-                if Code.cont_equal cont1' cont2'
-                then (
-                  let decr_usage x = st.live.(Var.idx x) <- st.live.(Var.idx x) - 1 in
-                  decr_usage x;
-                  let body =
-                    List.fold_right
-                      ~f:(fun i rem ->
-                        if live_instr st i
-                        then
-                          match i, rem with
-                          | Event _, Event _ :: _ -> rem
-                          | _ -> i :: rem
-                        else (
-                          Freevars.iter_instr_free_vars decr_usage i;
-                          rem))
-                      body
-                      ~init:[]
-                  in
-                  let block = { block with body; branch = Branch cont1' } in
-                  register_block_if_empty pc block;
-                  block)
-                else { block with branch = Cond (x, cont1', cont2') }
-            | _ ->
-                { block with
-                  branch =
-                    (let branch = block.branch in
-                     match branch with
-                     | Branch cont -> Branch (resolve cont)
-                     | Switch (x, a1) -> Switch (x, Array.map ~f:resolve a1)
-                     | Pushtrap (cont1, x, cont2) ->
-                         Pushtrap (resolve cont1, x, resolve cont2)
-                     | Poptrap cont -> Poptrap (resolve cont)
-                     | Cond _ | Return _ | Raise _ | Stop -> assert false)
-                })
-            blocks)
-      p.blocks
-      (Addr.Map.to_rev_seq p.blocks)
-  in
-  { p with blocks }
+  Addr.Map.iter register_block_if_empty (Code.blocks p);
+  (* We are relying on the fact that forward branches target blocks
+     with higher addresses in the code generated by the OCaml
+     compiler. By processing the blocks in descending address order,
+     simplifying a branch can make it possible to simplify earlier
+     branches. *)
+  Seq.fold_left
+    (fun p (pc, block) ->
+      if
+        match block.branch with
+        | Branch (pc, _) | Poptrap (pc, _) -> not (Addr.Hashtbl.mem shortcuts pc)
+        | Cond (_, (pc1, _), (pc2, _)) | Pushtrap ((pc1, _), _, (pc2, _)) ->
+            not (Addr.Hashtbl.mem shortcuts pc1 || Addr.Hashtbl.mem shortcuts pc2)
+        | Switch (_, a) ->
+            not (Array.exists ~f:(fun (pc, _) -> Addr.Hashtbl.mem shortcuts pc) a)
+        | Return _ | Raise _ | Stop -> true
+      then p
+      else
+        Code.add_block
+          pc
+          (match block with
+          | { body; branch = Cond (x, cont1, cont2); _ } ->
+              let cont1' = resolve cont1 in
+              let cont2' = resolve cont2 in
+              if Code.cont_equal cont1' cont2'
+              then (
+                let decr_usage x = st.live.(Var.idx x) <- st.live.(Var.idx x) - 1 in
+                decr_usage x;
+                let body =
+                  List.fold_right
+                    ~f:(fun i rem ->
+                      if live_instr st i
+                      then
+                        match i, rem with
+                        | Event _, Event _ :: _ -> rem
+                        | _ -> i :: rem
+                      else (
+                        Freevars.iter_instr_free_vars decr_usage i;
+                        rem))
+                    body
+                    ~init:[]
+                in
+                let block = { block with body; branch = Branch cont1' } in
+                register_block_if_empty pc block;
+                block)
+              else { block with branch = Cond (x, cont1', cont2') }
+          | _ ->
+              { block with
+                branch =
+                  (let branch = block.branch in
+                   match branch with
+                   | Branch cont -> Branch (resolve cont)
+                   | Switch (x, a1) -> Switch (x, Array.map ~f:resolve a1)
+                   | Pushtrap (cont1, x, cont2) ->
+                       Pushtrap (resolve cont1, x, resolve cont2)
+                   | Poptrap cont -> Poptrap (resolve cont)
+                   | Cond _ | Return _ | Raise _ | Stop -> assert false)
+              })
+          p)
+    p
+    (Addr.Map.to_rev_seq (Code.blocks p))
 
-let f pure_funs ({ blocks; _ } as p : Code.program) =
+let f pure_funs (p : Code.program) =
   let previous_p = p in
   Code.invariant p;
   let t = Timer.make () in
@@ -446,21 +438,21 @@ let f pure_funs ({ blocks; _ } as p : Code.program) =
               ());
       match block.branch with
       | Return _ | Raise _ | Stop -> ()
-      | Branch cont -> add_cont_dep blocks defs cont
+      | Branch cont -> add_cont_dep p defs cont
       | Cond (_, cont1, cont2) ->
-          add_cont_dep blocks defs cont1;
-          add_cont_dep blocks defs cont2
-      | Switch (_, a1) -> Array.iter a1 ~f:(fun cont -> add_cont_dep blocks defs cont)
+          add_cont_dep p defs cont1;
+          add_cont_dep p defs cont2
+      | Switch (_, a1) -> Array.iter a1 ~f:(fun cont -> add_cont_dep p defs cont)
       | Pushtrap (cont, _, cont_h) ->
-          add_cont_dep blocks defs cont_h;
-          add_cont_dep blocks defs cont
-      | Poptrap cont -> add_cont_dep blocks defs cont)
-    blocks;
+          add_cont_dep p defs cont_h;
+          add_cont_dep p defs cont
+      | Poptrap cont -> add_cont_dep p defs cont)
+    (Code.blocks p);
   let st =
     { live
     ; defs
-    ; blocks
-    ; reachable_blocks = BitSet.create' p.free_pc
+    ; p
+    ; reachable_blocks = BitSet.create' (Code.free_pc p)
     ; pure_funs
     ; deleted_instrs = 0
     ; deleted_blocks = 0
@@ -468,10 +460,9 @@ let f pure_funs ({ blocks; _ } as p : Code.program) =
     ; block_shortcut = 0
     }
   in
-  mark_reachable st p.start;
+  mark_reachable st (Code.start p);
   if debug () then Print.program Format.err_formatter (fun pc xi -> annot st pc xi) p;
   let p =
-    let all_blocks = blocks in
     let blocks =
       Addr.Map.filter_map
         (fun pc block ->
@@ -490,16 +481,16 @@ let f pure_funs ({ blocks; _ } as p : Code.program) =
                           i :: prev
                       | _ ->
                           if live_instr st i
-                          then filter_closure all_blocks st i :: acc
+                          then filter_closure p st i :: acc
                           else (
                             st.deleted_instrs <- st.deleted_instrs + 1;
                             acc))
                   |> List.rev
-              ; branch = filter_live_last all_blocks st block.branch
+              ; branch = filter_live_last p st block.branch
               })
-        blocks
+        (Code.blocks p)
     in
-    { p with blocks }
+    Code.program (Code.start p) blocks
   in
   let p = remove_empty_blocks st p in
   if times () then Format.eprintf "  dead code elim.: %a@." Timer.print t;
diff --git a/compiler/lib/duplicate.ml b/compiler/lib/duplicate.ml
index e54684aa06..2f32b85d10 100644
--- a/compiler/lib/duplicate.ml
+++ b/compiler/lib/duplicate.ml
@@ -19,36 +19,36 @@
 open! Stdlib
 open Code
 
-let bound_variables { blocks; _ } ~f ~params ~cont:(pc, _) =
+let bound_variables p ~f ~params ~cont:(pc, _) =
   let bound_vars = ref Var.Map.empty in
   let add_var x = bound_vars := Var.Map.add x (Var.fork x) !bound_vars in
   List.iter ~f:add_var (f :: params);
-  let rec traverse blocks pc =
+  let rec traverse pc =
     Code.traverse
       { fold = fold_children }
       (fun pc _ ->
-        let block = Addr.Map.find pc blocks in
+        let block = Code.block pc p in
         Freevars.iter_block_bound_vars add_var block;
         List.iter
           ~f:(fun i ->
             match i with
             | Let (_, Closure (params, (pc', _), _)) ->
                 List.iter ~f:add_var params;
-                traverse blocks pc'
+                traverse pc'
             | _ -> ())
           block.body)
       pc
-      blocks
+      p
       ()
   in
-  traverse blocks pc;
+  traverse pc;
   !bound_vars
 
 let rec blocks_to_rename p pc lst =
   Code.traverse
     { fold = Code.fold_children }
     (fun pc lst ->
-      let block = Addr.Map.find pc p.blocks in
+      let block = Code.block pc p in
       List.fold_left
         ~f:(fun lst i ->
           match i with
@@ -57,27 +57,25 @@ let rec blocks_to_rename p pc lst =
         ~init:(pc :: lst)
         block.body)
     pc
-    p.blocks
+    p
     lst
 
 let closure p ~f ~params ~cont =
   let s = Subst.from_map (bound_variables p ~f ~params ~cont) in
   let pc, args = cont in
   let blocks = blocks_to_rename p pc [] in
-  let free_pc, m =
+  let _free_pc, m =
     List.fold_left
       ~f:(fun (pc', m) pc -> pc' + 1, Addr.Map.add pc pc' m)
-      ~init:(p.free_pc, Addr.Map.empty)
+      ~init:(Code.free_pc p, Addr.Map.empty)
       blocks
   in
-  let blocks =
+  let p =
     List.fold_left
-      ~f:(fun blocks pc ->
-        let b = Addr.Map.find pc blocks in
-        let b = Subst.Including_Binders.And_Continuations.block m s b in
-        Addr.Map.add (Addr.Map.find pc m) b blocks)
-      ~init:p.blocks
+      ~f:(fun p pc ->
+        let b = Subst.Including_Binders.And_Continuations.block m s (Code.block pc p) in
+        Code.add_block (Addr.Map.find pc m) b p)
+      ~init:p
       blocks
   in
-  let p = { p with blocks; free_pc } in
   p, s f, List.map ~f:s params, (Addr.Map.find pc m, List.map ~f:s args)
diff --git a/compiler/lib/effects.ml b/compiler/lib/effects.ml
index 8bec7e0a7e..0e776f0f93 100644
--- a/compiler/lib/effects.ml
+++ b/compiler/lib/effects.ml
@@ -62,7 +62,7 @@ type control_flow_graph =
   ; block_order : int Addr.Hashtbl.t
   }
 
-let build_graph blocks pc =
+let build_graph p pc =
   let succs = Addr.Hashtbl.create 16 in
   let l = ref [] in
   let visited = Addr.Hashtbl.create 16 in
@@ -70,7 +70,7 @@ let build_graph blocks pc =
     if not (Addr.Hashtbl.mem visited pc)
     then (
       Addr.Hashtbl.add visited pc ();
-      let successors = Code.fold_children blocks pc Addr.Set.add Addr.Set.empty in
+      let successors = Code.fold_children p pc Addr.Set.add Addr.Set.empty in
       Addr.Hashtbl.add succs pc successors;
       Addr.Set.iter traverse successors;
       l := pc :: !l)
@@ -171,7 +171,7 @@ also mark blocks that correspond to function continuations or
 exception handlers. And we keep track of the exception handler
 associated to each Poptrap, and possibly Raise.
 *)
-let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
+let compute_needed_transformations ~cfg ~idom ~cps_needed p start =
   let frontiers = dominance_frontier cfg idom in
   let transformation_needed = ref Addr.Set.empty in
   let matching_exn_handler = Addr.Hashtbl.create 16 in
@@ -197,7 +197,7 @@ let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
     then visited
     else
       let visited = Addr.Set.add pc visited in
-      let block = Addr.Map.find pc blocks in
+      let block = Code.block pc p in
       (match block.branch with
       | Branch (dst, _) -> (
           match Code.last_instr block.body with
@@ -219,7 +219,7 @@ let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
           | _ -> ())
       | _ -> ());
       Code.fold_children
-        blocks
+        p
         pc
         (fun pc visited ->
           let englobing_exn_handlers =
@@ -277,7 +277,7 @@ type in_cps = Var.Set.t
 type st =
   { mutable new_blocks : Code.block Addr.Map.t
   ; mutable free_pc : Code.Addr.t
-  ; blocks : Code.block Addr.Map.t
+  ; p : program
   ; cfg : control_flow_graph
   ; jc : jump_closures
   ; closure_info : (Var.t list * (Addr.t * Var.t list)) Addr.Hashtbl.t
@@ -588,7 +588,7 @@ let cps_block ~st ~k ~orig_pc block =
     | to_allocate ->
         List.map to_allocate ~f:(fun (cname, jump_pc) ->
             let params =
-              let jump_block = Addr.Map.find jump_pc st.blocks in
+              let jump_block = Code.block jump_pc st.p in
               (* For a function to be used as a continuation, it needs
                  exactly one parameter. So, we add a parameter if
                  needed. *)
@@ -776,20 +776,22 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
   let p =
     Code.fold_closures_innermost_first
       p
-      (fun name_opt params (start, args) _cloc ({ Code.blocks; free_pc; _ } as p) ->
+      (fun name_opt params (start, args) _cloc p ->
         (* We speculatively add a block at the beginning of the
            function. In case of tail-recursion optimization, the
            function implementing the loop body may have to be placed
            there. *)
         let initial_start = start in
-        let start', blocks' =
-          ( free_pc
-          , Addr.Map.add
-              free_pc
+        let p', start' =
+          let free_pc = Code.free_pc p in
+          let new_start = free_pc in
+          ( Code.add_block
+              new_start
               { params = []; body = []; branch = Branch (start, args) }
-              blocks )
+              p
+          , new_start )
         in
-        let cfg = build_graph blocks' start' in
+        let cfg = build_graph p' start' in
         let idom = dominator_tree cfg in
         let should_compute_needed_transformations =
           match name_opt with
@@ -802,28 +804,22 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
         in
         let blocks_to_transform, matching_exn_handler, is_continuation =
           if should_compute_needed_transformations
-          then
-            compute_needed_transformations
-              ~cfg
-              ~idom
-              ~cps_needed
-              ~blocks:blocks'
-              ~start:start'
+          then compute_needed_transformations ~cfg ~idom ~cps_needed p' start'
           else Addr.Set.empty, Addr.Hashtbl.create 1, Addr.Hashtbl.create 1
         in
         let closure_jc = jump_closures blocks_to_transform idom in
-        let start, args, blocks, free_pc =
+        let start, args, p =
           (* Insert an initial block if needed. *)
           if
             should_compute_needed_transformations
             && Addr.Map.mem start' closure_jc.closures_of_alloc_site
-          then start', [], blocks', free_pc + 1
-          else start, args, blocks, free_pc
+          then start', [], p'
+          else start, args, p
         in
         let st =
           { new_blocks = Addr.Map.empty
-          ; free_pc
-          ; blocks
+          ; free_pc = Code.free_pc p + 1
+          ; p
           ; cfg
           ; jc = closure_jc
           ; closure_info
@@ -857,16 +853,16 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
             { fold = Code.fold_children }
             (fun pc _ ->
               if Addr.Set.mem pc blocks_to_transform then Format.eprintf "CPS@.";
-              let block = Addr.Map.find pc blocks in
+              let block = Code.block pc p in
               Code.Print.block
                 Format.err_formatter
                 (fun _ xi -> Partial_cps_analysis.annot cps_needed xi)
                 pc
                 block)
             start
-            blocks
+            p
             ());
-        let blocks =
+        let p =
           (* For every block in the closure,
              1. CPS-translate it if needed. If we double-translate, add its CPS
                 translation to the block map at a fresh address. Otherwise,
@@ -915,18 +911,18 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
           in
           Code.traverse
             { fold = Code.fold_children }
-            (fun pc blocks ->
-              let block, cps_block_opt = transform_block pc (Addr.Map.find pc blocks) in
-              let blocks = Addr.Map.add pc block blocks in
+            (fun pc p ->
+              let block, cps_block_opt = transform_block pc (Code.block pc p) in
+              let p = Code.add_block pc block p in
               match cps_block_opt with
-              | None -> blocks
+              | None -> p
               | Some b ->
                   let cps_pc = mk_cps_pc_of_direct ~st pc in
                   st.new_blocks <- Addr.Map.add cps_pc b st.new_blocks;
-                  Addr.Map.add cps_pc b blocks)
+                  Code.add_block cps_pc b p)
             start
-            st.blocks
-            st.blocks
+            st.p
+            st.p
         in
         (* If double-translating, all variables bound in the CPS version will have to be
              subst with fresh ones to avoid clashing with the definitions in the original
@@ -937,25 +933,25 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
             Code.traverse
               Code.{ fold = fold_children }
               (fun pc () ->
-                let block = Addr.Map.find pc p.blocks in
+                let block = Code.block pc p in
                 Freevars.iter_block_bound_vars
                   (fun v -> subst_add_fresh cloned_vars v)
                   block)
               initial_start
-              p.blocks
+              p
               ();
             subst_bound_in_blocks st.new_blocks cloned_subst)
           else st.new_blocks
         in
-        let blocks =
+        let p =
           (* Remove the initial block added only for the CPS transformation *)
           if double_translate () && start <> initial_start
-          then Addr.Map.remove start blocks
-          else blocks
+          then Code.remove_block start p
+          else p
         in
-        let blocks = Addr.Map.fold Addr.Map.add new_blocks blocks in
+        let p = Addr.Map.fold Code.add_block new_blocks p in
         if debug () then Format.eprintf "@.";
-        { p with blocks; free_pc = st.free_pc })
+        p)
       p
   in
   (* Also apply our substitution to the sets of trampolined calls, and cps call sites *)
@@ -965,11 +961,11 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
     if double_translate ()
     then p
     else
-      match Addr.Hashtbl.find_opt closure_info p.start with
+      match Addr.Hashtbl.find_opt closure_info (Code.start p) with
       | None -> p
       | Some (cps_params, cps_cont) ->
           (* Call [caml_cps_trampoline] to set up the execution context. *)
-          let new_start = p.free_pc in
+          let new_start = Code.free_pc p in
           let blocks =
             let main = Var.fresh () in
             let args = Var.fresh () in
@@ -984,9 +980,9 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
                   ]
               ; branch = Return res
               }
-              p.blocks
+              (Code.blocks p)
           in
-          { start = new_start; blocks; free_pc = new_start + 1 }
+          Code.program new_start blocks
   in
   p, !trampolined_calls, !in_cps
 
@@ -1011,16 +1007,12 @@ let wrap_call ~cps_needed p x f args accu =
 
 let wrap_primitive ~cps_needed (p : program) x e accu =
   let f = Var.fresh () in
-  let closure_pc = p.free_pc in
-  ( { p with
-      free_pc = p.free_pc + 1
-    ; blocks =
-        Addr.Map.add
-          closure_pc
-          (let y = Var.fresh () in
-           { params = []; body = [ Let (y, e) ]; branch = Return y })
-          p.blocks
-    }
+  let closure_pc = Code.free_pc p in
+  ( Code.add_block
+      closure_pc
+      (let y = Var.fresh () in
+       { params = []; body = [ Let (y, e) ]; branch = Return y })
+      p
   , Var.Set.remove x (Var.Set.add f cps_needed)
   , let args = Var.fresh () in
     [ Let (f, Closure ([], (closure_pc, []), None))
@@ -1041,8 +1033,8 @@ let rewrite_toplevel_instr (p, cps_needed, accu) instr =
    unncessary function nestings. This is not done inside loops since
    using repeatedly [caml_cps_trampoline] can be costly. *)
 let rewrite_toplevel ~cps_needed p =
-  let { start; blocks; _ } = p in
-  let cfg = build_graph blocks start in
+  let start = Code.start p in
+  let cfg = build_graph p start in
   let idom = dominator_tree cfg in
   let frontiers = dominance_frontier cfg idom in
   let rec traverse visited (p : Code.program) cps_needed in_loop pc =
@@ -1054,16 +1046,16 @@ let rewrite_toplevel ~cps_needed p =
       let p, cps_needed =
         if Option.is_none in_loop
         then
-          let block = Addr.Map.find pc p.blocks in
+          let block = Code.block pc p in
           let p, cps_needed, body_rev =
             List.fold_left ~f:rewrite_toplevel_instr ~init:(p, cps_needed, []) block.body
           in
           let body = List.concat @@ List.rev body_rev in
-          { p with blocks = Addr.Map.add pc { block with body } p.blocks }, cps_needed
+          Code.add_block pc { block with body } p, cps_needed
         else p, cps_needed
       in
       Code.fold_children
-        blocks
+        p
         pc
         (fun pc (visited, p, cps_needed) -> traverse visited p cps_needed in_loop pc)
         (visited, p, cps_needed)
@@ -1089,20 +1081,20 @@ let split_blocks ~cps_needed (p : Code.program) =
           && Var.Set.mem x cps_needed
       | _ -> false
     in
-    let rec split (p : Code.program) pc block accu l branch =
+    let rec split (p : Code.program) free_pc pc block accu l branch =
       match l with
       | [] ->
           let block = { block with body = List.rev accu } in
-          { p with blocks = Addr.Map.add pc block p.blocks }
+          Code.add_block pc block p
       | (Let (x, e) as i) :: r when is_split_point i r branch ->
-          let pc' = p.free_pc in
+          let pc' = free_pc in
           let block' = { params = []; body = []; branch = block.branch } in
           let block =
             { block with body = List.rev (Let (x, e) :: accu); branch = Branch (pc', []) }
           in
-          let p = { p with blocks = Addr.Map.add pc block p.blocks; free_pc = pc' + 1 } in
-          split p pc' block' [] r branch
-      | i :: r -> split p pc block (i :: accu) r branch
+          let p = Code.add_block pc block p in
+          split p (free_pc + 1) pc' block' [] r branch
+      | i :: r -> split p free_pc pc block (i :: accu) r branch
     in
     let rec should_split l branch =
       match l with
@@ -1110,10 +1102,10 @@ let split_blocks ~cps_needed (p : Code.program) =
       | i :: r -> is_split_point i r branch || should_split r branch
     in
     if should_split block.body block.branch
-    then split p pc block [] block.body block.branch
+    then split p (Code.free_pc p) pc block [] block.body block.branch
     else p
   in
-  Addr.Map.fold split_block p.blocks p
+  Addr.Map.fold split_block (Code.blocks p) p
 
 (****)
 
@@ -1127,6 +1119,7 @@ let f ~flow_info ~live_vars p =
     if double_translate ()
     then (
       let p, liftings = Lambda_lifting_simple.f ~to_lift:cps_needed p in
+      Code.invariant p;
       let cps_needed =
         Var.Set.map
           (fun f -> try Subst.from_map liftings f with Not_found -> f)
@@ -1148,6 +1141,7 @@ let f ~flow_info ~live_vars p =
       p, cps_needed
   in
   let p = split_blocks ~cps_needed p in
+  Code.invariant p;
   let p, trampolined_calls, in_cps = cps_transform ~live_vars ~flow_info ~cps_needed p in
   if Debug.find "times" () then Format.eprintf "  effects: %a@." Timer.print t;
   Code.invariant p;
diff --git a/compiler/lib/eval.ml b/compiler/lib/eval.ml
index 5ec43d2cb3..3ada9f502c 100644
--- a/compiler/lib/eval.ml
+++ b/compiler/lib/eval.ml
@@ -727,12 +727,12 @@ let eval_branch update_branch info l =
 
 exception May_raise
 
-let rec do_not_raise pc visited rewrite blocks =
+let rec do_not_raise pc visited rewrite p =
   if Addr.Set.mem pc visited
   then visited, rewrite
   else
     let visited = Addr.Set.add pc visited in
-    let b = Addr.Map.find pc blocks in
+    let b = Code.block pc p in
     List.iter b.body ~f:(fun i ->
         match i with
         | Event _
@@ -752,53 +752,45 @@ let rec do_not_raise pc visited rewrite blocks =
     | Raise _ -> raise May_raise
     | Stop | Return _ -> visited, rewrite
     | Poptrap _ -> visited, pc :: rewrite
-    | Branch (pc, _) -> do_not_raise pc visited rewrite blocks
+    | Branch (pc, _) -> do_not_raise pc visited rewrite p
     | Cond (_, (pc1, _), (pc2, _)) ->
-        let visited, rewrite = do_not_raise pc1 visited rewrite blocks in
-        let visited, rewrite = do_not_raise pc2 visited rewrite blocks in
+        let visited, rewrite = do_not_raise pc1 visited rewrite p in
+        let visited, rewrite = do_not_raise pc2 visited rewrite p in
         visited, rewrite
     | Switch (_, a1) ->
         let visited, rewrite =
           Array.fold_left
             a1
             ~init:(visited, rewrite)
-            ~f:(fun (visited, rewrite) (pc, _) -> do_not_raise pc visited rewrite blocks)
+            ~f:(fun (visited, rewrite) (pc, _) -> do_not_raise pc visited rewrite p)
         in
         visited, rewrite
     | Pushtrap _ -> raise May_raise
 
-let drop_exception_handler drop_count blocks =
+let drop_exception_handler drop_count p =
   Addr.Map.fold
-    (fun pc _ blocks ->
-      match Addr.Map.find pc blocks with
+    (fun pc _ p ->
+      match Code.block pc p with
       | { branch = Pushtrap (((addr, _) as cont1), _x, _cont2); _ } as b -> (
-          match do_not_raise addr Addr.Set.empty [] blocks with
-          | exception May_raise -> blocks
+          match do_not_raise addr Addr.Set.empty [] p with
+          | exception May_raise -> p
           | _visited, rewrite ->
               incr drop_count;
               let b = { b with branch = Branch cont1 } in
-              let blocks = Addr.Map.add pc b blocks 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:blocks
-              in
-              blocks)
-      | _ -> blocks)
-    blocks
-    blocks
-
-let eval update_count update_branch inline_constant ~target info blocks =
-  Addr.Map.map
-    (fun block ->
+              let p = Code.add_block pc b p in
+              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
+
+let eval update_count update_branch inline_constant ~target info p =
+  Code.map_blocks p ~f:(fun block ->
       let body =
         List.concat_map
           block.body
@@ -806,7 +798,6 @@ let eval update_count update_branch inline_constant ~target info blocks =
       in
       let branch = eval_branch update_branch info block.branch in
       { block with Code.body; Code.branch })
-    blocks
 
 let f info p =
   Code.invariant p;
@@ -816,17 +807,10 @@ let f info p =
   let inline_constant = ref 0 in
   let drop_count = ref 0 in
   let t = Timer.make () in
-  let blocks =
-    eval
-      update_count
-      update_branch
-      inline_constant
-      ~target:(Config.target ())
-      info
-      p.blocks
+  let p =
+    eval update_count update_branch inline_constant ~target:(Config.target ()) info p
   in
-  let blocks = drop_exception_handler drop_count blocks in
-  let p = { p with blocks } in
+  let p = drop_exception_handler drop_count p in
   if times () then Format.eprintf "  eval: %a@." Timer.print t;
   if stats ()
   then
diff --git a/compiler/lib/flow.ml b/compiler/lib/flow.ml
index adf814ec9e..051e552a3a 100644
--- a/compiler/lib/flow.ml
+++ b/compiler/lib/flow.ml
@@ -113,7 +113,7 @@ let expr_deps blocks vars deps defs x e =
   | Block (_, a, _, _) -> Array.iter a ~f:(fun y -> add_dep deps x y)
   | Field (y, _, _) -> add_dep deps x y
 
-let program_deps { blocks; _ } =
+let program_deps blocks =
   let nv = Var.count () in
   let vars = Var.ISet.empty () in
   let deps = Array.make nv Var.Set.empty in
@@ -266,7 +266,7 @@ let expr_escape st _x e =
       in
       loop l ka
 
-let program_escape defs known_origins { blocks; _ } =
+let program_escape defs known_origins blocks =
   let may_escape = Var.ISet.empty () in
   let possibly_mutable = Var.ISet.empty () in
   let st = { defs; known_origins; may_escape; possibly_mutable } in
@@ -485,13 +485,14 @@ let f p =
   Code.invariant p;
   let t = Timer.make () in
   let t1 = Timer.make () in
-  let vars, deps, defs = program_deps p in
+  let blocks = Code.blocks p in
+  let vars, deps, defs = program_deps blocks in
   if times () then Format.eprintf "    flow analysis 1: %a@." Timer.print t1;
   let t2 = Timer.make () in
   let known_origins = solver1 vars deps defs in
   if times () then Format.eprintf "    flow analysis 2: %a@." Timer.print t2;
   let t3 = Timer.make () in
-  let possibly_mutable = program_escape defs known_origins p in
+  let possibly_mutable = program_escape defs known_origins blocks in
   if times () then Format.eprintf "    flow analysis 3: %a@." Timer.print t3;
   let t4 = Timer.make () in
   let maybe_unknown = solver2 vars deps defs known_origins possibly_mutable in
diff --git a/compiler/lib/freevars.ml b/compiler/lib/freevars.ml
index c378572d61..092b9773b6 100644
--- a/compiler/lib/freevars.ml
+++ b/compiler/lib/freevars.ml
@@ -112,7 +112,7 @@ let find_loops p in_loop pc =
     incr index;
     Stack.push pc stack;
     Code.fold_children
-      p.blocks
+      p
       pc
       (fun pc' () ->
         try
@@ -155,12 +155,12 @@ let find_all_loops p =
 
 let mark_variables in_loop p =
   let vars = Var.Tbl.make () (-1) in
-  let visited = BitSet.create' p.free_pc in
+  let visited = BitSet.create' (Code.free_pc p) in
   let rec traverse pc =
     if not (BitSet.mem visited pc)
     then (
       BitSet.set visited pc;
-      let block = Addr.Map.find pc p.blocks in
+      let block = Code.block pc p in
       (try
          let pc' = Addr.Map.find pc in_loop in
          iter_block_bound_vars (fun x -> Var.Tbl.set vars x pc') block
@@ -169,20 +169,20 @@ let mark_variables in_loop p =
           match i with
           | Let (_, Closure (_, (pc', _), _)) -> traverse pc'
           | _ -> ());
-      Code.fold_children p.blocks pc (fun pc' () -> traverse pc') ())
+      Code.fold_children p pc (fun pc' () -> traverse pc') ())
   in
-  traverse p.start;
+  traverse (Code.start p);
   vars
 
 let free_variables vars in_loop p =
   let all_freevars = ref Addr.Map.empty in
   let freevars = ref Addr.Map.empty in
-  let visited = BitSet.create' p.free_pc in
+  let visited = BitSet.create' (Code.free_pc p) in
   let rec traverse pc =
     if not (BitSet.mem visited pc)
     then (
       BitSet.set visited pc;
-      let block = Addr.Map.find pc p.blocks in
+      let block = Code.block pc p in
       iter_block_free_vars
         (fun x ->
           let pc' = Var.Tbl.get vars x in
@@ -211,16 +211,16 @@ let free_variables vars in_loop p =
                 all_freevars := Addr.Map.remove pc'' !all_freevars
               with Not_found -> freevars := Addr.Map.add pc' Var.Set.empty !freevars)
           | _ -> ());
-      Code.fold_children p.blocks pc (fun pc' () -> traverse pc') ())
+      Code.fold_children p pc (fun pc' () -> traverse pc') ())
   in
-  traverse p.start;
+  traverse (Code.start p);
   !freevars
 
 let f p =
   Code.invariant p;
   let t = Timer.make () in
   let bound = Code.Var.ISet.empty () in
-  let visited = BitSet.create' p.free_pc in
+  let visited = BitSet.create' (Code.free_pc p) in
   let free_vars =
     Code.fold_closures_innermost_first
       p
@@ -233,14 +233,14 @@ let f p =
           if not (BitSet.mem visited pc)
           then (
             BitSet.set visited pc;
-            let block = Addr.Map.find pc p.blocks in
+            let block = Code.block pc p in
             iter_block_bound_vars (fun x -> Code.Var.ISet.add bound x) block;
             iter_block_free_vars using block;
             List.iter block.body ~f:(function
               | Let (_, Closure (_, (pc_clo, _), _)) ->
                   Code.Var.Set.iter using (Code.Addr.Map.find pc_clo acc)
               | _ -> ());
-            Code.fold_children p.blocks pc (fun pc' () -> traverse pc') ())
+            Code.fold_children p pc (fun pc' () -> traverse pc') ())
         in
         List.iter params ~f:(fun x -> Code.Var.ISet.add bound x);
         List.iter args ~f:using;
diff --git a/compiler/lib/generate.ml b/compiler/lib/generate.ml
index aaf394c337..fd979c1852 100644
--- a/compiler/lib/generate.ml
+++ b/compiler/lib/generate.ml
@@ -155,7 +155,7 @@ module Share = struct
       ?alias_strings
       ?(alias_prims = false)
       ?(alias_apply = true)
-      { blocks; _ } : t =
+      p : t =
     let alias_strings =
       match alias_strings with
       | None -> Config.Flag.use_js_string () && not (Config.Flag.share_constant ())
@@ -190,7 +190,7 @@ module Share = struct
                   add_args args share
               | Let (_, Prim (_, args)) -> add_args args share
               | _ -> share))
-        blocks
+        (Code.blocks p)
         empty_aux
     in
     let count =
@@ -293,7 +293,7 @@ end
 
 module Ctx = struct
   type t =
-    { blocks : block Addr.Map.t
+    { p : program
     ; live : Deadcode.variable_uses
     ; share : Share.t
     ; exported_runtime : (Code.Var.t * bool ref) option
@@ -314,11 +314,11 @@ module Ctx = struct
       ~mutated_vars
       ~freevars
       ~in_cps
-      blocks
+      p
       live
       trampolined_calls
       share =
-    { blocks
+    { p
     ; live
     ; share
     ; exported_runtime
@@ -831,7 +831,7 @@ end
 
 let build_graph ctx pc cloc =
   let visited_blocks = ref Addr.Set.empty in
-  let structure = Structure.build_graph ctx.Ctx.blocks pc in
+  let structure = Structure.build_graph ctx.Ctx.p pc in
   let dom = Structure.dominator_tree structure in
   { visited_blocks; structure; dom; ctx; cloc }
 
@@ -1710,14 +1710,11 @@ and translate_instrs_rev (ctx : Ctx.t) loc expr_queue instrs acc_rev muts_map =
         then ctx
         else
           let subst = Subst.from_map muts_map in
-          let p, _visited =
-            List.fold_left
-              pcs
-              ~init:(ctx.blocks, Addr.Set.empty)
-              ~f:(fun (blocks, visited) pc ->
-                Subst.Excluding_Binders.cont' subst pc blocks visited)
+          let p =
+            List.fold_left pcs ~init:ctx.p ~f:(fun p pc ->
+                Subst.Excluding_Binders.cont subst pc p)
           in
-          { ctx with blocks = p }
+          { ctx with p }
       in
       let vd kind = function
         | [] -> []
@@ -1852,7 +1849,7 @@ and compile_block_no_loop st loc queue (pc : Addr.t) ~fall_through scope_stack =
     assert false);
   if debug () then Format.eprintf "Compiling block %d@;" pc;
   st.visited_blocks := Addr.Set.add pc !(st.visited_blocks);
-  let block = Addr.Map.find pc st.ctx.blocks in
+  let block = Code.block pc st.ctx.p in
   let loc, seq, queue = translate_instrs st.ctx loc queue block.body in
   let nbbranch =
     match block.branch with
@@ -2124,7 +2121,7 @@ and compile_argument_passing ctx loc queue (pc, args) back_edge continuation =
   if List.is_empty args
   then continuation queue
   else
-    let block = Addr.Map.find pc ctx.Ctx.blocks in
+    let block = Code.block pc ctx.Ctx.p in
     parallel_renaming ctx loc back_edge block.params args continuation queue
 
 and compile_branch st loc queue ((pc, _) as cont) scope_stack ~fall_through : bool * _ =
@@ -2207,7 +2204,7 @@ and compile_closure ctx (pc, args) (cloc : Parse_info.t option) =
   if debug () then Format.eprintf "@[<hv 2>closure {@;";
   let scope_stack = [] in
   let start_loc =
-    let block = Addr.Map.find pc ctx.Ctx.blocks in
+    let block = Code.block pc ctx.Ctx.p in
     match block.body with
     | Event loc :: _ -> J.Pi loc
     | _ -> J.U
@@ -2300,12 +2297,12 @@ let f
       ~mutated_vars
       ~freevars
       ~in_cps
-      p.blocks
+      p
       live_vars
       trampolined_calls
       share
   in
-  let p = compile_program ctx p.start in
+  let p = compile_program ctx (Code.start p) in
   if times () then Format.eprintf "  code gen.: %a@." Timer.print t';
   p
 
diff --git a/compiler/lib/generate_closure.ml b/compiler/lib/generate_closure.ml
index 6771a59473..5f40bd9303 100644
--- a/compiler/lib/generate_closure.ml
+++ b/compiler/lib/generate_closure.ml
@@ -37,12 +37,12 @@ let add_multi k v map =
   let set = try Var.Map.find k map with Not_found -> Addr.Set.empty in
   Var.Map.add k (Addr.Set.add v set) map
 
-let rec collect_apply pc blocks visited tc =
+let rec collect_apply pc p visited tc =
   if Addr.Set.mem pc visited
   then visited, tc
   else
     let visited = Addr.Set.add pc visited in
-    let block = Addr.Map.find pc blocks in
+    let block = Code.block pc p in
     let tc_opt =
       match block.branch with
       | Return x -> (
@@ -57,16 +57,16 @@ let rec collect_apply pc blocks visited tc =
     | Some tc -> visited, tc
     | None ->
         Code.fold_children
-          blocks
+          p
           pc
-          (fun pc (visited, tc) -> collect_apply pc blocks visited tc)
+          (fun pc (visited, tc) -> collect_apply pc p visited tc)
           (visited, tc)
 
-let rec collect_closures blocks l pos =
+let rec collect_closures p l pos =
   match l with
   | Let (f_name, Closure (args, ((pc, _) as cont), cloc)) :: rem ->
-      let _, tc = collect_apply pc blocks Addr.Set.empty Var.Map.empty in
-      let l, rem = collect_closures blocks rem (succ pos) in
+      let _, tc = collect_apply pc p Addr.Set.empty Var.Map.empty in
+      let l, rem = collect_closures p rem (succ pos) in
       { f_name; args; cont; tc; pos; cloc } :: l, rem
   | rem -> [], rem
 
@@ -158,12 +158,12 @@ module Trampoline = struct
 
   let wrapper_closure pc args cloc = Closure (args, (pc, []), cloc)
 
-  let f free_pc blocks closures_map component =
+  let f p closures_map component =
     match component with
     | SCC.No_loop id ->
         let ci = Var.Map.find id closures_map in
         let instr = Let (ci.f_name, Closure (ci.args, ci.cont, ci.cloc)) in
-        free_pc, blocks, [ One { name = ci.f_name; code = instr } ]
+        p, [ One { name = ci.f_name; code = instr } ]
     | SCC.Has_loop all ->
         if debug_tc ()
         then (
@@ -180,20 +180,16 @@ module Trampoline = struct
               ( (if tailcall_max_depth = 0 then None else Some (Code.Var.fresh_n "counter"))
               , Var.Map.find id closures_map ))
         in
-        let blocks, free_pc, closures =
-          List.fold_left
-            all
-            ~init:(blocks, free_pc, [])
-            ~f:(fun (blocks, free_pc, closures) (counter, ci) ->
+        let p, closures =
+          List.fold_left all ~init:(p, []) ~f:(fun (p, closures) (counter, ci) ->
               if debug_tc ()
               then Format.eprintf "Rewriting for %a\n%!" Var.print ci.f_name;
               let new_f = Code.Var.fork ci.f_name in
               let new_args = List.map ci.args ~f:Code.Var.fork in
-              let wrapper_pc = free_pc in
-              let free_pc = free_pc + 1 in
+              let wrapper_pc = Code.free_pc p in
               let new_counter = Option.map counter ~f:Code.Var.fork in
               let start_loc =
-                let block = Addr.Map.find (fst ci.cont) blocks in
+                let block = Code.block (fst ci.cont) p in
                 match block.body with
                 | Event loc :: _ -> loc
                 | _ -> Parse_info.zero
@@ -201,7 +197,7 @@ module Trampoline = struct
               let wrapper_block =
                 wrapper_block new_f ~args:new_args ~counter:new_counter start_loc
               in
-              let blocks = Addr.Map.add wrapper_pc wrapper_block blocks in
+              let p = Code.add_block wrapper_pc wrapper_block p in
               let instr_wrapper =
                 Let (ci.f_name, wrapper_closure wrapper_pc new_args ci.cloc)
               in
@@ -218,30 +214,26 @@ module Trampoline = struct
                       List.map pcs ~f:(fun x -> counter, x) @ acc
                     with Not_found -> acc)
               in
-              let blocks, free_pc =
-                List.fold_left
-                  counter_and_pc
-                  ~init:(blocks, free_pc)
-                  ~f:(fun (blocks, free_pc) (counter, pc) ->
+              let p =
+                List.fold_left counter_and_pc ~init:p ~f:(fun p (counter, pc) ->
                     if debug_tc () then Format.eprintf "Rewriting tc in %d\n%!" pc;
-                    let block = Addr.Map.find pc blocks in
-                    let direct_call_pc = free_pc in
-                    let bounce_call_pc = free_pc + 1 in
-                    let free_pc = free_pc + 2 in
+                    let block = Code.block pc p in
+                    let direct_call_pc = Code.free_pc p in
+                    let bounce_call_pc = direct_call_pc + 1 in
                     match List.rev block.body with
                     | Let (x, Apply { f; args; exact = true }) :: rem_rev ->
                         assert (Var.equal f ci.f_name);
-                        let blocks =
-                          Addr.Map.add
+                        let p =
+                          Code.add_block
                             direct_call_pc
                             (direct_call_block ~counter ~x ~f:new_f ~args)
-                            blocks
+                            p
                         in
-                        let blocks =
-                          Addr.Map.add
+                        let p =
+                          Code.add_block
                             bounce_call_pc
                             (bounce_call_block ~x ~f:new_f ~args)
-                            blocks
+                            p
                         in
                         let block =
                           match counter with
@@ -265,37 +257,30 @@ module Trampoline = struct
                               in
                               { block with body = List.rev (last :: rem_rev); branch }
                         in
-                        let blocks = Addr.Map.remove pc blocks in
-                        Addr.Map.add pc block blocks, free_pc
+                        Code.add_block pc block p
                     | _ -> assert false)
               in
-              ( blocks
-              , free_pc
+              ( p
               , Wrapper { name = ci.f_name; code = instr_real; wrapper = instr_wrapper }
                 :: closures ))
         in
-        free_pc, blocks, closures
+        p, closures
 end
 
-let rec rewrite_closures free_pc blocks body : int * _ * _ list =
+let rec rewrite_closures p body : _ * _ list =
   match body with
   | Let (_, Closure _) :: _ ->
-      let closures, rem = collect_closures blocks body 0 in
+      let closures, rem = collect_closures p body 0 in
       let closures_map =
         List.fold_left closures ~init:Var.Map.empty ~f:(fun closures_map x ->
             Var.Map.add x.f_name x closures_map)
       in
       let components = group_closures closures_map in
-      let free_pc, blocks, closures =
-        List.fold_left
-          components
-          ~init:(free_pc, blocks, [])
-          ~f:(fun (free_pc, blocks, acc) component ->
-            let free_pc, blocks, closures =
-              Trampoline.f free_pc blocks closures_map component
-            in
+      let p, closures =
+        List.fold_left components ~init:(p, []) ~f:(fun (p, acc) component ->
+            let p, closures = Trampoline.f p closures_map component in
             let intrs = closures :: acc in
-            free_pc, blocks, intrs)
+            p, intrs)
       in
       let closures =
         let pos_of_var x = (Var.Map.find x closures_map).pos in
@@ -309,26 +294,25 @@ let rec rewrite_closures free_pc blocks body : int * _ * _ list =
              | One { code; _ } -> [ code ]
              | Wrapper { code; wrapper; _ } -> [ code; wrapper ])
       in
-      let free_pc, blocks, rem = rewrite_closures free_pc blocks rem in
-      free_pc, blocks, closures @ rem
+      let p, rem = rewrite_closures p rem in
+      p, closures @ rem
   | i :: rem ->
-      let free_pc, blocks, rem = rewrite_closures free_pc blocks rem in
-      free_pc, blocks, i :: rem
-  | [] -> free_pc, blocks, []
+      let p, rem = rewrite_closures p rem in
+      p, i :: rem
+  | [] -> p, []
 
 let f p : Code.program =
   Code.invariant p;
-  let blocks, free_pc =
+  let p =
     Addr.Map.fold
-      (fun pc _ (blocks, free_pc) ->
+      (fun pc _ p ->
         (* make sure we have the latest version *)
-        let block = Addr.Map.find pc blocks in
-        let free_pc, blocks, body = rewrite_closures free_pc blocks block.body in
-        Addr.Map.add pc { block with body } blocks, free_pc)
-      p.blocks
-      (p.blocks, p.free_pc)
+        let block = Code.block pc p in
+        let p, body = rewrite_closures p block.body in
+        Code.add_block pc { block with body } p)
+      (Code.blocks p)
+      p
   in
-  let p = { p with blocks; free_pc } in
   Code.invariant p;
   p
 
diff --git a/compiler/lib/global_deadcode.ml b/compiler/lib/global_deadcode.ml
index 991205370a..b3b68d474c 100644
--- a/compiler/lib/global_deadcode.ml
+++ b/compiler/lib/global_deadcode.ml
@@ -99,9 +99,9 @@ let iter_with_scope prog f =
     (fun scope _ (pc, _) _ () ->
       Code.traverse
         { fold = fold_children }
-        (fun pc () -> f scope (Addr.Map.find pc prog.blocks))
+        (fun pc () -> f scope (Code.block pc prog))
         pc
-        prog.blocks
+        prog
         ())
     ()
 
@@ -119,7 +119,7 @@ let definitions prog =
           | Event _ | Set_field (_, _, _, _) | Offset_ref (_, _) | Array_set (_, _, _) ->
               ())
         block.body)
-    prog.blocks;
+    (Code.blocks prog);
   defs
 
 let variable_may_escape x (global_info : Global_flow.info) =
@@ -158,9 +158,9 @@ let usages prog (global_info : Global_flow.info) scoped_live_vars :
     List.iter2 ~f:(fun x y -> add_use (Propagate { scope = []; src = x }) x y) params args
   in
   let add_cont_deps (pc, args) =
-    match try Some (Addr.Map.find pc prog.blocks) with Not_found -> None with
-    | Some block -> add_arg_dep block.params args
-    | None -> () (* Dead continuation *)
+    match Code.block pc prog with
+    | block -> add_arg_dep block.params args
+    | exception Not_found -> () (* Dead continuation *)
   in
   let add_expr_uses scope x e : unit =
     match e with
@@ -503,8 +503,7 @@ let zero prog pure_funs sentinal live_table =
     in
     { block with body; branch }
   in
-  let blocks = prog.blocks |> Addr.Map.map zero_block in
-  { prog with blocks }
+  Code.map_blocks ~f:zero_block prog
 
 module Print = struct
   let rec live_to_string = function
diff --git a/compiler/lib/global_flow.ml b/compiler/lib/global_flow.ml
index 1e167a47ef..e1441dcaec 100644
--- a/compiler/lib/global_flow.ml
+++ b/compiler/lib/global_flow.ml
@@ -58,12 +58,12 @@ let return_values p =
             Code.traverse
               { fold = fold_children }
               (fun pc s ->
-                let block = Addr.Map.find pc p.blocks in
+                let block = Code.block pc p in
                 match block.branch with
                 | Return x -> Var.Set.add x s
                 | _ -> s)
               pc
-              p.blocks
+              p
               Var.Set.empty
           in
           Var.Map.add name s rets)
@@ -156,8 +156,8 @@ let rec arg_deps st ?ignore params args =
   | [], [] -> ()
   | _ -> assert false
 
-let cont_deps blocks st ?ignore (pc, args) =
-  let block = Addr.Map.find pc blocks in
+let cont_deps p st ?ignore (pc, args) =
+  let block = Code.block pc p in
   arg_deps st ?ignore block.params args
 
 let do_escape st level x = st.variable_may_escape.(Var.idx x) <- level
@@ -249,15 +249,15 @@ let expr_deps blocks st x e =
       cont_deps blocks st cont
   | Field (y, _, _) -> add_dep st x y
 
-let program_deps st { start; blocks; _ } =
+let program_deps st p =
   Code.traverse
     { Code.fold = Code.fold_children }
     (fun pc () ->
-      match Addr.Map.find pc blocks with
+      match Code.block pc p with
       | { branch = Return x; _ } -> do_escape st Escape x
       | _ -> ())
-    start
-    blocks
+    (Code.start p)
+    p
     ();
   Addr.Map.iter
     (fun _ block ->
@@ -265,7 +265,7 @@ let program_deps st { start; blocks; _ } =
           match i with
           | Let (x, e) ->
               add_expr_def st x e;
-              expr_deps blocks st x e
+              expr_deps p st x e
           | Assign (x, y) -> add_assign_def st x y
           | Set_field (x, _, _, y) | Array_set (x, _, y) ->
               possibly_mutable st x;
@@ -274,12 +274,12 @@ let program_deps st { start; blocks; _ } =
       match block.branch with
       | Return _ | Stop -> ()
       | Raise (x, _) -> do_escape st Escape x
-      | Branch cont | Poptrap cont -> cont_deps blocks st cont
+      | Branch cont | Poptrap cont -> cont_deps p st cont
       | Cond (x, cont1, cont2) ->
-          cont_deps blocks st cont1;
-          cont_deps blocks st ~ignore:x cont2
+          cont_deps p st cont1;
+          cont_deps p st ~ignore:x cont2
       | Switch (x, a1) -> (
-          Array.iter a1 ~f:(fun cont -> cont_deps blocks st cont);
+          Array.iter a1 ~f:(fun cont -> cont_deps p st cont);
           if not st.fast
           then
             (* looking up the def of x is fine here, because the tag
@@ -296,7 +296,7 @@ let program_deps st { start; blocks; _ } =
                       (i :: (try Addr.Hashtbl.find h pc with Not_found -> [])));
                 Addr.Hashtbl.iter
                   (fun pc tags ->
-                    let block = Addr.Map.find pc blocks in
+                    let block = Code.block pc p in
                     List.iter
                       ~f:(fun i ->
                         match i with
@@ -309,9 +309,9 @@ let program_deps st { start; blocks; _ } =
       | Pushtrap (cont, x, cont_h) ->
           add_var st x;
           st.defs.(Var.idx x) <- Phi { known = Var.Set.empty; others = true };
-          cont_deps blocks st cont_h;
-          cont_deps blocks st cont)
-    blocks
+          cont_deps p st cont_h;
+          cont_deps p st cont)
+    (Code.blocks p)
 
 (* For each variable, we keep track of which values, function or
    block, it may contain. Other kinds of values are not relevant and
diff --git a/compiler/lib/inline.ml b/compiler/lib/inline.ml
index fc14c81c2e..939454d91c 100644
--- a/compiler/lib/inline.ml
+++ b/compiler/lib/inline.ml
@@ -45,7 +45,7 @@ let collect_closures p =
     Code.traverse
       { fold = Code.fold_children }
       (fun pc () ->
-        let block = Addr.Map.find pc p.blocks in
+        let block = Code.block pc p in
         List.iter
           ~f:(fun i ->
             match i with
@@ -55,10 +55,10 @@ let collect_closures p =
             | _ -> ())
           block.body)
       pc
-      p.blocks
+      p
       ()
   in
-  traverse p None p.start;
+  traverse p None (Code.start p);
   closures
 
 let collect_deps p closures =
@@ -74,7 +74,7 @@ let collect_deps p closures =
     Code.traverse
       { fold = Code.fold_children }
       (fun pc () ->
-        let block = Addr.Map.find pc p.blocks in
+        let block = Code.block pc p in
         Freevars.iter_block_free_vars add_dep block;
         List.iter
           ~f:(fun i ->
@@ -83,7 +83,7 @@ let collect_deps p closures =
             | _ -> ())
           block.body)
       pc
-      p.blocks
+      p
       ()
   in
   Var.Hashtbl.iter (fun f (_, (pc, _), _) -> traverse p f pc) closures;
@@ -137,7 +137,7 @@ let visit_closures p ~live_vars f acc =
     ~enclosing_function:None
     ~current_function:None
     ~params:[]
-    ~cont:(p.start, [])
+    ~cont:(Code.start p, [])
     acc
 
 (****)
@@ -149,9 +149,9 @@ let blocks_in_loop p pc =
     Code.traverse
       { fold = Code.fold_children }
       (fun pc g ->
-        Addr.Map.add pc (Code.fold_children p.blocks pc Addr.Set.add Addr.Set.empty) g)
+        Addr.Map.add pc (Code.fold_children p pc Addr.Set.add Addr.Set.empty) g)
       pc
-      p.blocks
+      p
       Addr.Map.empty
   in
   let scc = SCC.component_graph g in
@@ -218,23 +218,18 @@ let contains_loop ~context info =
         then visited, Addr.Map.find pc visited
         else
           let visited, loop =
-            Code.fold_children
-              context.p.blocks
-              pc
-              traverse
-              (Addr.Map.add pc true visited, false)
+            Code.fold_children context.p pc traverse (Addr.Map.add pc true visited, false)
           in
           Addr.Map.add pc false visited, loop
       in
       snd (traverse pc (Addr.Map.empty, false)))
 
 let sum ~context f pc =
-  let blocks = context.p.blocks in
   Code.traverse
     { fold = fold_children }
-    (fun pc acc -> f (Addr.Map.find pc blocks) + acc)
+    (fun pc acc -> f (Code.block pc context.p) + acc)
     pc
-    blocks
+    context.p
     0
 
 let rec block_size ~recurse ~context { branch; body; _ } =
@@ -294,13 +289,12 @@ let count_init_code ~context info =
 (** Whether the function returns a block. *)
 let returns_a_block ~context info =
   cache ~info info.returns_a_block (fun pc ->
-      let blocks = context.p.blocks in
       Code.traverse
         { fold = fold_children }
         (fun pc acc ->
           acc
           &&
-          let block = Addr.Map.find pc blocks in
+          let block = Code.block pc context.p in
           match block.branch with
           | Return x -> (
               match Code.last_instr block.body with
@@ -308,7 +302,7 @@ let returns_a_block ~context info =
               | _ -> false)
           | Raise _ | Stop | Branch _ | Cond _ | Switch _ | Pushtrap _ | Poptrap _ -> true)
         pc
-        blocks
+        context.p
         true)
 
 (** List of parameters that corresponds to functions called once in
@@ -320,11 +314,10 @@ let interesting_parameters ~context info =
       if List.is_empty params
       then []
       else
-        let blocks = context.p.blocks in
         Code.traverse
           { fold = fold_children }
           (fun pc lst ->
-            let block = Addr.Map.find pc blocks in
+            let block = Code.block pc context.p in
             List.fold_left
               ~f:(fun lst i ->
                 match i with
@@ -334,7 +327,7 @@ let interesting_parameters ~context info =
               ~init:lst
               block.body)
           pc
-          blocks
+          context.p
           [])
 
 (*
@@ -480,77 +473,66 @@ let remove_dead_closures_from_block ~live_vars p pc block =
   in
   if List.exists ~f:is_dead_closure block.body
   then
-    { p with
-      blocks =
-        Addr.Map.add
-          pc
-          { block with
-            body =
-              List.fold_left block.body ~init:[] ~f:(fun acc i ->
-                  match i, acc with
-                  | Event _, Event _ :: prev ->
-                      (* Avoid consecutive events (keep just the last one) *)
-                      i :: prev
-                  | _ -> if is_dead_closure i then acc else i :: acc)
-              |> List.rev
-          }
-          p.blocks
-    }
+    Code.add_block
+      pc
+      { block with
+        body =
+          List.fold_left block.body ~init:[] ~f:(fun acc i ->
+              match i, acc with
+              | Event _, Event _ :: prev ->
+                  (* Avoid consecutive events (keep just the last one) *)
+                  i :: prev
+              | _ -> if is_dead_closure i then acc else i :: acc)
+          |> List.rev
+      }
+      p
   else p
 
 let remove_dead_closures ~live_vars p pc =
   Code.traverse
     { fold = fold_children }
     (fun pc p ->
-      let block = Addr.Map.find pc p.blocks in
+      let block = Code.block pc p in
       remove_dead_closures_from_block ~live_vars p pc block)
     pc
-    p.blocks
+    p
     p
 
 (****)
 
-let rewrite_block pc' pc blocks =
-  let block = Addr.Map.find pc blocks in
-  let block =
-    match block.branch, pc' with
-    | Return y, Some pc' -> { block with branch = Branch (pc', [ y ]) }
-    | _ -> block
-  in
-  Addr.Map.add pc block blocks
+let rewrite_block pc' pc p =
+  let block = Code.block pc p in
+  match block.branch, pc' with
+  | Return y, Some pc' -> Code.add_block pc { block with branch = Branch (pc', [ y ]) } p
+  | _ -> p
 
-let rewrite_closure blocks cont_pc clos_pc =
+let rewrite_closure p cont_pc clos_pc =
   Code.traverse
     { fold = Code.fold_children_skip_try_body }
     (rewrite_block cont_pc)
     clos_pc
-    blocks
-    blocks
+    p
+    p
 
 let rewrite_inlined_function p rem branch x params cont args =
-  let blocks, cont_pc, free_pc =
+  let p, cont_pc =
     match rem, branch with
     | [], Return y when Var.equal x y ->
         (* We do not need a continuation block for tail calls *)
-        p.blocks, None, p.free_pc
+        p, None
     | _ ->
-        let fresh_addr = p.free_pc in
-        let free_pc = fresh_addr + 1 in
-        ( Addr.Map.add fresh_addr { params = [ x ]; body = rem; branch } p.blocks
-        , Some fresh_addr
-        , free_pc )
+        let fresh_addr = Code.free_pc p in
+        ( Code.add_block fresh_addr { params = [ x ]; body = rem; branch } p
+        , Some fresh_addr )
   in
-  let blocks = rewrite_closure blocks cont_pc (fst cont) in
+  let p = rewrite_closure p cont_pc (fst cont) in
   (* We do not really need this intermediate block.
      It just avoids the need to find which function
      parameters are used in the function body. *)
-  let fresh_addr = free_pc in
-  let free_pc = fresh_addr + 1 in
+  let fresh_addr = Code.free_pc p in
   assert (List.compare_lengths args params = 0);
-  let blocks =
-    Addr.Map.add fresh_addr { params; body = []; branch = Branch cont } blocks
-  in
-  [], (Branch (fresh_addr, args), { p with blocks; free_pc })
+  let p = Code.add_block fresh_addr { params; body = []; branch = Branch cont } p in
+  [], (Branch (fresh_addr, args), p)
 
 let rec inline_recursively ~context ~info p params (pc, _) args =
   let relevant_args = relevant_arguments ~context info args in
@@ -570,7 +552,7 @@ let rec inline_recursively ~context ~info p params (pc, _) args =
     Code.traverse
       { fold = Code.fold_children }
       (fun pc p ->
-        let block = Addr.Map.find pc p.blocks in
+        let block = Code.block pc p in
         let body, (branch, p) =
           List.fold_right
             ~f:(fun i (rem, state) ->
@@ -590,9 +572,9 @@ let rec inline_recursively ~context ~info p params (pc, _) args =
             ~init:([], (block.branch, p))
             block.body
         in
-        { p with blocks = Addr.Map.add pc { block with body; branch } p.blocks })
+        Code.add_block pc { block with body; branch } p)
       pc
-      p.blocks
+      p
       p
 
 and inline_function ~context i x f args rem state =
@@ -627,7 +609,7 @@ let inline_in_block ~context pc block p =
       ~init:([], (block.branch, p))
       block.body
   in
-  { p with blocks = Addr.Map.add pc { block with body; branch } p.blocks }
+  Code.add_block pc { block with body; branch } p
 
 let inline ~profile ~inline_count p ~live_vars =
   if debug () then Format.eprintf "====== inlining ======@.";
@@ -651,7 +633,7 @@ let inline ~profile ~inline_count p ~live_vars =
          Code.traverse
            { fold = Code.fold_children }
            (fun pc p ->
-             let block = Addr.Map.find pc p.blocks in
+             let block = Code.block pc p in
              if
                (* Skip blocks with no call of known function *)
                List.for_all
@@ -669,7 +651,7 @@ let inline ~profile ~inline_count p ~live_vars =
                  block
                  p)
            pc
-           p.blocks
+           p
            p
        in
        let p = remove_dead_closures ~live_vars p pc in
diff --git a/compiler/lib/lambda_lifting.ml b/compiler/lib/lambda_lifting.ml
index 945c1e7512..6f47ddd1b0 100644
--- a/compiler/lib/lambda_lifting.ml
+++ b/compiler/lib/lambda_lifting.ml
@@ -73,7 +73,7 @@ let rec compute_depth program pc =
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc d ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       List.fold_left block.body ~init:d ~f:(fun d i ->
           match i with
           | Let (_, Closure (_, (pc', _), _)) ->
@@ -81,7 +81,7 @@ let rec compute_depth program pc =
               max d (d' + 1)
           | _ -> d))
     pc
-    program.blocks
+    program
     0
 
 let collect_free_vars program var_depth depth pc =
@@ -91,7 +91,7 @@ let collect_free_vars program var_depth depth pc =
     Code.preorder_traverse
       { fold = Code.fold_children }
       (fun pc () ->
-        let block = Code.Addr.Map.find pc program.blocks in
+        let block = Code.block pc program in
         Freevars.iter_block_free_vars
           (fun x ->
             let idx = Var.idx x in
@@ -106,7 +106,7 @@ let collect_free_vars program var_depth depth pc =
             | Let (_, Closure (_, (pc', _), _)) -> traverse pc'
             | _ -> ()))
       pc
-      program.blocks
+      program
       ()
   in
   traverse pc;
@@ -125,7 +125,7 @@ let rec traverse var_depth (program, functions) pc depth limit =
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc (program, functions) ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       mark_bound_variables var_depth block depth;
       if depth = baseline
       then (
@@ -140,7 +140,7 @@ let rec traverse var_depth (program, functions) pc depth limit =
                   program, List.rev_append functions (i :: rem)
               | i -> program, i :: rem)
         in
-        { program with blocks = Addr.Map.add pc { block with body } program.blocks }, [])
+        Code.add_block pc { block with body } program, [])
       else if depth < limit
       then
         List.fold_left block.body ~init:(program, functions) ~f:(fun st i ->
@@ -189,14 +189,9 @@ let rec traverse var_depth (program, functions) pc depth limit =
                     depth
                     (Var.Set.cardinal free_vars)
                     (compute_depth program pc');
-                let pc'' = program.free_pc in
+                let pc'' = Code.free_pc program in
                 let bl = { params = []; body = [ Let (f', cl) ]; branch = Return f' } in
-                let program =
-                  { program with
-                    free_pc = pc'' + 1
-                  ; blocks = Addr.Map.add pc'' bl program.blocks
-                  }
-                in
+                let program = Code.add_block pc'' bl program in
                 let functions =
                   Let (f'', Closure (List.map s ~f:snd, (pc'', []), None)) :: functions
                 in
@@ -219,10 +214,9 @@ let rec traverse var_depth (program, functions) pc depth limit =
         let body, (program, functions) =
           rewrite_body true (program, functions) block.body
         in
-        ( { program with blocks = Addr.Map.add pc { block with body } program.blocks }
-        , functions ))
+        Code.add_block pc { block with body } program, functions)
     pc
-    program.blocks
+    program
     (program, functions)
 
 let f p =
@@ -232,7 +226,7 @@ let f p =
   let p, functions =
     let threshold = Config.Param.lambda_lifting_threshold () in
     let baseline = Config.Param.lambda_lifting_baseline () in
-    traverse var_depth (p, []) p.start 0 (baseline + threshold)
+    traverse var_depth (p, []) (Code.start p) 0 (baseline + threshold)
   in
   assert (List.is_empty functions);
   if Debug.find "times" () then Format.eprintf "  lambda lifting: %a@." Timer.print t;
diff --git a/compiler/lib/lambda_lifting_simple.ml b/compiler/lib/lambda_lifting_simple.ml
index 74f3be6b9b..af273c9aee 100644
--- a/compiler/lib/lambda_lifting_simple.ml
+++ b/compiler/lib/lambda_lifting_simple.ml
@@ -27,7 +27,7 @@ let rec compute_depth program pc =
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc d ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       List.fold_left block.body ~init:d ~f:(fun d i ->
           match i with
           | Let (_, Closure (_, (pc', _), _)) ->
@@ -35,7 +35,7 @@ let rec compute_depth program pc =
               max d (d' + 1)
           | _ -> d))
     pc
-    program.blocks
+    program
     0
 
 let collect_free_vars program var_depth depth pc =
@@ -44,7 +44,7 @@ let collect_free_vars program var_depth depth pc =
     Code.preorder_traverse
       { fold = Code.fold_children }
       (fun pc () ->
-        let block = Code.Addr.Map.find pc program.blocks in
+        let block = Code.block pc program in
         Freevars.iter_block_free_vars
           (fun x ->
             let idx = Var.idx x in
@@ -59,7 +59,7 @@ let collect_free_vars program var_depth depth pc =
             | Let (_, Closure (_, (pc', _), _)) -> traverse pc'
             | _ -> ()))
       pc
-      program.blocks
+      program
       ()
   in
   traverse pc;
@@ -90,7 +90,7 @@ let rec rewrite_blocks
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc (program, functions, lifters) ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       mark_bound_variables var_depth block depth;
       let body, (program, functions, lifters) =
         rewrite_body
@@ -103,11 +103,9 @@ let rec rewrite_blocks
           ~acc_instr:[]
           block.body
       in
-      ( { program with blocks = Addr.Map.add pc { block with body } program.blocks }
-      , functions
-      , lifters ))
+      Code.add_block pc { block with body } program, functions, lifters)
     pc
-    program.blocks
+    program
     (program, functions, lifters)
 
 and rewrite_body
@@ -159,11 +157,9 @@ and rewrite_body
           depth
           (Var.Set.cardinal free_vars)
           (compute_depth program pc');
-      let pc'' = program.free_pc in
+      let pc'' = Code.free_pc program in
       let bl = { params = []; body = [ Let (f', cl) ]; branch = Return f' } in
-      let program =
-        { program with free_pc = pc'' + 1; blocks = Addr.Map.add pc'' bl program.blocks }
-      in
+      let program = Code.add_block pc'' bl program in
       (* Add to returned list of lifter functions definitions *)
       let functions =
         Let (f'', Closure (List.map s ~f:snd, (pc'', []), None)) :: functions
@@ -259,7 +255,7 @@ and rewrite_body
                    f_tuple
                    depth
                    (Var.Set.cardinal free_vars)));
-            let pc_tuple = program.free_pc in
+            let pc_tuple = Code.free_pc program in
             let lifted_block =
               let tuple = Var.fresh_n "tuple" in
               { params = []
@@ -273,12 +269,7 @@ and rewrite_body
               ; branch = Return tuple
               }
             in
-            let program =
-              { program with
-                free_pc = pc_tuple + 1
-              ; blocks = Addr.Map.add pc_tuple lifted_block program.blocks
-              }
-            in
+            let program = Code.add_block pc_tuple lifted_block program in
             let functions =
               Let (f_tuple, Closure (List.map s ~f:snd, (pc_tuple, []), None))
               :: functions
@@ -328,7 +319,7 @@ let lift ~to_lift ~pc program : program * Var.t Var.Map.t =
   Code.preorder_traverse
     { fold = Code.fold_children }
     (fun pc (program, lifter_map) ->
-      let block = Code.Addr.Map.find pc program.blocks in
+      let block = Code.block pc program in
       mark_bound_variables var_depth block 0;
       let program, body, lifter_map' =
         List.fold_right
@@ -349,10 +340,10 @@ let lift ~to_lift ~pc program : program * Var.t Var.Map.t =
                 program, List.rev_append functions (i :: rem), lifters
             | i -> program, i :: rem, lifters)
       in
-      ( { program with blocks = Addr.Map.add pc { block with body } program.blocks }
+      ( Code.add_block pc { block with body } program
       , Var.Map.union (fun _ _ -> assert false) lifter_map lifter_map' ))
     pc
-    program.blocks
+    program
     (program, Var.Map.empty)
 
 let f ~to_lift program =
@@ -362,6 +353,6 @@ let f ~to_lift program =
     Code.Print.program Format.err_formatter (fun _ _ -> "") program;
     Format.eprintf "@]");
   let t = Timer.make () in
-  let program, liftings = lift ~to_lift ~pc:program.start program in
+  let program, liftings = lift ~to_lift ~pc:(Code.start program) program in
   if Debug.find "times" () then Format.eprintf "  lambda lifting: %a@." Timer.print t;
   program, liftings
diff --git a/compiler/lib/parse_bytecode.ml b/compiler/lib/parse_bytecode.ml
index 0ce79a4f6e..9a22795a72 100644
--- a/compiler/lib/parse_bytecode.ml
+++ b/compiler/lib/parse_bytecode.ml
@@ -2531,8 +2531,7 @@ let parse_bytecode code globals debug_data =
             })
           !compiled_blocks
       in
-      let free_pc = (Addr.Map.max_binding blocks |> fst) + 1 in
-      { start; blocks; free_pc })
+      Code.program start blocks)
     else Code.empty
   in
   compiled_blocks := Addr.Map.empty;
@@ -3130,7 +3129,7 @@ let predefined_exceptions () =
     ; aliases = []
     }
   in
-  { start = 0; blocks = Addr.Map.singleton 0 block; free_pc = 1 }, unit_info
+  Code.program 0 (Addr.Map.singleton 0 block), unit_info
 
 let link_info ~symbols ~primitives ~crcs =
   let gdata = Code.Var.fresh_n "global_data" in
@@ -3172,4 +3171,4 @@ let link_info ~symbols ~primitives ~crcs =
     Let (gdata, Prim (Extern "caml_get_global_data", [])) :: body
   in
   let block = { params = []; body; branch = Stop } in
-  { start = 0; blocks = Addr.Map.singleton 0 block; free_pc = 1 }
+  Code.program 0 (Addr.Map.singleton 0 block)
diff --git a/compiler/lib/partial_cps_analysis.ml b/compiler/lib/partial_cps_analysis.ml
index df7eeb8fcd..ae1e55fe77 100644
--- a/compiler/lib/partial_cps_analysis.ml
+++ b/compiler/lib/partial_cps_analysis.ml
@@ -113,9 +113,9 @@ let program_deps ~info ~vars ~tail_deps ~deps p =
       traverse
         { fold = Code.fold_children }
         (fun pc () ->
-          block_deps ~info ~vars ~tail_deps ~deps ~blocks:p.blocks ~fun_name pc)
+          block_deps ~info ~vars ~tail_deps ~deps ~blocks:(Code.blocks p) ~fun_name pc)
         pc
-        p.blocks
+        p
         ())
     ()
 
diff --git a/compiler/lib/phisimpl.ml b/compiler/lib/phisimpl.ml
index 5f633db638..051efa703e 100644
--- a/compiler/lib/phisimpl.ml
+++ b/compiler/lib/phisimpl.ml
@@ -60,7 +60,7 @@ let expr_deps blocks vars deps defs x e =
   | Block (_, a, _, _) -> Array.iter a ~f:(fun y -> add_dep deps x y)
   | Field (y, _, _) -> add_dep deps x y
 
-let program_deps { blocks; _ } =
+let program_deps blocks =
   let nv = Var.count () in
   let vars = Var.ISet.empty () in
   let deps = Array.make nv Var.Set.empty in
@@ -157,7 +157,8 @@ let f p =
   Code.invariant p;
   let t = Timer.make () in
   let t' = Timer.make () in
-  let vars, deps, defs = program_deps p in
+  let blocks = Code.blocks p in
+  let vars, deps, defs = program_deps blocks in
   if times () then Format.eprintf "    phi-simpl. 1: %a@." Timer.print t';
   let t' = Timer.make () in
   let subst = solver1 vars deps defs in
diff --git a/compiler/lib/pure_fun.ml b/compiler/lib/pure_fun.ml
index da4a9c59aa..5a73b5b188 100644
--- a/compiler/lib/pure_fun.ml
+++ b/compiler/lib/pure_fun.ml
@@ -45,26 +45,26 @@ let pure_instr pure_funs i =
 
 (****)
 
-let rec traverse blocks pc visited pure_blocks funs =
+let rec traverse p pc visited pure_blocks funs =
   if BitSet.mem visited pc
   then BitSet.mem pure_blocks pc
   else (
     BitSet.set visited pc;
-    let pure = block blocks pc visited pure_blocks funs in
+    let pure = block p pc visited pure_blocks funs in
     let pure =
       fold_children
-        blocks
+        p
         pc
         (fun pc pure ->
-          let pure' = traverse blocks pc visited pure_blocks funs in
+          let pure' = traverse p pc visited pure_blocks funs in
           pure && pure')
         pure
     in
     if pure then BitSet.set pure_blocks pc;
     pure)
 
-and block blocks pc visited pure_blocks funs =
-  let b = Addr.Map.find pc blocks in
+and block p pc visited pure_blocks funs =
+  let b = Code.block pc p in
   let pure =
     match b.branch with
     | Raise _ -> false
@@ -73,7 +73,7 @@ and block blocks pc visited pure_blocks funs =
   List.fold_left b.body ~init:pure ~f:(fun pure i ->
       (match i with
       | Let (x, Closure (_, (pc, _), _)) ->
-          let pure = traverse blocks pc visited pure_blocks funs in
+          let pure = traverse p pc visited pure_blocks funs in
           if pure then funs := Var.Set.add x !funs
       | _ -> ());
       pure && pure_instr !funs i)
@@ -82,10 +82,10 @@ type t = Var.Set.t
 
 let f p =
   let t = Timer.make () in
-  let visited = BitSet.create' p.free_pc in
-  let pure = BitSet.create' p.free_pc in
+  let visited = BitSet.create' (Code.free_pc p) in
+  let pure = BitSet.create' (Code.free_pc p) in
   let funs = ref Var.Set.empty in
-  let _ = traverse p.blocks p.start visited pure funs in
+  let _ = traverse p (Code.start p) visited pure funs in
   if times () then Format.eprintf "    pure funs.: %a@." Timer.print t;
   if stats () then Format.eprintf "Stats - pure functions: %d@." (Var.Set.cardinal !funs);
   !funs
diff --git a/compiler/lib/specialize.ml b/compiler/lib/specialize.ml
index 769e510c41..c4115a1123 100644
--- a/compiler/lib/specialize.ml
+++ b/compiler/lib/specialize.ml
@@ -64,15 +64,15 @@ let unknown_apply = function
   | Let (_, Apply { f = _; args = _; exact = false }) -> true
   | _ -> false
 
-let specialize_apply opt_count function_arity ((acc, free_pc, extra), loc) i =
+let specialize_apply opt_count function_arity ((acc, p), loc) i =
   match i with
   | Let (x, Apply { f; args; exact = false }) -> (
       let n' = List.length args in
       match function_arity f with
-      | None -> i :: acc, free_pc, extra
+      | None -> i :: acc, p
       | Some n when n = n' ->
           incr opt_count;
-          Let (x, Apply { f; args; exact = true }) :: acc, free_pc, extra
+          Let (x, Apply { f; args; exact = true }) :: acc, p
       | Some n when n < n' ->
           incr opt_count;
           let v = Code.Var.fresh () in
@@ -80,8 +80,7 @@ let specialize_apply opt_count function_arity ((acc, free_pc, extra), loc) i =
           ( (* Reversed *)
             Let (x, Apply { f = v; args = rest; exact = false })
             :: add_event loc (Let (v, Apply { f; args; exact = true }) :: acc)
-          , free_pc
-          , extra )
+          , p )
       | Some n when n > n' ->
           incr opt_count;
           let missing = Array.init (n - n') ~f:(fun _ -> Code.Var.fresh ()) in
@@ -98,39 +97,32 @@ let specialize_apply opt_count function_arity ((acc, free_pc, extra), loc) i =
             ; branch = Return return'
             }
           in
+          let free_pc = Code.free_pc p in
           ( Let (x, Closure (missing, (free_pc, missing), None)) :: acc
-          , free_pc + 1
-          , (free_pc, block) :: extra )
+          , Code.add_block free_pc block p )
       | Some _ -> assert false)
-  | _ -> i :: acc, free_pc, extra
+  | _ -> i :: acc, p
 
 let specialize_instrs ~function_arity opt_count p =
-  let blocks, free_pc =
-    Addr.Map.fold
-      (fun pc block (blocks, free_pc) ->
-        if List.exists ~f:unknown_apply block.body
-        then
-          let (body, free_pc, extra), _ =
-            List.fold_left
-              block.body
-              ~init:(([], free_pc, []), None)
-              ~f:(fun acc i ->
-                match i with
-                | Event loc ->
-                    let (body, free_pc, extra), _ = acc in
-                    (i :: body, free_pc, extra), Some loc
-                | _ -> specialize_apply opt_count function_arity acc i, None)
-          in
-          let blocks =
-            List.fold_left extra ~init:blocks ~f:(fun blocks (pc, b) ->
-                Addr.Map.add pc b blocks)
-          in
-          Addr.Map.add pc { block with Code.body = List.rev body } blocks, free_pc
-        else blocks, free_pc)
-      p.blocks
-      (p.blocks, p.free_pc)
-  in
-  { p with blocks; free_pc }
+  Addr.Map.fold
+    (fun pc block p ->
+      if List.exists ~f:unknown_apply block.body
+      then
+        let (body, p), _ =
+          List.fold_left
+            block.body
+            ~init:(([], p), None)
+            ~f:(fun acc i ->
+              match i with
+              | Event loc ->
+                  let (body, p), _ = acc in
+                  (i :: body, p), Some loc
+              | _ -> specialize_apply opt_count function_arity acc i, None)
+        in
+        Code.add_block pc { block with Code.body = List.rev body } p
+      else p)
+    (Code.blocks p)
+    p
 
 let f ~function_arity p =
   Code.invariant p;
@@ -183,74 +175,68 @@ let switches p =
   let t = Timer.make () in
   let opt_count = ref 0 in
   let p =
-    { p with
-      blocks =
-        Addr.Map.fold
-          (fun pc block blocks ->
-            match block.branch with
-            | Switch (x, l) -> (
-                match find_outlier_index l with
-                | `All_equals ->
-                    incr opt_count;
-                    Addr.Map.add pc { block with branch = Branch l.(0) } blocks
-                | `Distinguished i ->
-                    incr opt_count;
-                    let block =
-                      let c = Var.fresh () in
-                      { block with
-                        body =
-                          block.body
-                          @ [ Let
-                                (c, Prim (Eq, [ Pc (Int (Targetint.of_int_exn i)); Pv x ]))
-                            ]
-                      ; branch = Cond (c, l.(i), l.((i + 1) mod Array.length l))
-                      }
-                    in
-                    Addr.Map.add pc block blocks
-                | `Splitted i ->
-                    incr opt_count;
-                    let block =
-                      let c = Var.fresh () in
-                      { block with
-                        body =
-                          block.body
-                          @ [ Let
-                                (c, Prim (Lt, [ Pv x; Pc (Int (Targetint.of_int_exn i)) ]))
-                            ]
-                      ; branch = Cond (c, l.(i - 1), l.(i))
-                      }
-                    in
-                    Addr.Map.add pc block blocks
-                | `Splitted_shifted (i, j) ->
-                    incr opt_count;
-                    let block =
-                      let shifted = Var.fresh () in
-                      let c = Var.fresh () in
-                      { block with
-                        body =
-                          block.body
-                          @ [ Let
-                                ( shifted
-                                , Prim
-                                    ( Extern "%int_sub"
-                                    , [ Pv x; Pc (Int (Targetint.of_int_exn i)) ] ) )
-                            ; Let
-                                ( c
-                                , Prim
-                                    ( Ult
-                                    , [ Pv shifted
-                                      ; Pc (Int (Targetint.of_int_exn (j - i)))
-                                      ] ) )
-                            ]
-                      ; branch = Cond (c, l.(i), l.(j))
-                      }
-                    in
-                    Addr.Map.add pc block blocks
-                | `Many_cases -> blocks)
-            | _ -> blocks)
-          p.blocks
-          p.blocks
-    }
+    Addr.Map.fold
+      (fun pc block p ->
+        match block.branch with
+        | Switch (x, l) -> (
+            match find_outlier_index l with
+            | `All_equals ->
+                incr opt_count;
+                Code.add_block pc { block with branch = Branch l.(0) } p
+            | `Distinguished i ->
+                incr opt_count;
+                let block =
+                  let c = Var.fresh () in
+                  { block with
+                    body =
+                      block.body
+                      @ [ Let (c, Prim (Eq, [ Pc (Int (Targetint.of_int_exn i)); Pv x ]))
+                        ]
+                  ; branch = Cond (c, l.(i), l.((i + 1) mod Array.length l))
+                  }
+                in
+                Code.add_block pc block p
+            | `Splitted i ->
+                incr opt_count;
+                let block =
+                  let c = Var.fresh () in
+                  { block with
+                    body =
+                      block.body
+                      @ [ Let (c, Prim (Lt, [ Pv x; Pc (Int (Targetint.of_int_exn i)) ]))
+                        ]
+                  ; branch = Cond (c, l.(i - 1), l.(i))
+                  }
+                in
+                Code.add_block pc block p
+            | `Splitted_shifted (i, j) ->
+                incr opt_count;
+                let block =
+                  let shifted = Var.fresh () in
+                  let c = Var.fresh () in
+                  { block with
+                    body =
+                      block.body
+                      @ [ Let
+                            ( shifted
+                            , Prim
+                                ( Extern "%int_sub"
+                                , [ Pv x; Pc (Int (Targetint.of_int_exn i)) ] ) )
+                        ; Let
+                            ( c
+                            , Prim
+                                ( Ult
+                                , [ Pv shifted; Pc (Int (Targetint.of_int_exn (j - i))) ]
+                                ) )
+                        ]
+                  ; branch = Cond (c, l.(i), l.(j))
+                  }
+                in
+                Code.add_block pc block p
+            | `Many_cases -> p)
+        | _ -> p)
+      (Code.blocks p)
+      p
   in
   if times () then Format.eprintf "  switches: %a@." Timer.print t;
   if stats () then Format.eprintf "Stats - switches: %d@." !opt_count;
diff --git a/compiler/lib/specialize_js.ml b/compiler/lib/specialize_js.ml
index dda64b1063..adf5ff70eb 100644
--- a/compiler/lib/specialize_js.ml
+++ b/compiler/lib/specialize_js.ml
@@ -392,20 +392,15 @@ let specialize_instrs ~target opt_count info l =
   aux info [] l []
 
 let specialize_all_instrs ~target opt_count info p =
-  let blocks =
-    Addr.Map.map
-      (fun block ->
-        { block with
-          Code.body =
-            specialize_instrs
-              ~target
-              opt_count
-              info
-              (specialize_string_concat opt_count block.body)
-        })
-      p.blocks
-  in
-  { p with blocks }
+  Code.map_blocks p ~f:(fun block ->
+      { block with
+        Code.body =
+          specialize_instrs
+            ~target
+            opt_count
+            info
+            (specialize_string_concat opt_count block.body)
+      })
 
 (****)
 
@@ -445,10 +440,9 @@ let f_once_before p =
             loop acc r
         | _ -> loop (i :: acc) r)
   in
-  let blocks =
-    Addr.Map.map (fun block -> { block with Code.body = loop [] block.body }) p.blocks
+  let p =
+    Code.map_blocks p ~f:(fun block -> { block with Code.body = loop [] block.body })
   in
-  let p = { p with blocks } in
   Code.invariant p;
   p
 
@@ -467,7 +461,7 @@ let f_once_after p =
   in
   let f = function
     | Let (x, Closure (l, (pc, []), _)) as i -> (
-        let block = Addr.Map.find pc p.blocks in
+        let block = Code.block pc p in
         match block with
         | { body =
               ( [ Let (y, Prim (Extern prim, args)) ]
@@ -488,12 +482,11 @@ let f_once_after p =
   in
   if first_class_primitives
   then (
-    let blocks =
-      Addr.Map.map
-        (fun block -> { block with Code.body = List.map block.body ~f })
-        p.blocks
+    let p =
+      Code.map_blocks p ~f:(fun block ->
+          { block with Code.body = List.map block.body ~f })
     in
-    let p = Deadcode.remove_unused_blocks { p with blocks } in
+    let p = Deadcode.remove_unused_blocks p in
     Code.invariant p;
     p)
   else p
diff --git a/compiler/lib/structure.ml b/compiler/lib/structure.ml
index 403d758f67..86db477fcc 100644
--- a/compiler/lib/structure.ml
+++ b/compiler/lib/structure.ml
@@ -67,17 +67,17 @@ let empty_body body =
       | _ -> false)
     body
 
-let rec leave_try_body block_order preds blocks pc =
+let rec leave_try_body block_order preds p pc =
   if is_merge_node' block_order preds pc
   then false
   else
-    match Addr.Map.find pc blocks with
+    match Code.block pc p with
     | { body; branch = Return _ | Stop; _ } when empty_body body -> false
     | { body; branch = Branch (pc', _); _ } when empty_body body ->
-        leave_try_body block_order preds blocks pc'
+        leave_try_body block_order preds p pc'
     | _ -> true
 
-let build_graph blocks pc =
+let build_graph p pc =
   let succs = Addr.Hashtbl.create 16 in
   let l = ref [] in
   let visited = Addr.Hashtbl.create 16 in
@@ -86,9 +86,9 @@ let build_graph blocks pc =
     if not (Addr.Hashtbl.mem visited pc)
     then (
       Addr.Hashtbl.add visited pc ();
-      let successors = Code.fold_children blocks pc Addr.Set.add Addr.Set.empty in
+      let successors = Code.fold_children p pc Addr.Set.add Addr.Set.empty in
       Addr.Hashtbl.add succs pc successors;
-      let block = Addr.Map.find pc blocks in
+      let block = Code.block pc p in
       Addr.Set.iter
         (fun pc' ->
           let englobing_exn_handlers =
@@ -112,7 +112,7 @@ let build_graph blocks pc =
   List.iteri !l ~f:(fun i pc -> Addr.Hashtbl.add block_order pc i);
   let preds = reverse_graph succs in
   List.iter !poptraps ~f:(fun (enter_pc, leave_pc) ->
-      if leave_try_body block_order preds blocks leave_pc
+      if leave_try_body block_order preds p leave_pc
       then (
         (* Add an edge to limit the [try] body *)
         Addr.Hashtbl.replace
@@ -210,11 +210,11 @@ let mark_loops g =
     g.preds;
   in_loop
 
-let rec measure blocks g pc limit =
+let rec measure p g pc limit =
   if is_loop_header g pc
   then -1
   else
-    let b = Addr.Map.find pc blocks in
+    let b = Code.block pc p in
     let limit =
       List.fold_left b.body ~init:limit ~f:(fun acc x ->
           match x with
@@ -227,13 +227,13 @@ let rec measure blocks g pc limit =
     then limit
     else
       Addr.Set.fold
-        (fun pc limit -> if limit < 0 then limit else measure blocks g pc limit)
+        (fun pc limit -> if limit < 0 then limit else measure p g pc limit)
         (get_edges g.succs pc)
         limit
 
-let is_small blocks g pc = measure blocks g pc 20 >= 0
+let is_small p g pc = measure p g pc 20 >= 0
 
-let shrink_loops blocks ({ succs; preds; reverse_post_order; _ } as g) =
+let shrink_loops p ({ succs; preds; reverse_post_order; _ } as g) =
   let add_edge pred succ =
     Addr.Hashtbl.replace succs pred (Addr.Set.add succ (Addr.Hashtbl.find succs pred));
     Addr.Hashtbl.replace preds succ (Addr.Set.add pred (Addr.Hashtbl.find preds succ))
@@ -244,7 +244,7 @@ let shrink_loops blocks ({ succs; preds; reverse_post_order; _ } as g) =
   let rec traverse ignored pc =
     let succs = get_edges dom pc in
     let loops = get_edges in_loop pc in
-    let block = Addr.Map.find pc blocks in
+    let block = Code.block pc p in
     Addr.Set.iter
       (fun pc' ->
         (* Whatever is in the scope of an exception handler should not be
@@ -260,7 +260,7 @@ let shrink_loops blocks ({ succs; preds; reverse_post_order; _ } as g) =
         (* If we leave a loop, we add an edge from predecessors of
            the loop header to the current block, so that it is
            considered outside of the loop. *)
-        if not (Addr.Set.is_empty left_loops || is_small blocks g pc')
+        if not (Addr.Set.is_empty left_loops || is_small p g pc')
         then
           Addr.Set.iter
             (fun pc0 ->
@@ -273,7 +273,7 @@ let shrink_loops blocks ({ succs; preds; reverse_post_order; _ } as g) =
   in
   traverse Addr.Set.empty root
 
-let build_graph blocks pc =
-  let g = build_graph blocks pc in
-  shrink_loops blocks g;
+let build_graph p pc =
+  let g = build_graph p pc in
+  shrink_loops p g;
   g
diff --git a/compiler/lib/structure.mli b/compiler/lib/structure.mli
index 1aa1a10940..d9fc57f8e6 100644
--- a/compiler/lib/structure.mli
+++ b/compiler/lib/structure.mli
@@ -11,7 +11,7 @@ val is_backward : t -> Addr.t -> Addr.t -> bool
 
 val is_forward : t -> Addr.t -> Addr.t -> bool
 
-val build_graph : block Addr.Map.t -> Addr.t -> t
+val build_graph : program -> Addr.t -> t
 
 val dominator_tree : t -> graph
 
diff --git a/compiler/lib/subst.ml b/compiler/lib/subst.ml
index cfa7cd1593..03e93c012c 100644
--- a/compiler/lib/subst.ml
+++ b/compiler/lib/subst.ml
@@ -66,33 +66,27 @@ 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) p.blocks in
-    { p with blocks }
+  let program s p = Code.map_blocks p ~f:(fun b -> block s b)
 
-  let rec cont' s pc blocks visited =
+  let rec cont' s pc p visited =
     if Addr.Set.mem pc visited
-    then blocks, visited
+    then p, visited
     else
       let visited = Addr.Set.add pc visited in
-      let b = Addr.Map.find pc blocks in
+      let b = Code.block pc p in
       let b = block s b in
-      let blocks = Addr.Map.add pc b blocks in
-      let blocks, visited =
-        List.fold_left b.body ~init:(blocks, visited) ~f:(fun (blocks, visited) instr ->
+      let p = Code.add_block pc b p in
+      let p, visited =
+        List.fold_left b.body ~init:(p, visited) ~f:(fun (p, visited) instr ->
             match instr with
-            | Let (_, Closure (_, (pc, _), _)) -> cont' s pc blocks visited
-            | _ -> blocks, visited)
+            | Let (_, Closure (_, (pc, _), _)) -> cont' s pc p visited
+            | _ -> p, visited)
       in
-      Code.fold_children
-        blocks
-        pc
-        (fun pc (blocks, visited) -> cont' s pc blocks visited)
-        (blocks, visited)
+      Code.fold_children p pc (fun pc (p, visited) -> cont' s pc p visited) (p, visited)
 
   let cont s addr p =
-    let blocks, _ = cont' s addr p.blocks Addr.Set.empty in
-    { p with blocks }
+    let p, _ = cont' s addr p Addr.Set.empty in
+    p
 end
 
 (****)
diff --git a/compiler/lib/subst.mli b/compiler/lib/subst.mli
index a3920f0650..2ab93d12ab 100644
--- a/compiler/lib/subst.mli
+++ b/compiler/lib/subst.mli
@@ -37,13 +37,6 @@ module Excluding_Binders : sig
   val last : (Var.t -> Var.t) -> last -> last
 
   val cont : (Var.t -> Var.t) -> int -> program -> program
-
-  val cont' :
-       (Var.t -> Var.t)
-    -> int
-    -> block Addr.Map.t
-    -> Addr.Set.t
-    -> block Addr.Map.t * Addr.Set.t
 end
 
 val from_array : Var.t array -> Var.t -> Var.t
diff --git a/compiler/lib/tailcall.ml b/compiler/lib/tailcall.ml
index 377407e544..16c87d76c7 100644
--- a/compiler/lib/tailcall.ml
+++ b/compiler/lib/tailcall.ml
@@ -45,8 +45,8 @@ let rec tail_call x f l =
     -> Some args
   | _ :: rem -> tail_call x f rem
 
-let rewrite_block update_count (f, f_params, f_pc, used) pc blocks =
-  let block = Addr.Map.find pc blocks in
+let rewrite_block update_count (f, f_params, f_pc, used) pc p =
+  let block = Code.block pc p in
   match block.branch with
   | Return x -> (
       match tail_call x f block.body with
@@ -57,98 +57,91 @@ let rewrite_block update_count (f, f_params, f_pc, used) pc blocks =
             List.iter2 f_params f_args ~f:(fun p a -> Code.Var.propagate_name p a);
             used := true;
             Some
-              (Addr.Map.add
+              (Code.add_block
                  pc
                  { params = block.params
                  ; body = remove_last block.body
                  ; branch = Branch (f_pc, f_args)
                  }
-                 blocks))
+                 p))
           else None
       | None -> None)
   | _ -> None
 
-let rec traverse update_count f pc visited blocks =
+let rec traverse update_count f pc visited p =
   if not (Addr.Set.mem pc visited)
   then
     let visited = Addr.Set.add pc visited in
-    match rewrite_block update_count f pc blocks with
-    | Some blocks ->
+    match rewrite_block update_count f pc p with
+    | Some p ->
         (* The block was rewritten with a branch to the top of the function.
            No need to visit children. *)
-        visited, blocks
+        visited, p
     | None ->
-        let visited, blocks =
-          Code.fold_children_skip_try_body
-            blocks
-            pc
-            (fun pc (visited, blocks) ->
-              let visited, blocks = traverse update_count f pc visited blocks in
-              visited, blocks)
-            (visited, blocks)
-        in
-        visited, blocks
-  else visited, blocks
+        Code.fold_children_skip_try_body
+          p
+          pc
+          (fun pc (visited, p) -> traverse update_count f pc visited p)
+          (visited, p)
+  else visited, p
 
 let f p =
   let previous_p = p in
   Code.invariant p;
-  let free_pc = ref p.free_pc in
-  let blocks = ref p.blocks in
+  let p = ref p in
   let update_count = ref 0 in
   let t = Timer.make () in
   Addr.Map.iter
     (fun pc _ ->
-      let block = Addr.Map.find pc !blocks in
+      let block = Code.block pc !p in
       let rewrite_body = ref false in
       let body =
         List.map block.body ~f:(function
           | Let (f, Closure (params, (pc_head, args), cloc)) as i ->
               if List.equal ~eq:Code.Var.equal params args
               then (
-                blocks :=
+                p :=
                   snd
                     (traverse
                        update_count
                        (f, params, pc_head, ref false)
                        pc_head
                        Addr.Set.empty
-                       !blocks);
+                       !p);
                 i)
               else
-                let intermediate_pc = !free_pc in
+                let intermediate_pc = Code.free_pc !p in
                 let need_to_create_intermediate_block = ref false in
-                blocks :=
+                p :=
                   snd
                     (traverse
                        update_count
                        (f, params, intermediate_pc, need_to_create_intermediate_block)
                        pc_head
                        Addr.Set.empty
-                       !blocks);
+                       !p);
                 if !need_to_create_intermediate_block
                 then (
-                  incr free_pc;
                   let new_params = List.map params ~f:Code.Var.fork in
                   let body =
                     (* duplicate the debug event before the loop header. *)
-                    match (Addr.Map.find pc_head !blocks).body with
+                    match (Code.block pc_head !p).body with
                     | (Event _ as e) :: _ -> [ e ]
                     | _ -> []
                   in
-                  blocks :=
-                    Addr.Map.add
+                  p :=
+                    Code.add_block
                       intermediate_pc
                       { params; body; branch = Branch (pc_head, args) }
-                      !blocks;
+                      !p;
                   rewrite_body := true;
                   Let (f, Closure (new_params, (intermediate_pc, new_params), cloc)))
                 else i
           | i -> i)
       in
-      if !rewrite_body then blocks := Addr.Map.add pc { block with body } !blocks)
-    p.blocks;
-  let p = { p with blocks = !blocks; free_pc = !free_pc } in
+      if !rewrite_body then p := Code.add_block pc { block with body } !p)
+    (Code.blocks !p);
+  let p = !p in
   if times () then Format.eprintf "  tail calls: %a@." Timer.print t;
   if stats () then Format.eprintf "Stats - tail calls: %d optimizations@." !update_count;
   if debug_stats ()