(***********************************************************************)
(*                                                                     *)
(*                          OCamlP3l                                   *)
(*                                                                     *)
(* (C) 2004-2007                                                       *)
(*             Roberto Di Cosmo (dicosmo@dicosmo.org)                  *)
(*             Zheng Li (zli@lip6.fr)                                  *)
(*             Pierre Weis (Pierre.Weis@inria.fr)                      *)
(*             Francois Clement (Francois.Clement@inria.fr)            *)
(*                                                                     *)
(* Based on original Ocaml P3L System                                  *)
(* (C) 1997 by                                                         *)
(*             Roberto Di Cosmo (dicosmo@ens.fr)                       *)
(*             Marco Danelutto (marcod@di.unipi.it)                    *)
(*             Xavier Leroy  (Xavier.Leroy@inria.fr)                   *)
(*             Susanna Pelagatti (susanna@di.unipi.it)                 *)
(*                                                                     *)
(* This program is free software; you can redistribute it and/or       *)
(* modify it under the terms of the GNU Library General Public License *)
(* as published by the Free Software Foundation; either version 2      *)
(* of the License, or (at your option) any later version.              *)
(*                                                                     *)
(* This program is distributed in the hope that it will be useful,     *)
(* but WITHOUT ANY WARRANTY; without even the implied warranty of      *)
(* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the       *)
(* GNU Library General Public License for more details.                *)
(*                                                                     *)
(***********************************************************************)

(* $Id: seqp3l.ml,v 1.20 2007-01-23 15:50:04 weis Exp $ *)

open P3lstream;;
open Streams;;

type ('a, 'b) io = unit;;

type ('a,'b) skel = ?ecol:color -> unit -> 'a s -> 'b s;;

let rindex n =
  Random.self_init ();
  let st = ref (Random.get_state ()) in
  fun _ -> 
    let () = Random.set_state !st in
    let num = Random.int n in
    let () = st := Random.get_state () in
    num
;;

let sindex n =
  let r = ref (-1) in
  fun _ ->
    r := (!r + 1) mod n; !r
;;

let seq ?(col:color = 0) (f:unit -> 'a -> 'b) =
  fun ?(ecol:color = 0) ->
    let fcol = match col with 0 -> ecol | _ -> col in
    fun () -> 
	  let pf =
	    let inst_f = f () in
	    function 
	      | Pac(t,v) -> Pac (t, lazy (inst_f (Lazy.force v)))
	      | Sign _ as x  -> x in
	  streamer(pf, fcol)
;;


(* Actually, we don't need a pipe. This is an enhanced version serves
   the compatability purpose *)
let pipeline
    (spna: ?ecol:color -> unit -> 'a -> 'b) 
    (spnb: ?ecol:color -> unit -> 'b -> 'c) =
  fun ?(ecol:color=0) () ->
    let inst_spna = spna ~ecol:ecol () 
    and inst_spnb = spnb ~ecol:ecol () in
    fun s -> inst_spnb (inst_spna s)
;;

let ( ||| ) = pipeline;;

