Compiler: speedup subst

compiler/lib/code.ml

@@ -918,15 +918,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
@@ -962,6 +962,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') =
@@ -970,6 +989,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

@@ -317,14 +317,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/subst.ml

@@ -67,7 +67,25 @@ module Excluding_Binders = struct
     { 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
+    let count = ref 0 in
+    let s' x =
+      let y = s x in
+      if not (Code.Var.equal x y) then incr count;
+      y
+    in
+    let blocks =
+      Addr.Map.fold
+        (fun pc b blocks ->
+          let saved = !count in
+          let b' = block s' b in
+          if !count = saved
+          then (
+            Code.assert_block_equal ~name:"subst" b b';
+            blocks)
+          else Addr.Map.add pc b' blocks)
+        p.blocks
+        p.blocks
+    in
     { p with blocks }
 
   let rec cont' s pc blocks visited =