Merge pull request #60 from CraigFe/type-attributes

Add support for extensible type attributes

src/repr/attribute.ml

@@ -0,0 +1,93 @@
+include Attribute_intf
+open Higher
+
+module Key = struct
+  type 'f t = { uid : int; name : string; wit : 'f Witness.t }
+
+  let uid =
+    let counter = ref (-1) in
+    fun () ->
+      incr counter;
+      !counter
+
+  let create ~name =
+    let uid = uid () in
+    let wit = Witness.make () in
+    { uid; name; wit }
+
+  let name t = t.name
+
+  type 'a ty = 'a t
+
+  module Boxed = struct
+    type t = E : _ ty -> t [@@ocaml.unboxed]
+
+    let compare (E k1) (E k2) = Int.compare k1.uid k2.uid
+  end
+end
+
+module Map = struct
+  open Map.Make (Key.Boxed)
+
+  type ('a, 'f) data = ('a, 'f) app
+  type 'a binding = B : 'f Key.t * ('a, 'f) data -> 'a binding
+  type nonrec 'a t = 'a binding t
+
+  let empty = empty
+  let is_empty = is_empty
+  let mem t k = mem (E k) t
+  let add t ~key ~data = add (E key) (B (key, data)) t
+
+  let update :
+      type a f.
+      a t -> f Key.t -> ((a, f) data option -> (a, f) data option) -> a t =
+   fun t k f ->
+    update (E k)
+      (fun b ->
+        let v =
+          f
+            (match b with
+            | None -> None
+            | Some (B (k', v)) -> (
+                match Witness.eq k.wit k'.wit with
+                | None -> None
+                | Some Refl -> Some v))
+        in
+        match v with None -> None | Some v -> Some (B (k, v)))
+      t
+
+  let singleton k v = singleton (E k) (B (k, v))
+  let iter t ~f = iter (fun _ b -> f b) t
+  let for_all t ~f = for_all (fun _ b -> f b) t
+  let exists t ~f = exists (fun _ b -> f b) t
+  let cardinal t = cardinal t
+  let bindings t = bindings t |> List.map snd
+
+  let find : type a f. a t -> f Key.t -> (a, f) data option =
+   fun t k ->
+    match find_opt (E k) t with
+    | None -> None
+    | Some (B (k', v)) -> (
+        match Witness.eq k.wit k'.wit with None -> None | Some Refl -> Some v)
+end
+
+module Make1 (T : sig
+  type 'a t
+
+  val name : string
+end) =
+struct
+  include T
+  include Branded.Make (T)
+
+  let key : br Key.t = Key.create ~name
+
+  let find map =
+    match Map.find map key with None -> None | Some x -> Some (prj x)
+
+  let add data map = Map.add map ~key ~data:(inj data)
+end
+
+include Key
+
+module type S1 = S1 with type 'a attr := 'a t and type 'a map := 'a Map.t

src/repr/attribute.mli

@@ -0,0 +1,2 @@
+include Attribute_intf.Attribute
+(** @inline *)

src/repr/attribute_intf.ml

@@ -0,0 +1,76 @@
+open Higher
+
+(** An attribute key is a value that can be used to attach polymorphic data to a
+    heterogeneous attribute map. *)
+module type S1 = sig
+  type 'a attr
+  type 'a map
+
+  type 'a t
+  (** The type of data associated with the {!attr} attribute key. *)
+
+  val add : 'a t -> 'a map -> 'a map
+  (** Attach data for {!attr} to a given map. *)
+
+  val find : 'a map -> 'a t option
+  (** Search for data corresponding to the key {!attr} in the given map. *)
+
+  include Branded.S with type 'a t := 'a t
+
+  val key : br attr
+end
+
+module type Attribute = sig
+  type 'f t
+  (** An ['f t] is an attribute key that can be used to pack polymorphic data
+      into a heterogeneous {!Map} (and then recover it again).
+
+      The type parameter ['f] is the brand of a type operator [f : * ⇒ *]
+      which, when applied to the type parameter ['a] of a {!Map.t}, gives the
+      type ['a f] of the associated data. This allows a single attribute key to
+      store {i polymorphic} data. *)
+
+  val create : name:string -> _ t
+  (** [create ~name] is a fresh attribute key with the given string name. *)
+
+  val name : _ t -> string
+  (** Get the string name of an attribute key. *)
+
+  module Map : sig
+    type 'f key := 'f t
+
+    type 'a t
+    (** The type of polymorphic, heterogeneous maps. *)
+
+    type ('a, 'f) data := ('a, 'f) app
+    (** Given an ['a t] map and an ['f key] attribute key, the type of the
+        corresponding data is [('a, 'f) Higher.app]. *)
+
+    val empty : _ t
+    val is_empty : _ t -> bool
+    val mem : 'a t -> 'f key -> bool
+    val add : 'a t -> key:'f key -> data:('a, 'f) data -> 'a t
+
+    val update :
+      'a t -> 'f key -> (('a, 'f) data option -> ('a, 'f) data option) -> 'a t
+
+    val singleton : 'f key -> ('a, 'f) data -> 'a t
+
+    type 'a binding = B : 'f key * ('a, 'f) data -> 'a binding
+
+    val iter : 'a t -> f:('a binding -> unit) -> unit
+    val for_all : 'a t -> f:('a binding -> bool) -> bool
+    val exists : 'a t -> f:('a binding -> bool) -> bool
+    val cardinal : 'a t -> int
+    val find : 'a t -> 'f key -> ('a, 'f) data option
+    val bindings : 'a t -> 'a binding list
+  end
+
+  module type S1 = S1 with type 'a attr := 'a t and type 'a map := 'a Map.t
+
+  module Make1 (T : sig
+    type 'a t
+
+    val name : string
+  end) : S1 with type 'a t = 'a T.t
+end

src/repr/higher.ml

@@ -0,0 +1,24 @@
+(** Defunctionalised higher-kinded types. See "Lightweight Higher-Kinded
+    Polymorphism" (Yallop and White, 2014) for more details. *)
+
+type ('a, 'f) app
+
+module Branded = struct
+  module type S = sig
+    type 'a t
+    type br
+
+    external inj : 'a t -> ('a, br) app = "%identity"
+    external prj : ('a, br) app -> 'a t = "%identity"
+  end
+
+  module Make (T : sig
+    type 'a t
+  end) : S with type 'a t := 'a T.t = struct
+    type 'a t = 'a T.t
+    type br
+
+    external inj : 'a t -> ('a, br) app = "%identity"
+    external prj : ('a, br) app -> 'a t = "%identity"
+  end
+end

src/repr/type.ml

@@ -58,7 +58,7 @@ let array ?(len = `Int) v = Array { v; len }
 let pair a b = Tuple (Pair (a, b))
 let triple a b c = Tuple (Triple (a, b, c))
 let option a = Option a
-let boxed t = Boxed t
+let boxed t = annotate ~add:Type_binary.Boxed.add ~data:() t
 
 let abstract ~pp ~of_string ~json ~bin ?unboxed_bin ~equal ~compare ~short_hash
     ~pre_hash () =
@@ -73,8 +73,6 @@ let abstract ~pp ~of_string ~json ~bin ?unboxed_bin ~equal ~compare ~short_hash
       pp;
       of_string;
       pre_hash;
-      encode_json;
-      decode_json;
       encode_bin;
       decode_bin;
       size_of;
@@ -85,6 +83,8 @@ let abstract ~pp ~of_string ~json ~bin ?unboxed_bin ~equal ~compare ~short_hash
       unboxed_decode_bin;
       unboxed_size_of;
     }
+  |> annotate ~add:Encode_json.add ~data:encode_json
+  |> annotate ~add:Decode_json.add ~data:decode_json
 
 (* fix points *)
 
@@ -342,13 +342,11 @@ let partially_abstract ~pp ~of_string ~json ~bin ~unboxed_bin ~equal ~compare
       ~undefined:(fun () -> undefined' "pre_hash")
       ~structural:(fun () -> encode_bin)
   in
-  Custom
+  Type_core.Custom
     {
       cwit = `Type t;
       pp;
       of_string;
-      encode_json;
-      decode_json;
       encode_bin;
       decode_bin;
       size_of;
@@ -360,6 +358,8 @@ let partially_abstract ~pp ~of_string ~json ~bin ~unboxed_bin ~equal ~compare
       unboxed_decode_bin;
       unboxed_size_of;
     }
+  |> annotate ~add:Encode_json.add ~data:encode_json
+  |> annotate ~add:Decode_json.add ~data:decode_json
 
 let like ?pp ?of_string ?json ?bin ?unboxed_bin ?equal ?compare ?short_hash
     ?pre_hash t =
@@ -438,6 +438,10 @@ module Json = struct
     list (pair string a) |> like ~json
 end
 
+module Attribute = struct
+  let set_random f ty = annotate ~add:Type_random.Attr.add ~data:f ty
+end
+
 let ref : type a. a t -> a ref t = fun a -> map a ref (fun t -> !t)
 let lazy_t : type a. a t -> a Lazy.t t = fun a -> map a Lazy.from_val Lazy.force
 

src/repr/type_binary.ml

@@ -18,6 +18,12 @@ open Type_core
 open Staging
 open Utils
 
+module Boxed = Attribute.Make1 (struct
+  type _ t = unit
+
+  let name = "boxed"
+end)
+
 module Encode = struct
   let chars =
     Array.init 256 (fun i -> Bytes.unsafe_to_string (Bytes.make 1 (Char.chr i)))
@@ -129,7 +135,7 @@ module Encode = struct
     | Custom c -> c.encode_bin
     | Map b -> map ~boxed:true b
     | Prim t -> prim ~boxed:true t
-    | Boxed b -> t b
+    | Attributes { attr_type = x; _ } -> t x
     | List l -> list (t l.v) l.len
     | Array a -> array (t a.v) a.len
     | Tuple t -> tuple t
@@ -143,7 +149,8 @@ module Encode = struct
     | Custom c -> c.unboxed_encode_bin
     | Map b -> map ~boxed:false b
     | Prim t -> prim ~boxed:false t
-    | Boxed b -> t b
+    | Attributes { attr_type = x; attrs } -> (
+        match Boxed.find attrs with Some () -> t x | None -> unboxed x)
     | List l -> list (t l.v) l.len
     | Array a -> array (t a.v) a.len
     | Tuple t -> tuple t
@@ -332,7 +339,7 @@ module Decode = struct
     | Custom c -> c.decode_bin
     | Map b -> map ~boxed:true b
     | Prim t -> prim ~boxed:true t
-    | Boxed b -> t b
+    | Attributes { attr_type = x; _ } -> t x
     | List l -> list (t l.v) l.len
     | Array a -> array (t a.v) a.len
     | Tuple t -> tuple t
@@ -346,7 +353,8 @@ module Decode = struct
     | Custom c -> c.unboxed_decode_bin
     | Map b -> map ~boxed:false b
     | Prim t -> prim ~boxed:false t
-    | Boxed b -> t b
+    | Attributes { attr_type = x; attrs } -> (
+        match Boxed.find attrs with Some () -> t x | None -> unboxed x)
     | List l -> list (t l.v) l.len
     | Array a -> array (t a.v) a.len
     | Tuple t -> tuple t

src/repr/type_binary.mli

@@ -16,6 +16,7 @@
 
 open Type_core
 open Staging
+module Boxed : Attribute.S1 with type _ t = unit
 
 val encode_bin : 'a t -> 'a encode_bin
 val decode_bin : 'a t -> 'a decode_bin

src/repr/type_core.ml

@@ -33,6 +33,23 @@ module Json = struct
     | [] -> Jsonm.decode e.d
 end
 
+module Encode_json = Attribute.Make1 (struct
+  type 'a t = Jsonm.encoder -> 'a -> unit
+
+  let name = "encode_json"
+end)
+
+module Decode_json = Attribute.Make1 (struct
+  type 'a t = json_decoder -> ('a, [ `Msg of string ]) result
+
+  let name = "decode_json"
+end)
+
+let annotate t ~add ~data =
+  match t with
+  | Attributes t -> Attributes { t with attrs = add data t.attrs }
+  | t -> Attributes { attrs = add data Attribute.Map.empty; attr_type = t }
+
 let partial ?(pp = fun _ -> failwith "`pp` not implemented")
     ?(of_string = fun _ -> failwith "`of_string` not implemented")
     ?(encode_json = fun _ -> failwith "`encode_json` not implemented")
@@ -56,8 +73,6 @@ let partial ?(pp = fun _ -> failwith "`pp` not implemented")
       cwit = `Witness (Witness.make ());
       pp;
       of_string;
-      encode_json;
-      decode_json;
       short_hash;
       pre_hash;
       compare;
@@ -69,6 +84,8 @@ let partial ?(pp = fun _ -> failwith "`pp` not implemented")
       unboxed_decode_bin;
       unboxed_size_of;
     }
+  |> annotate ~add:Encode_json.add ~data:encode_json
+  |> annotate ~add:Decode_json.add ~data:decode_json
 
 let rec fields_aux : type a b. (a, b) fields -> a a_field list = function
   | F0 -> []