let farm_template: (int -> 'a pkg -> int) -> (int -> unit -> int) ->
  ?col:color -> ?colv:color list -> ('a, 'b) skel * int -> ('a, 'b) skel =
  fun emit_fun collect_fun ->
    fun ?(col: color = 0) ?(colv: color list = []) 
      ((spnx: ('a, 'b) skel), (n: int)) ->
	fun ?(ecol: color = 0) ->
	  let fcol = match col with 0 -> ecol | _ -> col in
	  let fcola = match colv with 
	    | [] -> Array.make n fcol
	    | _ -> Array.init n
		(fun i -> match List.nth colv i with 0 -> fcol | x -> x) in
	  fun () ->
	    let inst_spnxa = 
	      Array.init n (fun i -> spnx ~ecol:fcola.(i) ()) in
	    fun s ->
	      let sai = split (n, emit_fun n, fcol) s in
	      let sao = 
		Array.init n (fun i -> (inst_spnxa.(i) sai.(i))) in
	      combine (collect_fun n, fcol) sao
;;


let farm ?(col=0) ?(colv=[]) (spn, n) = 
  farm_template sindex rindex ~col ~colv (spn,n) ;;

let sfarm ?(col=0) ?(colv=[]) (spn, n) = 
  farm_template sindex sindex ~col ~colv (spn,n) ;;


let do_while ?(col=0) ((cf: 'a pkg -> bool), (spnx: ('a, 'a) skel)) = 
  let mcf = fun p -> if cf p then 1 else 0 in
  fun ?(ecol=0) ->
    let fcol = match col with 0 -> ecol | _ -> col in
    fun () ->
      let inst_spnx = spnx ~ecol:fcol () in
      let spnr = recur (1,1) 
	(fun sa -> split (2, mcf, col) 
	   (inst_spnx (combine (rindex 2, fcol) sa))) in
      fun s -> 
	let si = [|s|] in
	let so =  spnr si in
	so.(0)
;;


let while_do ?(col=0) ((cf: 'a pkg -> bool), (spnx: ('a, 'a) skel)) =
  let mcf = fun p -> if cf p then 1 else 0 in
  fun ?(ecol=0) ->
    let fcol = match col with 0 -> ecol | _ -> col in
    fun () ->
      let inst_spnx = spnx ~ecol:fcol () in
      let spnr = recur (1,1)
	(fun sa ->
	   (function 
	      | [|s0;s1|] -> [|s0; inst_spnx s1|] 
	      | _ -> assert false)
	     (split (2, mcf, fcol) 
		(combine (rindex 2, fcol) sa))) in
      fun s -> 
	let si = [|s|] in
	let so =  spnr si in
	so.(0)
;;


let loop ?(col=0) (cf, spnx) =
  let pcf = function
    | Pac(_,v) -> cf (Lazy.force v)
    | Sign _ -> false in
  do_while ~col (pcf, spnx);;


(* To be compatible with the old version, the inside farm of mapvector
   use sfarm, in which the emittor sequently spread packge to the
   workers and collector get them in the same order *)
let mapvector ?(col=0) ?(colv=[]) ((spnx: ('a, 'b) skel), (n: int)) =
  fun ?(ecol=0) ->
    let fcol = match col with 0 -> ecol | _ -> col in
    let df = function
      | Sign _ as x ->  [|x|]
      | Pac(t,v) -> 
	      let a = Lazy.force v in
	      let l = Array.length a - 1 in
	      Array.mapi (fun i x -> Pac((i,l)::t, lazy x)) a in
    let uf = function
      | (Pac((_, lb)::tb, _))::t as pl -> 
          if lb = List.length pl - 1 then
            let nl = List.sort 
              ( fun x y -> match x, y with
                  | Pac((x1,_)::_,_), Pac((x2,_)::_,_) -> compare x1 x2
                  | _ -> assert false
              ) pl in
            let ev = Array.of_list
              (List.map 
                 (function
                    | Pac(_,v) -> Lazy.force v
                    | _ -> assert false
                 ) nl
              ) in
            Some (Pac(tb,lazy ev),[]) 
          else None
      | _ -> assert false in
    let subfarm = sfarm ~col ~colv (spnx, n) in
    fun () ->
      let inst_subfarm = subfarm ~ecol:ecol () in
      fun s -> unif (uf,fcol) (inst_subfarm (dive (df,fcol) s))
;;


let reducevector 
    ?(col=0) ?(colv=[]) ((spnx: (('a * 'a), 'a) skel), (n: int)) = 
  fun ?(ecol=0) ->
    let fcol = match col with 0 -> ecol | _ -> col in
    let df = function
      | Sign _ as x ->  [|x|]
      | Pac(t,v) -> 
	      let a = Lazy.force v in
	      let l = Array.length a in
          Array.map (fun x -> Pac((1,l)::t, lazy x)) a in
    let cf = function
      | Sign _ -> false
      | Pac(t,v) -> 
	      match t with
	        | (x,l)::_ when x >= l -> false
	        | _ -> true  in
    let uf = function
      | (Pac((x1,l1)::t1, v1))::(Pac((x2,l2)::t2, v2))::t 
          when (l1 = l2) ->
          let x = x1 + x2 in
          let ev = lazy (Lazy.force v1, Lazy.force v2) in
          Some ((Pac((x,l1)::t1, ev)), t)
      | (Pac _)::[] -> None
      | _ -> assert false in
    let sf = function 
      | Pac (t,v) -> Pac (List.tl t, v)
      | x -> x in
    let subfarm = farm ~col ~colv (spnx, n) in
    let subloop = while_do ~col
      (cf, fun ?(ecol=0) () ->
	     let inst_subfarm = subfarm ~ecol () in
	     fun s -> inst_subfarm (unif(uf, fcol) s)) in
    fun () ->
      let inst_subloop = subloop ~ecol () in
      fun s ->
	    streamer (sf, fcol) (inst_subloop (dive (df, fcol) s))
;;


let parfun (expr: unit -> ('a, 'b) skel) = expr () ();;

let pardo expr =
  Arg.parse
    (Command_options.all ())
    (fun s -> prerr_endline ("Unknown anonymous argument " ^ s); exit 2)
    ("Usage: " ^ Sys.argv.(0) ^ " [options]");
  expr ()
;;