WIP

Author: hhugo

compiler/lib/driver.ml

@@ -608,29 +608,31 @@ let simplify_js js =
 
 let collects_shapes p =
   let l = ref StringMap.empty in
-  let t = Timer.make () in
-  let _, info = Flow.f p in
-  let pure = Pure_fun.f p in
-  let return_values = Code.return_values p in
-  Code.Addr.Map.iter
-    (fun _ block ->
-      List.iter block.Code.body ~f:(fun i ->
-          match i with
-          | Code.Let
-              ( _
-              , Prim
-                  ( Extern "caml_register_global"
-                  , [ _code; Pv block; Pc (NativeString name) ] ) ) ->
-              let shape = Flow.the_shape_of ~return_values ~pure info block in
-              let name =
-                match name with
-                | Byte s -> s
-                | Utf (Utf8 s) -> s
-              in
-              l := StringMap.add name shape !l
-          | _ -> ()))
-    p.blocks;
-  if times () then Format.eprintf "  shapes: %a@." Timer.print t;
+  if true
+  then (
+    let t = Timer.make () in
+    let _, info = Flow.f p in
+    let pure = Pure_fun.f p in
+    let return_values = Code.return_values p in
+    Code.Addr.Map.iter
+      (fun _ block ->
+        List.iter block.Code.body ~f:(fun i ->
+            match i with
+            | Code.Let
+                ( _
+                , Prim
+                    ( Extern "caml_register_global"
+                    , [ _code; Pv block; Pc (NativeString name) ] ) ) ->
+                let shape = Flow.the_shape_of ~return_values ~pure info block in
+                let name =
+                  match name with
+                  | Byte s -> s
+                  | Utf (Utf8 s) -> s
+                in
+                l := StringMap.add name shape !l
+            | _ -> ()))
+      p.blocks;
+    if times () then Format.eprintf "  shapes: %a@." Timer.print t);
   !l
 
 let configure formatter =

compiler/lib/shape.ml

@@ -132,27 +132,34 @@ module Store = struct
 end
 
 module State = struct
-  type nonrec t = t Code.Var.Hashtbl.t
+  type nonrec t =
+    { table : t Code.Var.Hashtbl.t
+    ; cache : BitSet.t
+    }
 
-  let t : t = Code.Var.Hashtbl.create 17
+  let t : t = { table = Code.Var.Hashtbl.create 17; cache = BitSet.create () }
 
-  let assign x shape = Code.Var.Hashtbl.replace t x shape
+  let assign x shape =
+    Code.Var.Hashtbl.replace t.table x shape;
+    BitSet.set t.cache (Code.Var.idx x)
 
   let propagate x offset target =
-    match Code.Var.Hashtbl.find_opt t x with
+    match Code.Var.Hashtbl.find_opt t.table x with
     | None -> ()
     | Some (Top _ | Function _) -> ()
     | Some (Block l) -> assign target (List.nth l offset)
 
-  let get x = Code.Var.Hashtbl.find_opt t x
+  let get x = Code.Var.Hashtbl.find_opt t.table x
 
-  let mem x = Code.Var.Hashtbl.mem t x
+  let mem x = BitSet.mem t.cache (Code.Var.idx x)
 
   let is_pure_fun x =
-    match Code.Var.Hashtbl.find_opt t x with
+    match Code.Var.Hashtbl.find_opt t.table x with
     | None -> false
     | Some (Top _ | Block _) -> false
     | Some (Function { pure; _ }) -> pure
 
-  let reset () = Code.Var.Hashtbl.clear t
+  let reset () =
+    Code.Var.Hashtbl.clear t.table;
+    BitSet.clear t.cache
 end