GADT-based index map experiment

indexMapGADT.ml

type read
type write
type no

type (_, _, _) t =
  | Read : ('i -> 'e) -> ('i, 'e, (read * no)) t
  | Write : ('i -> 'e -> unit) -> ('i, 'e, (no * write)) t
  | ReadWrite : ('i -> 'e) * ('i -> 'e -> unit) -> ('i, 'e, (read * write)) t

let make_ro get = Read get
let make_wo set = Write set
let make_rw get set = ReadWrite (get, set)

let read : type i e w. (i, e, (read * w)) t -> i -> e =
  fun ct ->
    match ct with
    | Read g -> g
    | ReadWrite (g, _) -> g

let write : type i e r. (i, e, (r * write)) t -> i -> e -> unit =
  fun ct ->
    match ct with
    | Write s -> s
    | ReadWrite (_, s) -> s

let get = read
let set = write

let map_ro :
  type i e w z. (i, e, (read * w)) t ->
    (e -> z) ->
    (i, z, (read * no)) t
    = fun ct f ->
      match ct with
      | Read g ->
        let get i = f (g i) in
        Read get
      | ReadWrite (g, _) ->
        let get i = f (g i) in
        Read get

let map_wo : type i e r z. (i, e, (r * write)) t ->
  (z -> e) ->
  (i, z, (no * write)) t
  = fun ct f ->
    match ct with
    | Write s ->
      let set i x = s i (f x) in
      Write set
    | ReadWrite (_, s) ->
      let set i x = s i (f x) in
      Write set

let map_rw :
  type i e z. (i, e, (read * write)) t ->
    (e -> z) ->
    (z -> e) ->
    (i, z, (read * write)) t
    = fun ct fr fw ->
      match ct with
      | ReadWrite (g, s) ->
        let get i = fr (g i) in
        let set i x = s i (fw x) in
        ReadWrite (get, set)

let mapi :
  type i e r w j. (i, e, (r * w)) t ->
    (j -> i) ->
    (j, e, (r * w)) t
    = fun ct f ->
      match ct with
      | Read g ->
        let get j = g (f j) in
        Read get
      | Write s ->
        let set j x = s (f j) x in
        Write set
      | ReadWrite (g, s) ->
        let get j = g (f j) in
        let set j x = s (f j) x in
        ReadWrite (get, set)

let to_ro :
  type i e w. (i, e, (read * w)) t ->
    (i, e, (read * no)) t
    = function
      | Read _ as r -> r
      | ReadWrite (g, _) -> Read g

let to_wo :
  type i e r. (i, e, (r * write)) t ->
    (i, e, (no * write)) t
    = function
      | Write _ as w -> w
      | ReadWrite (_, s) -> Write s

let of_array a = make_rw (Array.get a) (Array.set a)
let of_arrays m = make_rw (fun (i, j) -> m.(i).(j)) (fun (i, j) x -> m.(i).(j) <- x)

let of_array1 ba = make_rw (Bigarray.Array1.get ba) (Bigarray.Array1.set ba)
let of_array2 ba = make_rw (fun (i, j) -> Bigarray.Array2.get ba i j) (fun (i, j) x -> Bigarray.Array2.set ba i j x)
let of_array3 ba = make_rw (fun (i, j, k) -> Bigarray.Array3.get ba i j k) (fun (i, j, k) x -> Bigarray.Array3.set ba i j k x)
let of_genarray ba = make_rw (Bigarray.Genarray.get ba) (Bigarray.Genarray.set ba)

let to_row_major ct ~columns = mapi ct (fun (i, j) -> i * columns + j)
let to_column_major ct ~rows = mapi ct (fun (i, j) -> i + rows * j)

module Tuple2 = struct
  let fix_first ct i = mapi ct (fun j -> (i, j))
  let fix_second ct j = mapi ct (fun i -> (i, j))

  let transpose ct = mapi ct (fun (i, j) -> (j, i))
end

module Tuple3 = struct
  let fix_first ct i = mapi ct (fun (j, k) -> (i, j, k))
  let fix_second ct j = mapi ct (fun (i, k) -> (i, j, k))
  let fix_third ct k = mapi ct (fun (i, j) -> (i, j, k))

  let fix_first_second ct i j = mapi ct (fun k -> (i, j, k))
  let fix_first_third ct i k = mapi ct (fun j -> (i, j, k))
  let fix_second_third ct j k = mapi ct (fun i -> (i, j, k))
