fun_cost_rec.ml

module type Arrow =
sig
  type ('a,'b) t
  val arr   : ('a -> 'b) -> ('a, 'b) t
  val (>>>) : ('a,'b) t -> ('b,'c) t -> ('a,'c) t
  val first : ('a,'b) t -> ('a * 'c, 'b * 'c) t
end

module type Arrow_choice =
sig
  include Arrow
  val (<+>) : ('a,'b) t -> ('a,'b) t -> ('a,'b) t
  val apply : ('a,'b) t -> 'a -> 'b
  val cost  : ('a,'b) t -> int
end

module Fun_cost : Arrow_choice =
struct
  type ('a,'b) t =
    {apply   : 'a -> 'b;
     compose : 'c. ('b,'c) t -> ('a,'c) t;
     cost    : int}

  (* Analogous to [] *)
  let id =
    {apply = (fun x -> x);
     compose = (fun g -> g);
     cost = 0}

  (** arr primitive *)
  let rec arr : 'a 'b 'r. ('a -> 'b) -> ('b,'r) t -> ('a,'r) t =
    fun f k (* continuation i.e) rest of the list *) ->
      {apply = (fun x -> k.apply (f x));
       compose = (fun g -> arr f (k.compose g));
       cost = 1 + k.cost}

  (* Analogous to (fun x -> [x]) *)
  let arr f =
    arr f id

  (** Function composition. Analogous to (fun x y -> x @ y) *)
  let (>>>) f g = f.compose g

  (** first primitive *)
  let rec first' : 'a 'b 'c 'r. ('a,'b) t -> ('b * 'c, 'r) t -> ('a * 'c, 'r) t =
    fun f k ->
      {apply = (fun (x,y (* additional argument *)) -> (k.apply (f.apply x, y)));
       compose = (fun g -> first' f (k.compose g));
       cost = f.cost}

  let first f = first' f id

  (** Function application *)
  let apply f x = f.apply x

  (** Choose the function with the minimum cost *)
  let (<+>) : 'a 'b. ('a,'b) t -> ('a,'b) t -> ('a,'b) t =
    fun f g ->
      let h = if f.cost < g.cost then f else g in
      {compose = h.compose;
       apply = h.apply;
       cost = h.cost}

   (** Fetch cost *)
   let cost {cost; _} = cost
end

(*************************)

open Fun_cost

let () = Random.self_init ()
let c1 = arr (fun x -> List.length x > 0)
let _ = Printf.printf "cost of c1 = %d\n" @@ cost c1

let c2 = arr (fun y -> Printf.sprintf "%B" y)
let _ = Printf.printf "cost of c2 = %d\n" @@ cost c2

let c12 = c1 >>> c2
let _ = Printf.printf "cost of c12 = %d\n" @@ cost c12

let c3 = arr (fun l -> String.concat " " (List.map string_of_int l))
let _ = Printf.printf "cost of c3 = %d\n" @@ cost c3

let c12or3 = c12 <+> c3
let _ = Printf.printf "cost of c12or3 = %d\n" @@ cost c12or3

let () = print_endline @@ apply c12or3 [1;2;3]