(* Lifted map-reduce, nested loops, multiple prompts *)
(* This is a different approach, more in the spirit of HANSEI.
   It partitions the work between reify and map_reduce_sw differently.
   The end results are identical to those achieved in lifted_nested_mp.ml.

   Although this code is not described in the paper, it should be easy
   to understand from the description in the paper. After all, it solves
   the same problem.
*)

open Delimcc;;

(* The `Logic' (or loop) variable of type 'i *)
(* We use the value of the type `generation' to impose the total order on
   logic variables. The variable of the greatest generation has been
   created by the most recent reify. Of course that holds only
   if no loop variable escapes from the scope of the reify that created
   it. Well, if it does, the corresponding prompt will no longer
   be set so that any attempt to use such a variable will lead to
   failure.
*)
type generation = int
let new_generation = 
  let cnt = ref 0 in
  fun () ->
    incr cnt; !cnt
;;

(* The variable and the interface to the reifier, 
   which `runs' the lifted computation *)
(* The request is delivered to the closest reify and always concerns
   its own variable. Therefore, we don't need to mention that
   variable in the request.
*)
type 'i var = Val of 'i | Var of generation * 'i req prompt
and  'i req = C_eq of 'i var * (bool -> 'i req)
;;


(* Polymorphic equality over loop variables *)

(*--start-equ--*)
let equ : 'i var -> 'i var -> bool = fun x y ->
  match (x,y) with
  | (Var (gx,px), Var (gy,py)) ->
      if gx > gy then
	shift px (fun k -> C_eq (y,k))
      else if gx < gy then
	shift py (fun k -> C_eq (x,k))
      else true
  | (Val x, Val y) -> x = y
  | (Var (_,px), y) | (y, Var (_,px)) ->
      shift px (fun k -> C_eq (y,k))
(*--end-equ--*)
;;

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

(* The result of reification: the lazy tree of `cases' terminated with
   the default case. Compare with the result of HANSEI reification, the
   lazy tree of weighted choices.
   It is a binary tree (we try to be general so to easily add other
   constraints).
*)
type ('i,'r) switch = Default of 'r 
  | Case of 'i var * ('i,'r) switch * ('i,'r) switch
;;


(* The reifier *)
(* Unlike the reifier in the paper, it does not try to solve the
   constraints. It merely reifies the loop body to a tree.
   It is quite like reify0 in HANSEI: it is very dumb.
   All smarts are in map_reduce, in the function that uses the
   reified data structure.
*)

let reify : ('i var -> 'r) -> ('i,'r) switch = fun f ->
  let var_prompt = new_prompt () in
  let var = Var (new_generation (), var_prompt) in
  let presult = new_prompt () in
  let rec loop = function
    | C_eq (y,k) ->
	abort presult 
	  (Case (y,
		 (push_prompt presult (fun () -> loop (k true))),
		 (push_prompt presult (fun () -> loop (k false)))))
  in 
  push_prompt presult (fun () ->
    loop (push_prompt var_prompt (fun () -> 
      abort presult (Default (f var)))))
;;


(* The map-reduce part *)

type 'i coll = {coll_member : 'i var -> bool;
                coll_cardinality : int};;


(* A sample collection used in all of the examples *)
(* The sample domain of people of cardinality 100. *)
let people =
  {coll_member = (fun v -> true);  (* Assume this is the whole domain *)
   coll_cardinality = 100;
 };;

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

let sum_monoid = {union = (+); empty = 0}
;;

let rec repeat monoid r = function	
  | 0 -> monoid.empty
  | 1 -> r
  | n when n mod 2 = 0 -> 
      repeat monoid (monoid.union r r) (n/2)
  | n ->
      monoid.union r (repeat monoid r (n-1))
;;

(* Fold over the collection coll. We apply the map function represented
   by the switch, and reduce the result with the commutative monoid.
   This is quite like explore of HANSEI
*)
let map_reduce_sw monoid coll fsw =
  let rec loop coll = function
    | Default x -> repeat monoid x coll.coll_cardinality
    | Case (v,br_t,br_f) -> 
	(match v with Var _ -> print_endline "case-var" 
	|  _ -> print_endline "case-val");
	if coll.coll_member v 
	then 
	  monoid.union 
	    (loop {coll_member = (fun x -> equ x v); coll_cardinality = 1}
	       br_t)
	    (loop {coll_member = 
		   (fun x -> if equ x v then false else coll.coll_member x);
		   coll_cardinality = coll.coll_cardinality - 1}
	       br_f)
	 else
	  loop coll br_f
  in loop coll fsw;;

