(* Lifted map-reduce, non-nested loops *)
(* Since this is an expository code, we limit the type
   of loop variables to strings.
   Polymorphic loop variables are possible, but the encoding
   will be distracting. We obtain the polymorphism when we move to
   multiple prompts, in lifted_nested_mp.ml
*)

open Delimcc_simple;;

(* Preliminaries. Code for Sec 2.1 *)

(* A commutative monoid *)
(*--start-monoid--*)
type 'r monoid = {union : 'r -> 'r -> 'r;
                  empty : 'r}
(*--end-monoid--*)

(*--start-sum-monoid--*)
let sum_monoid = {union = (+); empty = 0}
(*--end-sum-monoid--*)
;;

(* The type of loop variables *)
type indiv = string;;

(* The sample domain of people of cardinality 100.
   We distinguish three special people
*)
let people = 
  let rec loop acc = function
    | 0 -> acc
    | n -> loop (("p" ^ string_of_int n) :: acc) (n-1)
  in loop ["alice"; "bob"; "carol"] 97
;;

(* Generic map-reduce *)
let map_reduce_gen domain monoid f =
  List.fold_left (fun acc x -> monoid.union (f x) acc) monoid.empty domain;;


let testmp1 = 
  map_reduce_gen people sum_monoid
    (fun x -> 1 + if x = "alice" then 0 else 1)
;;
(* 199 *)


(* Switch: the list of `cases' terminated with
   the default case. Compare with the result of HANSEI reification, the
   lazy tree of weighted choices.
   The `tail' of Switch can be made lazy; we omit this optimization
   for the sake of simplicity of the example.
*)
type 'r switch = Default of 'r | Case of indiv * 'r * 'r switch

(*--start-reflect--*)
let rec reflect : 'r switch -> (indiv -> 'r)
  = function
  | Default r -> fun _ -> r
  | Case (x,r,fsw) -> fun x' ->
      if x = x' then r else reflect fsw x'
(*--end-reflect--*)

let 1 = reflect (Case ("alice",1,Default 2)) "alice"
let 2 = reflect (Case ("alice",1,Default 2)) "bob"
;;

(* monoid.union n copies of r together *)

(*--start-repeat--*)
let rec repeat monoid r = function	
  | 0 -> monoid.empty
  | n when n mod 2 = 0 -> 
      repeat monoid (monoid.union r r) (n/2)
  | n ->
      monoid.union r (repeat monoid r (n-1))
(*--end-repeat--*)
;;
    
(* Map-reduce with the switch *)

(* val map_reduce_sw : 'a monoid -> 'a switch -> 'a *)
(*--start-mr-sw--*)
let map_reduce_sw monoid fsw =
  let rec loop card = function 
    | Default r -> repeat monoid r card
    | Case (_,r,tail) -> 
	monoid.union r (loop (card-1) tail)
  in loop 100 fsw
(*--end-mr-sw--*)
;;

let testmp2 = 
  map_reduce_sw sum_monoid
    (Case ("alice",1,Default 2))
;;
(* 199 *)




(* The `Logic' variable of type indiv *)
type var = Val of indiv | Var
;;

(* The interface to the reifier, which `runs' the lifted computation *)
type 'r req = Done of 'r | Compare of var * var * (bool -> 'r req);;


(* Equality over logic variables *)

let equ : var -> var -> bool = fun x y ->
  shift (fun k -> Compare (x,y,k));;

(* The reifier *)

(*--start-reify--*)
let reify (f : var -> 'r) : 'r switch =
  let var = Var in
  let rec loop_known w = function
    | Done r -> r
    | Compare (Var,Var,k) -> 
	loop_known w (k true)
    | Compare (Var,Val y,k) ->
	loop_known w (k (y = w)) 
    | Compare (Val y,Var,k) ->
	loop_known w (k (y = w)) 
    | Compare (Val x,Val y,k) ->
	loop_known w (k (x = y)) in
  let rec loop ws = function
    | Done r -> Default r
    | Compare (Var,Var,k) ->
	loop ws (k true)
    | Compare (Var,Val y,k) ->
	make_case ws k y
    | Compare (Val y,Var,k) ->
	make_case ws k y
    | Compare (Val x,Val y,k) ->
	loop ws (k (x = y)) 
  and make_case ws k y =
    if List.mem y ws then loop ws (k false)
    else Case (y,loop_known y (k true),
               loop (y::ws) (k false))
  in loop [] (reset (fun () -> Done (f var)))
(*--end-reify--*)
;;

let rec member x = function
  | [] -> false
  | (h::t) -> equ x h || member x t
;;


(* Tests *)

let Case ("alice", true, Case ("bob", true, Default false)) =
  reify (fun x -> member x [Val "alice"; Val "bob"; Val "alice"])
;;

(* Our running example here *)
let Case ("alice", 1, Default 2) =
  reify (fun x -> 1 + if equ x (Val "alice") then 0 else 1)
;;


let 97 =
(*--start-large-ex--*)
let xs = [Val "alice"; Val "bob"; Val "alice"]
in map_reduce_sw sum_monoid 
    (reify (fun x ->
     (if member x xs then 0 else 1) *
     (if equ x (Val "alice") then 4 else 1) *
     (if equ x (Val "carol") then 0 else 1)))
(*--end-large-ex--*)
;;

(* Demonstration why the code in this file is too simple to handle nested
   loops
*)

(* The result is not correct for the nested loop: there has to be
   99 * 99 = 9801
*)
let 9800 =
(*--start-nested-ex--*)
map_reduce_sw sum_monoid (reify (fun x ->
  map_reduce_sw sum_monoid (reify (fun y ->
    if equ x (Val "alice") ||
       equ y (Val "bob")
    then 0 else 1))))
(*--end-nested-ex--*)
;;


(* and this is totally wrong since we just can't compare two variables.
   Before we assumed that only one variable exist and hence any two
   variables are equal.
*)

let 0 =
(*--start-xy-ex--*)
map_reduce_sw sum_monoid (reify (fun x ->
  map_reduce_sw sum_monoid (reify (fun y ->
    if equ x y then 0 else 1))))
(*--end-xy-ex--*)
;;

let 0 =
map_reduce_sw sum_monoid (reify (fun x ->
 map_reduce_sw sum_monoid (reify (fun y ->
  if equ x (Val "alice") ||
     equ y (Val "bob") ||
     equ x y
  then 0 else 1))))
;;


print_endline "\nAll done";;