Source file type.ml

open Bigarray

module type TYPE = sig
  type t
  type elt
  type kind = (t, elt) Bigarray.kind

  val name : string
  val kind : kind
  val of_float : float -> t
  val to_float : t -> float
end

type ('a, 'b) t = (module TYPE with type t = 'a and type elt = 'b)

exception Unsupported

module U8 = struct
  type t = int
  type elt = int8_unsigned_elt
  type kind = (t, elt) Bigarray.kind

  let name = "u8"
  let kind = Int8_unsigned
  let of_float = int_of_float
  let to_float = float_of_int
end

module U16 = struct
  type t = int
  type elt = int16_unsigned_elt
  type kind = (t, elt) Bigarray.kind

  let name = "u16"
  let kind = Int16_unsigned
  let of_float = int_of_float
  let to_float = float_of_int
end

module F32 = struct
  type t = float
  type elt = float32_elt
  type kind = (t, elt) Bigarray.kind

  let name = "f32"
  let kind = Float32
  let of_float x = x
  let to_float x = x
end

module F64 = struct
  type t = float
  type elt = float64_elt
  type kind = (t, elt) Bigarray.kind

  let name = "f64"
  let kind = Float64
  let of_float x = x
  let to_float x = x
end

module I32 = struct
  type t = int32
  type elt = int32_elt
  type kind = (t, elt) Bigarray.kind

  let name = "i32"
  let kind = Int32
  let of_float = Int32.of_float
  let to_float = Int32.to_float
end

module I64 = struct
  type t = int64
  type elt = int64_elt
  type kind = (t, elt) Bigarray.kind

  let name = "i64"
  let kind = Int64
  let of_float = Int64.of_float
  let to_float = Int64.to_float
end

type u8 = int8_unsigned_elt
type u16 = int16_unsigned_elt
type i32 = int32_elt
type i64 = int64_elt
type f32 = float32_elt
type f64 = float64_elt

let u8 : (int, u8) t = (module U8)
let u16 : (int, u16) t = (module U16)
let i32 : (int32, i32) t = (module I32)
let i64 : (int64, i64) t = (module I64)
let f32 : (float, f32) t = (module F32)
let f64 : (float, f64) t = (module F64)
let name : type a b. (a, b) t -> string = fun (module T) -> T.name

let depth : type a b. (a, b) t -> int =
 fun (module T) ->
  match T.kind with
  | Int8_unsigned -> 8
  | Int16_unsigned -> 16
  | Int32 -> 32
  | Int64 -> 64
  | Float32 -> 32
  | Float64 -> 32
  | _ -> raise Unsupported

let[@inline] max : type a b. (a, b) t -> a =
 fun (module T) ->
  match T.kind with
  | Int8_unsigned -> 255
  | Int16_unsigned -> 65535
  | Int32 -> Int32.max_int
  | Int64 -> Int64.max_int
  | Float32 -> 1.0
  | Float64 -> 1.0
  | _ -> raise Unsupported

let[@inline] min : type a b. (a, b) t -> a =
 fun (module T) ->
  match T.kind with
  | Int8_unsigned -> 0
  | Int16_unsigned -> 0
  | Int32 -> Int32.min_int
  | Int64 -> Int64.min_int
  | Float32 -> 0.0
  | Float64 -> 0.0
  | _ -> raise Unsupported

let[@inline] max_f : type a b. (a, b) t -> float =
 fun (module T) ->
  match T.kind with
  | Int8_unsigned -> 255.
  | Int16_unsigned -> 65535.
  | Int32 -> Int32.max_int |> Int32.to_float
  | Int64 -> Int64.max_int |> Int64.to_float
  | Float32 -> 1.0
  | Float64 -> 1.0
  | _ -> raise Unsupported

let[@inline] min_f : type a b. (a, b) t -> float =
 fun (module T) ->
  match T.kind with
  | Int8_unsigned -> 0.0
  | Int16_unsigned -> 0.0
  | Int32 -> Int32.min_int |> Int32.to_float
  | Int64 -> Int64.min_int |> Int64.to_float
  | Float32 -> 0.0
  | Float64 -> 0.0
  | _ -> raise Unsupported

let[@inline] to_float : type a b. (a, b) t -> a -> float =
 fun (module T) v -> T.to_float v

let[@inline] of_float : type a b. (a, b) t -> float -> a =
 fun (module T) v -> T.of_float v

let clamp kind f =
  let min = min_f kind in
  let max = max_f kind in
  if f < min then min else if f > max then max else f

let normalize kind f =
  let min = min_f kind in
  let max = max_f kind in
  (f -. min) /. (max -. min)

let denormalize kind f =
  let min = min_f kind in
  let max = max_f kind in
  f *. (max -. min)

let convert ~from to_kind f =
  to_float from f |> normalize from |> denormalize to_kind |> of_float to_kind

let kind : type a b. (a, b) t -> (a, b) Bigarray.kind = fun (module T) -> T.kind

let of_kind (type a b) : (a, b) Bigarray.kind -> (a, b) t = function
  | Int8_unsigned -> u8
  | Int16_unsigned -> u16
  | Int32 -> i32
  | Int64 -> i64
  | Float32 -> f32
  | Float64 -> f64
  | _ -> raise Unsupported