let sum_over coll = map_reduce_sw sum_monoid coll;;

(* Many tests follow *)

let 100 = sum_over people (reify (fun x -> 1));;

let 101 = sum_over people (reify (fun x -> 
  if equ x (Val "alice") then 2 else 1));;


let 103 = sum_over people (reify (fun x -> 
  if equ x (Val "alice") then 2 else if equ x (Val "bob") then 3 else 1));;

let 203 = sum_over people (reify (fun x -> 
  if equ x (Val "alice") then 
    (if equ (Val "alice") x then 102 else -100)
  else if equ x (Val "bob") then 3 else 1));;

let 1 = sum_over people (reify (fun x -> 
  if equ x (Val "alice") then 
    (if equ (Val "bob") x then 102 else -100)
  else if equ x (Val "bob") then 3 else 1));;

(* Double-nested loops *)
let 10000 = 
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 1))));;

let 9900 =
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
       if equ y (Val "alice") then 0 else 1))));;

let 9900 =
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
       if equ x (Val "alice") then 0 else 1))));;


let 9900 =
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
       if equ y x then 0 else 1))));;

let 9900 =
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
       if equ x y then 0 else 1))));;

let 29800 =  (* 99 * (99 + 200) + (99+100) *)
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
       if equ x y then 
	 (if equ x (Val "alice") then
	     100 else 200)
	 else 1))));;

let 29800 =  (* 99 * (99 + 200) + (99+100) *)
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
       if equ x y then 
	 (if equ x (Val "alice") then
	   (if equ y (Val "alice") then 100 else -100000)
	     else 200)
	 else 1))));;

(* Our running example in the paper *)
let 9703 =
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      if equ x (Val "alice") ||
         equ y (Val "bob") ||
	 equ x y
      then 0 else 1))))
;;


(* Triple-nested loops *)
let 990000 =      
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ x z then 0 else 1))))));;

let 990000 =      
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ z x then 0 else 1))))));;

let 990000 =      
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ x y then 0 else 1))))));;

let 990000 =      
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ y x then 0 else 1))))));;

let 999900 = (* 100 * (9900 + 99) *)
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ y x && equ x z then 0 else 1))))));;


let 999900 = (* 100 * (9900 + 99) *)
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ y x && equ x z && equ y z then 0 else 1))))));;

let 980100 = 
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ y x || equ x z then 0 else 1))))));;

let 970200 = (* 100 * 99 * 98 *)
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y -> 
      sum_over people (reify (fun z -> 
	if equ x z || equ z y || equ x y then 0 else 1))))));;


(*
let range_coll m n = 
  assert (n >= m);
  {coll_member = (fun v -> true); (* This assumes the collection is the whole
				     domain *)
   coll_cardinality = n - m + 1;}
;;
*)

(*
  Testing the complexity of dealing with the conjunction of disequality
  constraints in the star form.

  The conclusion of our tests below: if the user writes nested loops 
  explicitly, the user has selected the variable order -- which
  may be far from optimal. If the wrong order is selected, things can get
  very slow.
  No surprise here: when working with BDDs, the user, too, has to
  select the variable order. If the wrong (for the problem) order
  is selected, computations can become really, really slow.
*)

(* Time the execution *)
let timeit thunk =
  let time_start = Sys.time () in
  let r = thunk () in
  Printf.printf "\nTime spent: %g sec\n" (Sys.time () -. time_start);
  r;;

(*
let map_reduce_sw monoid coll fsw =
  let rec loop coll = function
    | Default x ->
	Printf.printf "repeating with counter %d\n" coll.coll_cardinality;
	repeat monoid x coll.coll_cardinality
    | Case (v,br_t,br_f) -> 
	if coll.coll_member v 
	then 
	  monoid.union 
	    (loop {coll_member = (fun x -> equ x v); coll_cardinality = 1}
	       br_t)
	    (loop {coll_member = 
		   (fun x -> if equ x v then false else coll.coll_member x);
		   coll_cardinality = coll.coll_cardinality - 1}
	       br_f)
	 else
	  loop coll br_f
  in loop coll fsw;;

  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y1 -> 
    sum_over people (reify (fun y2 -> 
	if equ x y1 || equ x y2 
	then 0 else 1))))))
;;

repeating with counter 99
repeating with counter 1
repeating with counter 100
repeating with counter 99
repeating with counter 1
repeating with counter 100

