WIP

Author: hhugo

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
+    (Code.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,16 +151,16 @@ 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
+    (Code.blocks program)
     program
 
 let f p =
   let t = Timer.make () in
   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

compiler/lib-wasm/generate.ml

@@ -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
+      (Code.blocks p)
       (false, false)
 
   let wrap_with_handler needed pc handler ~result_typ ~fall_through ~context body =
@@ -1193,7 +1193,7 @@ module Generate (Target : Target_sig.S) = struct
       { live = live_vars
       ; in_cps
       ; deadcode_sentinal
-      ; blocks = p.blocks
+      ; blocks = Code.blocks p
       ; closures
       ; global_context
       }
@@ -1275,16 +1275,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 +1291,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)

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 ->

compiler/lib/effects.ml

@@ -1082,28 +1082,28 @@ 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
           Code.add_block pc block p
       | (Let (x, e) as i) :: r when is_split_point i r branch ->
-          let pc' = Code.free_pc p 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 = Code.add_block pc block p in
-          split p pc' block' [] r branch
-      | i :: r -> split p pc block (i :: accu) r branch
+          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
       | [] -> false
       | 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 (Code.blocks p) p
@@ -1120,6 +1120,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)
@@ -1141,6 +1142,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;

compiler/lib/generate.ml

@@ -1711,14 +1711,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:(Code.program 0 ctx.blocks) ~f:(fun p pc ->
+                Subst.Excluding_Binders.cont subst pc p)
           in
-          { ctx with blocks = p }
+          { ctx with blocks = Code.blocks p }
       in
       let vd kind = function
         | [] -> []

compiler/lib/subst.ml

@@ -70,29 +70,29 @@ module Excluding_Binders = struct
     let blocks = Addr.Map.map (fun b -> block s b) (Code.blocks p) in
     Code.program (Code.start p) blocks
 
-  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
+        (Code.blocks p)
         pc
-        (fun pc (blocks, visited) -> cont' s pc blocks visited)
-        (blocks, visited)
+        (fun pc (p, visited) -> cont' s pc p visited)
+        (p, visited)
 
   let cont s addr p =
-    let blocks, _ = cont' s addr (Code.blocks p) Addr.Set.empty in
-    Code.program (Code.start p) blocks
+    let p, _ = cont' s addr p Addr.Set.empty in
+    p
 end
 
 (****)

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

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
+          (Code.blocks 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 (Code.free_pc p) in
-  let blocks = ref (Code.blocks p) 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)
-    (Code.blocks p);
-  let p = Code.program (Code.start p) !blocks 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 ()