(*(* Let's say I have those types *)

type ('input, 'output, 'datatype, 'make_argument) polymorphic_finite_partial_function_implementation =
  {
    make: 'make_argument -> 'datatype;
    set: 'input -> 'output -> 'datatype -> 'datatype;
    unset: 'input -> 'datatype -> 'datatype;
    get: 'input -> 'datatype -> 'output option
  }

type ('input, 'output, 'datatype, 'make_argument) t =
  {
    data: 'datatype;
    implementation: ('input, 'output, 'datatype, 'make_argument) polymorphic_finite_partial_function_implementation
  }

(* Where (i, o, d, a) t represents i -> o
(d is just the datatype used to actually represent it and a is the type of what is needed to initialize an element of d)

And I have two "subtypes":
*)

(* 'input1 * 'input2 -> 'output *)
type ('input, 'output, 'datatype, 'make_argument) uncurried_t = ('input, 'output, 'datatype, 'make_argument) t
constraint 'input = 'input1 * 'input2

(* 'input1 -> 'input2 -> 'output *)
type ('input, 'output, 'datatype, 'make_argument) curried_t = ('input, 'output, 'datatype, 'make_argument) t
constraint 'output = 'input2 -> 'output2

(*  And for some random reason, I would like to build a function between those types

uncurry : ('input1 -> 'input2 -> 'output) -> ('input1 * 'input2 -> 'output)

[ mainly because ('input1 -> 'input2 -> 'output) is more efficient for everything but
('input1 * 'input2 -> 'output) is as efficient and easier to use for some others ]

The problem is that only have access to the implementation of ('input1 -> 'input2 -> 'output) and have no
way of getting the implementation of ('input2 -> 'output) [unless I'm lucky and I get an element of ('input2 -> 'output) by
applying an element of ('input1 -> 'input2 -> 'output) to an element of ('input1), but in some cases I can't].

A possible solution would be to add something to the implementation that returns the implementation of the codomain but then
it becomes weird for non-function types. The only thing I thought of was to replace non-function types x by unit -> x when they
appear as a codomain to allow implementing that additionnal method but it feels weird.

Is there something else I could do? My only constraint is that operations on t must works on all types that represent functions

The rest of this file contains a version of what I want (functionality-wise, not convenience-wise) that uses modules that works
(or at least compiles) and the same thing in the version where implementations are carried arround as records
(for which I can't finish the uncurry function *)



*)






open Batteries
open Map

module type POLYMORPHIC_FINITE_PARTIAL_FUNCTION = sig
  type ('input, 'output) t

  type ('input, 'output) make_argument

  val make : ('input, 'output) make_argument -> ('input, 'output) t

  val set : 'input -> 'output -> ('input, 'output) t -> ('input, 'output) t

  val unset : 'input -> ('input, 'output) t -> ('input, 'output) t

  val get : 'input -> ('input, 'output) t -> 'output option

  val set_option : 'input -> 'output option -> ('input, 'output) t -> ('input, 'output) t
end

module FromConstant :  POLYMORPHIC_FINITE_PARTIAL_FUNCTION = struct
  type ('input, 'output) t = 'output option

  type ('input, 'output) make_argument = unit

  let make x = None

  let set _ output _ = Some output

  let unset _ _ = None

  let set_option _ output_option _ = output_option

  let get _ f = f
end

module FromPMap : POLYMORPHIC_FINITE_PARTIAL_FUNCTION = struct
  type ('input, 'output) t = ('input, 'output) PMap.t

  type ('input, 'output) make_argument = unit

  let make () =
    PMap.empty

  let set input output f =
    PMap.add input output f

  let unset input f =
    PMap.remove input f

  let set_option input output_option f =
    match output_option with
    | None -> unset input f
    | Some output -> set input output f

  let get input f =
    try
      Some (PMap.find input f)
    with
    | Not_found -> None
end

module Uncurry (Outer : POLYMORPHIC_FINITE_PARTIAL_FUNCTION) (Inner : POLYMORPHIC_FINITE_PARTIAL_FUNCTION) = struct
  type ('input, 'output) t = ('outer_input, ('inner_input, 'output) Inner.t) Outer.t * ('inner_input, 'output) Inner.make_argument
  constraint 'input = 'inner_input * 'outer_input

  type ('input, 'output) make_argument =  ('outer_input, ('inner_input, 'output) Inner.t) Outer.make_argument * ('inner_input, 'output) Inner.make_argument
  constraint 'input = 'inner_input * 'outer_input

  let make (outer_make_argument, inner_make_argument) =
    (Outer.make outer_make_argument, inner_make_argument)

  let set (outer_input, inner_input) output (outer_f, inner_make_argument) =
    let inner_f =
      match Outer.get outer_input outer_f with
      | None -> Inner.make inner_make_argument
      | Some inner_f' -> inner_f'
    in
    let new_inner_f = Inner.set inner_input output inner_f in
    let new_outer_f = Outer.set outer_input new_inner_f in
    (new_outer_f, inner_make_argument)

  let unset (outer_input, inner_input) (outer_f, inner_make_argument) =
    match Outer.get outer_input outer_f with
    | None -> (outer_f, inner_make_argument)
    | Some inner_f ->
      let new_inner_f = Inner.unset inner_input inner_f in
      let new_outer_f = Outer.set outer_input new_inner_f outer_f in
      (new_outer_f, inner_make_argument)

  let get (outer_input, inner_input) (outer_f, _) =
    match Outer.get outer_input outer_f with
    | None -> None
    | Some inner_f -> Inner.get inner_input inner_f

end





(* Nearly perfect but one case doesn't work because there is no way to access the implementation of
  B -> C from the implementation of A -> B -> C *)

type ('input, 'output, 'datatype) polymorphic_finite_partial_function_implementation =
  {
    make: unit -> 'datatype;
    set: 'input -> 'output -> 'datatype -> 'datatype;
    unset: 'input -> 'datatype -> 'datatype;
    get: 'input -> 'datatype -> 'output option
  }

type ('input, 'output, 'datatype) t =
  {
    data: 'datatype;
    implementation: ('input, 'output, 'datatype) polymorphic_finite_partial_function_implementation
  }

let set input output f =
  {
    data = f.implementation.set input output f.data;
    implementation = f.implementation
  }

let unset input f =
  {
    data = f.implementation.unset input f.data;
    implementation = f.implementation
  }

let set_option input output_option f =
  match output_option with
  | None -> unset input f
  | Some output -> set input output f

let get input f =
  f.implementation.get input f.data

let pmap_implementation =
  let make' () =
    PMap.empty
  in
  let set' input output f =
    PMap.add input output f
  in
  let unset' input f =
    PMap.remove input f
  in
  let get' input f =
    try
      Some (PMap.find input f)
    with
    | Not_found -> None
  in
  {
    make = make';
    set = set';
    unset = unset';
    get = get'
  }

type ('input, 'output, 'datatype) uncurried_t = ('input, 'output, 'datatype) t
constraint 'input = 'input1 * 'input2

type ('input, 'output, 'datatype) curried_t = ('input, 'output, 'datatype) t
constraint 'output = 'input2 -> 'output2

let uncurry (f : ('outer_input, 'inner_input -> 'output, 'datatype) curried_t) : ('outer_input * 'inner_input, 'output, 'datatype) t =
  let make' () =
      f.implementation.make ()
  in
  let set' (outer_input, inner_input) output outer_f =
    let inner_f =
      match get outer_input outer_f with
      | None -> failwith "TODO" (* INNER_IMPLEMENTATION.make inner_make_argument,
                                and I can't access it because the only element of type outer_input I have doesn't
                                give back a value when the function is applied to it *)
      | Some inner_f' -> inner_f'
    in
    let new_inner_f = set inner_input output inner_f in
    let new_outer_f = set outer_input new_inner_f outer_f in
    new_outer_f
  in
  let unset' (outer_input, inner_input) outer_f =
    match get outer_input outer_f with
    | None -> outer_f
    | Some inner_f ->
      let new_inner_f = unset inner_input inner_f in
      let new_outer_f = set outer_input new_inner_f outer_f in
      new_outer_f
  in
  let get' (outer_input, inner_input) outer_f =
    match get outer_input outer_f with
    | None -> None
    | Some inner_f -> get inner_input inner_f
  in
  let implementation' =
    {
      make = make';
      set = set';
      unset = unset';
      get = get'
    }
  in
  {
    data = f.data;
    implementation = implementation'
  }