The result shows that we really need laziness:

    sum_over people (reify (fun y1 -> 
    sum_over people (reify (fun y2 -> 
    sum_over people (reify (fun x -> 
	if equ x y1 || equ x y2 
	then 0 else 1))))))
;;

repeating with counter 98
repeating with counter 1
repeating with counter 1
repeating with counter 99
repeating with counter 1
repeating with counter 99
repeating with counter 1
repeating with counter 100


timeit (fun () ->
  sum_over people (reify (fun x -> 
    sum_over people (reify (fun y1 -> 
    sum_over people (reify (fun y2 -> 
    sum_over people (reify (fun y3 -> 
	if equ x y1 || equ x y2 || equ x y3
	then 0 else 1)))))))))
;;

repeating with counter 99
repeating with counter 1
repeating with counter 100
repeating with counter 99
repeating with counter 1
repeating with counter 100
repeating with counter 100
repeating with counter 99
repeating with counter 1
repeating with counter 100


timeit (fun () ->
    sum_over people (reify (fun y1 -> 
    sum_over people (reify (fun y2 -> 
    sum_over people (reify (fun y3 -> 
  sum_over people (reify (fun x -> 
	if equ x y1 || equ x y2 || equ x y3
	then 0 else 1)))))))))
;;


repeating with counter 97
repeating with counter 1
repeating with counter 1
repeating with counter 1
repeating with counter 98
repeating with counter 1
repeating with counter 1
repeating with counter 98
repeating with counter 1
repeating with counter 1
repeating with counter 98
repeating with counter 1
repeating with counter 1
repeating with counter 99
repeating with counter 1
repeating with counter 98
repeating with counter 1
repeating with counter 1
repeating with counter 99
repeating with counter 1
repeating with counter 99
repeating with counter 1
repeating with counter 99
repeating with counter 1
repeating with counter 100

for the loop y1..y2: 6 vs 8
for the loop y1..y3: 10 vs 25
for the loop y1..y4: 15 vs 230


For comparison: counting the number of invocations of repeat in the loops
(essentially, the number of Default leaves) in lifted_nested_mp.ml
The latter code handles Default for the `known' case specially and
hence it is not printed below.

let map_reduce_sw monoid fsw =
  let rec loop card = function 
    | Default r -> 
	Printf.printf "repeating with counter %d\n" card;
    repeat monoid r card
    | Case (_,r,tail) -> 
	monoid.union r (loop (card-1) tail) 
    (* assume always the member of the domain *)
  in loop 100 fsw
;;

  map_reduce_sw sum_monoid (reify (fun x -> 
  map_reduce_sw sum_monoid (reify (fun y1 -> 
  map_reduce_sw sum_monoid (reify (fun y2 -> 
	if equ x y1 || equ x y2 
	then 0 else 1))))))
;;

repeating with counter 99
repeating with counter 100
repeating with counter 99
repeating with counter 100
- : int = 980100

  map_reduce_sw sum_monoid (reify (fun y1 -> 
  map_reduce_sw sum_monoid (reify (fun y2 -> 
  map_reduce_sw sum_monoid (reify (fun x -> 
	if equ x y1 || equ x y2 
	then 0 else 1))))))
;;

repeating with counter 98
repeating with counter 99
repeating with counter 99
repeating with counter 100

  map_reduce_sw sum_monoid (reify (fun x -> 
  map_reduce_sw sum_monoid (reify (fun y1 -> 
  map_reduce_sw sum_monoid (reify (fun y2 -> 
  map_reduce_sw sum_monoid (reify (fun y3 -> 
	if equ x y1 || equ x y2 || equ x y3
	then 0 else 1))))))))
;;

repeating with counter 99
repeating with counter 100
repeating with counter 99
repeating with counter 100
repeating with counter 100
repeating with counter 99
repeating with counter 100
- : int = 97029900


  map_reduce_sw sum_monoid (reify (fun y1 -> 
  map_reduce_sw sum_monoid (reify (fun y2 -> 
  map_reduce_sw sum_monoid (reify (fun y3 -> 
  map_reduce_sw sum_monoid (reify (fun x -> 
	if equ x y1 || equ x y2 || equ x y3
	then 0 else 1))))))))
;;

for the loop y1..y2: 4 vs 4
for the loop y1..y3: 7 vs 9
for the loop y1..y4: 11 vs 24

*)

print_endline "\nAll done";;