Granular_marshal: type-safe caching

src/index-format/granular_map.ml

@@ -36,6 +36,7 @@ module type S = sig
   val mem : key -> 'a t -> bool
   val update : key -> ('a option -> 'a option) -> 'a t -> 'a t
   val schema :
+    'a t Type.Id.t ->
     Granular_marshal.iter ->
     (Granular_marshal.iter -> key -> 'a -> unit) ->
     'a t ->
@@ -296,12 +297,12 @@ module Make (Ord : Map.OrderedType) = struct
         let rr = update x f r in
         if r == rr then t else bal l v d rr
 
-  let rec schema iter f m =
-    iter.yield m @@ fun iter tree ->
+  let rec schema type_id iter f m =
+    iter.yield m type_id @@ fun iter tree ->
     match tree with
     | Empty -> ()
     | Node { l; v; d; r; _ } ->
-      schema iter f l;
+      schema type_id iter f l;
       f iter v d;
-      schema iter f r
+      schema type_id iter f r
 end

src/index-format/granular_map.mli

@@ -34,6 +34,7 @@ module type S = sig
   val mem : key -> 'a t -> bool
   val update : key -> ('a option -> 'a option) -> 'a t -> 'a t
   val schema :
+    'a t Type.Id.t ->
     Granular_marshal.iter ->
     (Granular_marshal.iter -> key -> 'a -> unit) ->
     'a t ->

src/index-format/granular_marshal.ml

@@ -2,7 +2,7 @@ module Cache = Hashtbl.Make (Int)
 
 type store = { filename : string; cache : any_link Cache.t }
 
-and any_link = Link : 'a link * 'a schema -> any_link
+and any_link = Link : 'a link * 'a link Type.Id.t -> any_link
 
 and 'a link = 'a repr ref
 
@@ -18,7 +18,7 @@ and 'a repr =
 
 and 'a schema = iter -> 'a -> unit
 
-and iter = { yield : 'a. 'a link -> 'a schema -> unit }
+and iter = { yield : 'a. 'a link -> 'a link Type.Id.t -> 'a schema -> unit }
 
 let schema_no_sublinks : _ schema = fun _ _ -> ()
 
@@ -34,21 +34,24 @@ let read_loc store fd loc schema =
   let v = Marshal.from_channel fd in
   let rec iter =
     { yield =
-        (fun lnk schema ->
+        (fun (type a) (lnk : a link) type_id schema ->
           match !lnk with
           | Small v ->
             schema iter v;
             lnk := In_memory v
           | Serialized { loc } -> lnk := On_disk { store; loc; schema }
           | Serialized_reused { loc } -> (
             match Cache.find store.cache loc with
-            | Link (lnk', schema') ->
-              let lnk' = normalize lnk' in
-              assert (schema == Obj.magic schema');
-              lnk := Duplicate (Obj.magic lnk')
+            | Link (type b) ((lnk', type_id') : b link * _) -> (
+              match Type.Id.provably_equal type_id type_id' with
+              | Some (Equal : (a link, b link) Type.eq) ->
+                lnk := Duplicate (normalize lnk')
+              | None ->
+                invalid_arg
+                  "Granular_marshal.read_loc: reuse of a different type")
             | exception Not_found ->
               lnk := On_disk { store; loc; schema };
-              Cache.add store.cache loc (Link (lnk, schema)))
+              Cache.add store.cache loc (Link (lnk, type_id)))
           | In_memory _ | In_memory_reused _ | On_disk _ | Duplicate _ -> ()
           | Placeholder -> invalid_arg "Granular_marshal.read_loc: Placeholder")
     }
@@ -131,7 +134,7 @@ let write ?(flags = []) fd root_schema root_value =
   output_string fd (String.make ptr_size '\000');
   let rec iter size ~placeholders ~restore =
     { yield =
-        (fun (type a) (lnk : a link) (schema : a schema) : unit ->
+        (fun (type a) (lnk : a link) _type_id (schema : a schema) : unit ->
           match !lnk with
           | Serialized _ | Serialized_reused _ | Small _ -> ()
           | Placeholder -> failwith "big nono"

src/index-format/granular_marshal.mli

@@ -17,7 +17,7 @@ val cache : 'a. (module Hashtbl.HashedType with type t = 'a) -> 'a link -> unit
     We of course have [fetch (link v) = v] and [link (fetch lnk) = lnk]. *)
 val fetch : 'a link -> 'a
 
-(** For this demo we can't depend on a PPX or external dependencies,
+(** For Merlin we can't depend on a PPX or external dependencies,
     so we require a user-defined {!schema} to describe where the links can be
     found.  This is just an iter traversal over the values, recursively
     yielding on any reachable link. Since links can point to values themselves
@@ -29,10 +29,14 @@ val fetch : 'a link -> 'a
     {[
       type t = { first : string link ; second : int link list link }
 
+      let type_first : string link Type.Id.t = Type.Id.make ()
+      let type_second : int link list link Type.Id.t = Type.Id.make ()
+      let type_v : int link Type.Id.t = Type.Id.make ()
+
       let schema : t schema = fun iter t ->
-        iter.yield t.first schema_no_sublinks ;
-        iter.yield t.second @@ fun iter lst ->
-          List.iter (fun v -> iter.yield v schema_no_sublinks) lst
+        iter.yield t.first type_first schema_no_sublinks ;
+        iter.yield t.second type_second @@ fun iter lst ->
+          List.iter (fun v -> iter.yield v type_v schema_no_sublinks) lst
     ]}
 
     where {!schema_no_sublinks} indicates that the yielded value contains
@@ -44,7 +48,7 @@ type 'a schema = iter -> 'a -> unit
 (** A callback to signal the reachable links and the schema of their pointed
     sub-value.  Since a value can contain multiple links each pointing to
     different types of values, the callback is polymorphic. *)
-and iter = { yield : 'a. 'a link -> 'a schema -> unit }
+and iter = { yield : 'a. 'a link -> 'a link Type.Id.t -> 'a schema -> unit }
 
 (** A schema usable when the ['a] value does not contain any links. *)
 val schema_no_sublinks : 'a schema

src/index-format/granular_set.ml

@@ -265,8 +265,10 @@ module Make (Ord : Set.OrderedType) = struct
       f v;
       iter f r
 
+  let type_id = Type.Id.make ()
+
   let rec schema iter f m =
-    iter.yield m @@ fun iter tree ->
+    iter.yield m type_id @@ fun iter tree ->
     match tree with
     | Empty -> ()
     | Node { l; v; r; _ } ->

src/index-format/index_format.ml

@@ -16,6 +16,11 @@ module Union_find = struct
   let union a b =
     Granular_marshal.(
       link (Union_find.union ~f:Uid_set.union (fetch a) (fetch b)))
+
+  let type_id : t Type.Id.t = Type.Id.make ()
+
+  let schema { Granular_marshal.yield } t =
+    yield t type_id Granular_marshal.schema_no_sublinks
 end
 
 let add map uid locs =
@@ -38,12 +43,18 @@ type index =
 
 let lidset_schema iter lidset = Lid_set.schema iter Lid.schema lidset
 
+let type_setmap : Lid_set.t Uid_map.t Type.Id.t = Type.Id.make ()
+let type_ufmap : Union_find.t Uid_map.t Type.Id.t = Type.Id.make ()
+
 let index_schema (iter : Granular_marshal.iter) index =
-  Uid_map.schema iter (fun iter _ v -> lidset_schema iter v) index.defs;
-  Uid_map.schema iter (fun iter _ v -> lidset_schema iter v) index.approximated;
-  Uid_map.schema iter
-    (fun iter _ v ->
-      iter.Granular_marshal.yield v Granular_marshal.schema_no_sublinks)
+  Uid_map.schema type_setmap iter
+    (fun iter _ v -> lidset_schema iter v)
+    index.defs;
+  Uid_map.schema type_setmap iter
+    (fun iter _ v -> lidset_schema iter v)
+    index.approximated;
+  Uid_map.schema type_ufmap iter
+    (fun iter _ v -> Union_find.schema iter v)
     index.related_uids
 
 let compress index =

src/index-format/lid.ml

@@ -51,9 +51,12 @@ let compare t1 t2 =
     | c -> c)
   | c -> c
 
+let type_string : string G.link Type.Id.t = Type.Id.make ()
+let type_longident : Longident.t G.link Type.Id.t = Type.Id.make ()
+
 let schema iter t =
-  iter.G.yield t.filename G.schema_no_sublinks;
-  iter.G.yield t.longident G.schema_no_sublinks
+  iter.G.yield t.filename type_string G.schema_no_sublinks;
+  iter.G.yield t.longident type_longident G.schema_no_sublinks
 
 module Li = struct
   include Longident