end

indexMapGADT.mli

type read
type write
type no

type ('i, 'e, 'p) t = private
  | Read : ('i -> 'e) -> ('i, 'e, (read * no)) t
  | Write : ('i -> 'e -> unit) -> ('i, 'e, (no * write)) t
  | ReadWrite : ('i -> 'e) * ('i -> 'e -> unit) -> ('i, 'e, (read * write)) t

val make_ro : ('i -> 'e) -> ('i, 'e, read * no) t
val make_wo : ('i -> 'e -> unit) -> ('i, 'e, no * write) t
val make_rw : ('i -> 'e) -> ('i -> 'e -> unit) -> ('i, 'e, read * write) t

val read : ('i, 'e, read * 'w) t -> 'i -> 'e
val write : ('i, 'e, 'r * write) t -> 'i -> 'e -> unit

val get : ('i, 'e, read * 'w) t -> 'i -> 'e
val set : ('i, 'e, 'r * write) t -> 'i -> 'e -> unit

val map_ro : ('i, 'e, read * 'w) t -> ('e -> 'z) -> ('i, 'z, read * no) t
val map_wo : ('i, 'e, 'r * write) t -> ('z -> 'e) -> ('i, 'z, no * write) t
val map_rw :
  ('i, 'e, read * write) t ->
  ('e -> 'z) -> ('z -> 'e) ->
  ('i, 'z, read * write) t

val mapi : ('i, 'e, 'r * 'w) t -> ('j -> 'i) -> ('j, 'e, 'r * 'w) t

val to_ro : ('i, 'e, read * 'w) t -> ('i, 'e, read * no) t
val to_wo : ('i, 'e, 'r * write) t -> ('i, 'e, no * write) t

val of_array : 'e array -> (int, 'e, read * write) t
val of_arrays : 'e array array -> (int * int, 'e, read * write) t

val of_array1 : ('e, _, _) Bigarray.Array1.t -> (int, 'e, read * write) t
val of_array2 : ('e, _, _) Bigarray.Array2.t -> (int * int, 'e, read * write) t
val of_array3 :
  ('e, _, _) Bigarray.Array3.t -> (int * int * int, 'e, read * write) t
val of_genarray :
  ('e, _, _) Bigarray.Genarray.t -> (int array, 'e, read * write) t

val to_row_major :
  (int, 'e, 'r * 'w) t -> columns:int -> (int * int, 'e, 'r * 'w) t
val to_column_major :
  (int, 'e, 'r * 'w) t -> rows:int -> (int * int, 'e, 'r * 'w) t

module Tuple2 : sig
  val fix_first : ('i * 'j, 'e, 'r * 'w) t -> 'i -> ('j, 'e, 'r * 'w) t
  val fix_second : ('i * 'j, 'e, 'r * 'w) t -> 'j -> ('i, 'e, 'r * 'w) t
  val transpose : ('i * 'j, 'e, 'r * 'w) t -> ('j * 'i, 'e, 'r * 'w) t
end
module Tuple3 : sig
  val fix_first :
    ('i * 'j * 'k, 'e, 'r * 'w) t -> 'i -> ('j * 'k, 'e, 'r * 'w) t
  val fix_second :
    ('i * 'j * 'k, 'e, 'r * 'w) t -> 'j -> ('i * 'k, 'e, 'r * 'w) t
  val fix_third :
    ('i * 'j * 'k, 'e, 'r * 'w) t -> 'k -> ('i * 'j, 'e, 'r * 'w) t
  val fix_first_second :
    ('i * 'j * 'k, 'e, 'r * 'w) t -> 'i -> 'j -> ('k, 'e, 'r * 'w) t
  val fix_first_third :
    ('i * 'j * 'k, 'e, 'r * 'w) t -> 'i -> 'k -> ('j, 'e, 'r * 'w) t
  val fix_second_third :
    ('i * 'j * 'k, 'e, 'r * 'w) t -> 'j -> 'k -> ('i, 'e, 'r * 'w) t
end