Preserve sharing of blocks

compiler/lib/code.ml

@@ -890,15 +890,15 @@ let return_values p =
           Var.Map.add name s rets)
     Var.Map.empty
 
+let block_equal b1 b2 =
+  phys_equal b1 b2
+  || List.equal ~eq:Var.equal b1.params b2.params
+     && Poly.equal b1.branch b2.branch
+     && List.equal ~eq:Poly.equal b1.body b2.body
+
 let equal p1 p2 =
   p1.start = p2.start
-  && Addr.Map.equal
-       (fun { params; body; branch } b ->
-         List.equal ~eq:Var.equal params b.params
-         && Poly.equal branch b.branch
-         && List.equal ~eq:Poly.equal body b.body)
-       p1.blocks
-       p2.blocks
+  && (phys_equal p1.blocks p2.blocks || Addr.Map.equal block_equal p1.blocks p2.blocks)
 
 let print_to_file p =
   let file = Filename.temp_file "jsoo" "prog" in
@@ -934,6 +934,25 @@ let check_updates ~name p1 p2 ~updates =
       print_diff p1 p2;
       assert false
 
+let print_block_sharing ~name p1 p2 =
+  let shared = ref 0 in
+  let updated = ref 0 in
+  Addr.Map.iter
+    (fun pc b2 ->
+      match Addr.Map.find_opt pc p1.blocks with
+      | Some b1 when phys_equal b1 b2 -> incr shared
+      | Some _ -> incr updated
+      | None -> incr updated)
+    p2.blocks;
+  let removed = Addr.Map.cardinal p1.blocks - !shared - !updated in
+  Format.eprintf
+    "Stats - %s sharing: %d/%d blocks shared, %d updated, %d removed@."
+    name
+    !shared
+    (Addr.Map.cardinal p2.blocks)
+    !updated
+    removed
+
 let cont_equal (pc, args) (pc', args') = pc = pc' && List.equal ~eq:Var.equal args args'
 
 let cont_compare (pc, args) (pc', args') =
@@ -942,6 +961,23 @@ let cont_compare (pc, args) (pc', args') =
 
 let with_invariant = Debug.find "invariant"
 
+let assert_block_equal ~name b_old b_new =
+  if with_invariant ()
+  then
+    if not (block_equal b_old b_new)
+    then (
+      Format.eprintf "ASSERT_BLOCK_EQUAL: %s: counter=0 but block differs.@." name;
+      assert false)
+
+let assert_program_equal ~name p_old p_new =
+  if with_invariant ()
+  then
+    if not (equal p_old p_new)
+    then (
+      Format.eprintf "ASSERT_PROGRAM_EQUAL: %s: counter=0 but program differs.@." name;
+      print_diff p_old p_new;
+      assert false)
+
 let do_compact { blocks; start; free_pc = _ } =
   let remap =
     let max = fst (Addr.Map.max_binding blocks) in

compiler/lib/code.mli

@@ -315,14 +315,22 @@ val compact : program -> program
 
 val is_empty : program -> bool
 
+val block_equal : block -> block -> bool
+
 val equal : program -> program -> bool
 
 val print_diff : program -> program -> unit
 
 val check_updates : name:string -> program -> program -> updates:int -> unit
 
+val print_block_sharing : name:string -> program -> program -> unit
+
 val invariant : program -> unit
 
+val assert_block_equal : name:string -> block -> block -> unit
+
+val assert_program_equal : name:string -> program -> program -> unit
+
 val cont_equal : cont -> cont -> bool
 
 val cont_compare : cont -> cont -> int

compiler/lib/deadcode.ml

@@ -206,14 +206,17 @@ let remove_unused_blocks' p =
         b)
       p.blocks
   in
-  { p with blocks }, !count
+  if !count = 0 then p, 0 else { p with blocks }, !count
 
 let remove_unused_blocks p =
   let previous_p = p in
   let t = Timer.make () in
   let p, count = remove_unused_blocks' p in
   if times () then Format.eprintf "  dead block: %a@." Timer.print t;
-  if stats () then Format.eprintf "Stats - dead block: deleted %d@." count;
+  if stats ()
+  then (
+    Format.eprintf "Stats - dead block: deleted %d@." count;
+    Code.print_block_sharing ~name:"dead block" previous_p p);
   if debug_stats () then Code.check_updates ~name:"dead block" previous_p p ~updates:count;
   p
 
@@ -326,7 +329,10 @@ let merge_blocks p =
       Subst.Excluding_Binders.program rename p
   in
   if times () then Format.eprintf "  merge block: %a@." Timer.print t;
-  if stats () then Format.eprintf "Stats - merge block: merged %d@." !merged;
+  if stats ()
+  then (
+    Format.eprintf "Stats - merge block: merged %d@." !merged;
+    Code.print_block_sharing ~name:"merge block" previous_p p);
   if debug_stats ()
   then Code.check_updates ~name:"merge block" previous_p p ~updates:!merged;
   p
@@ -480,38 +486,66 @@ let f pure_funs ({ blocks; _ } as p : Code.program) =
             st.deleted_blocks <- st.deleted_blocks + 1;
             None)
           else
