inline_sum.ml

module type INLINE_SUM = functor
  (Immediate: sig type t end) (Reference: sig type t end) ->
sig
  type t

  val of_immediate : Immediate.t -> t

  val of_reference : Reference.t -> t

  val case :
    immediate:(Immediate.t -> 'a) ->
    reference:(Reference.t -> 'a) ->
    t -> 'a
end



module Safe_not_inline_sum : INLINE_SUM = functor
  (Immediate: sig type t end) (Reference: sig type t end) ->
struct
  type t =
    | Immediate of Immediate.t (* [@inline] *)
    | Reference of Reference.t

  let of_immediate i = Immediate i

  let of_reference r = Reference r

  let case ~immediate ~reference = function
    | Immediate i -> immediate i
    | Reference r -> reference r
end

module Unsafe_inline_sum : INLINE_SUM = functor
  (Immediate: sig type t end) (Reference: sig type t end) ->
struct
  type t = Immediate.t

  let of_immediate t = t

  let of_reference = Obj.magic

  let case ~immediate ~reference t =
    if Obj.is_int (Obj.repr t) then
      immediate t
    else
      reference (Obj.magic t)
end


module Int = struct type t = int end

module Test = struct
  module Sum = Unsafe_inline_sum (Int) (String)

  let () = (* Test creation of immediate "variant" *)
    let t = Sum.of_immediate 0 in
    Sum.case t
      ~immediate:(fun n -> assert (n = 0))
      ~reference:(fun _ -> assert false)

  let () = (* Test creation of reference "variant" *)
    let t = Sum.of_reference "" in
    Sum.case t
      ~immediate:(fun _ -> assert false)
      ~reference:(fun s -> assert (s = ""))

  module Int_string_sum : sig
    type t

    val of_int : int -> t
    val of_string : string -> t

    val case :
      int:(int -> 'a) -> string:(string -> 'a) -> t -> 'a
  end = struct
    module Sum = Unsafe_inline_sum (Int) (String)

    type t = Sum.t

    let of_int = Sum.of_immediate
    let of_string = Sum.of_reference

    let case ~int ~string =
      Sum.case ~immediate:int ~reference:string
  end

  let () = (* Test creation of string "variant" *)
    let t = Int_string_sum.of_string "" in
    Int_string_sum.case t
      ~int:(fun _ -> assert false)
      ~string:(fun s -> assert (s = ""))
end



module Foo = struct
  type t =
    | Foo of string (* compiles to a pointer *)
    | Bar (* compiles to an immediate 0 *)
    | Baz (* compiles to an immediate 1 *)
    | Qux (* compiles to an immediate 2 *)
end




module Anti_example = struct
  open Foo

  module Sum = Unsafe_inline_sum (Int) (Foo)

  let () =
    let sum = Sum.of_reference Qux in
    Sum.case sum
      ~immediate:(fun n -> assert (n = 2))
      ~reference:(fun _ -> assert false)
end


module Anti_example_2 = struct
  open Foo

  module Sum = Unsafe_inline_sum (Foo) (String)

  let () =
    let sum = Sum.of_immediate Foo.(Foo "") in
    Sum.case sum
      ~immediate:(fun f -> assert false)
      ~reference:(fun s -> print_endline (String.escaped s))
end