-            Some
-              { params = List.filter block.params ~f:(fun x -> st.live.(Var.idx x) > 0)
-              ; 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 live_instr st i
-                          then filter_closure all_blocks st i :: acc
-                          else (
-                            st.deleted_instrs <- st.deleted_instrs + 1;
-                            acc))
-                  |> List.rev
-              ; branch = filter_live_last all_blocks st block.branch
-              })
+            let saved_instrs = st.deleted_instrs in
+            let saved_params = st.deleted_params in
+            let params_changed =
+              List.exists block.params ~f:(fun x -> st.live.(Var.idx x) = 0)
+            in
+            let params =
+              if params_changed
+              then List.filter block.params ~f:(fun x -> st.live.(Var.idx x) > 0)
+              else block.params
+            in
+            let 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 live_instr st i
+                      then filter_closure all_blocks st i :: acc
+                      else (
+                        st.deleted_instrs <- st.deleted_instrs + 1;
+                        acc))
+              |> List.rev
+            in
+            let branch = filter_live_last all_blocks st block.branch in
+            if
+              (not params_changed)
+              && st.deleted_instrs = saved_instrs
+              && st.deleted_params = saved_params
+            then (
+              Code.assert_block_equal ~name:"deadcode" block { params; body; branch };
+              Some block)
+            else Some { params; body; branch })
         blocks
     in
-    { p with blocks }
+    if st.deleted_instrs + st.deleted_blocks + st.deleted_params = 0
+    then (
+      Code.assert_program_equal ~name:"deadcode(filter)" p { p with blocks };
+      p)
+    else { p with blocks }
+  in
+  let p =
+    let p' = remove_empty_blocks st p in
+    if st.block_shortcut = 0
+    then (
+      Code.assert_program_equal ~name:"deadcode(shortcut)" p p';
+      p)
+    else p'
   in
-  let p = remove_empty_blocks st p in
   if times () then Format.eprintf "  dead code elim.: %a@." Timer.print t;
   if stats ()
-  then
+  then (
     Format.eprintf
       "Stats - dead code: deleted %d instructions, %d blocks, %d parameters, %d \
        branches@."
       st.deleted_instrs
       st.deleted_blocks
       st.deleted_params
       st.block_shortcut;
+    Code.print_block_sharing ~name:"deadcode" previous_p p);
   if debug_stats ()
   then
     Code.check_updates

compiler/lib/driver.ml

@@ -219,17 +219,43 @@ let rec loop max name round i (p : 'a) : 'a =
   else if Code.equal p' p
   then (
     if debug then Format.eprintf "%s#%d: fix-point reached.@." name i;
+    if stats () then Code.print_block_sharing ~name:(Printf.sprintf "%s#%d" name i) p p';
     p')
-  else loop max name round (i + 1) p'
+  else (
+    if stats () then Code.print_block_sharing ~name:(Printf.sprintf "%s#%d" name i) p p';
+    loop max name round (i + 1) p')
+
+let cached () =
+  let all_cached = ref [] in
+  fun f ->
+    let last = ref None in
+    all_cached := last :: !all_cached;
+    fun p ->
+      match !last with
+      | Some (blocks, result) when phys_equal p.Code.blocks blocks -> result
+      | _ ->
+          let result = f p in
+          List.iter !all_cached ~f:(fun r -> r := None);
+          last := Some (p.blocks, result);
+          result
 
 let round profile : 'a -> 'a =
-  print
-  +> tailcall
-  +> Ref_unboxing.f
-  +> (flow +> specialize +> eval +> fst)
-  +> inline profile
-  +> phi
-  +> deadcode
+  let cached = cached () in
+  let tailcall = cached tailcall in
+  let ref_unboxing = cached Ref_unboxing.f in
+  let flow_specialize_eval = cached (flow +> specialize +> eval +> fst) in
+  let inline = cached (inline profile) in
+  let phi = cached phi in
+  let deadcode = cached deadcode in
+  fun p ->
+    p
+    |> print
+    |> tailcall
+    |> ref_unboxing
+    |> flow_specialize_eval
+    |> inline
+    |> phi
+    |> deadcode
 
 (* o1 *)
 

compiler/lib/eval.ml

@@ -827,15 +827,26 @@ let drop_exception_handler drop_count blocks =
     blocks
 
 let eval update_count update_branch inline_constant ~target info blocks =
-  Addr.Map.map
-    (fun block ->
+  Addr.Map.fold
+    (fun pc block blocks ->
+      let saved_update = !update_count in
+      let saved_branch = !update_branch in
+      let saved_inline = !inline_constant in
       let body =
         List.concat_map
           block.body
           ~f:(eval_instr update_count inline_constant ~target info)
       in
       let branch = eval_branch update_branch info block.branch in
-      { block with Code.body; Code.branch })
+      if
+        !update_count = saved_update
+        && !update_branch = saved_branch
+        && !inline_constant = saved_inline
+      then (
+        Code.assert_block_equal ~name:"eval" block { block with Code.body; Code.branch };
+        blocks)
+      else Addr.Map.add pc { block with Code.body; Code.branch } blocks)
+    blocks
     blocks
 
 let f info p =
@@ -856,17 +867,24 @@ let f info p =
       p.blocks
   in
   let blocks = drop_exception_handler drop_count blocks in
-  let p = { p with blocks } in
+  let p =
+    if !update_count + !update_branch + !inline_constant + !drop_count = 0
+    then (
+      Code.assert_program_equal ~name:"eval" p { p with blocks };
+      p)
+    else { p with blocks }
+  in
   if times () then Format.eprintf "  eval: %a@." Timer.print t;
   if stats ()
-  then
+  then (
     Format.eprintf
       "Stats - eval: %d optimizations, %d inlined cst, %d dropped exception handlers, %d \
        branch updated@."
       !update_count
       !inline_constant
       !drop_count
       !update_branch;
+    Code.print_block_sharing ~name:"eval" previous_p p);
   if debug_stats ()
   then
     Code.check_updates