#load "pa_extend.cmo";;
#load "q_MLast.cmo" ;;

(* 
   pa_compr.ml:
   list comprehension  syntax extension with Camlp4
   see http://oandrieu.nerim.net/ocaml/index.html#pa_compr 
   © Olivier Andrieu, 2003
   
   uses some code by Anton Moscal
   http://caml.inria.fr/archives/200001/msg00066.html
*)

let rec is_irrefut_patt =
  function
    | <:patt< $lid:_$ >>
    | <:patt< () >>
    | <:patt< _ >> -> true
    | <:patt< ($x$ as $y$) >> -> is_irrefut_patt x && is_irrefut_patt y
    | <:patt< { $list:fpl$ } >> ->
	List.for_all (fun (_, p) -> is_irrefut_patt p) fpl
    | <:patt< ($p$ : $_$) >> -> is_irrefut_patt p
    | <:patt< ($list:pl$) >> -> List.for_all is_irrefut_patt pl
    | _ -> false 





let rec list_comp_hask _loc ?(acc= <:expr< [] >>) exp = function
  | [ `Guard <:expr< True >> ] -> <:expr< [ $exp$ :: $acc$ ] >>
  | [ qual ] ->
      list_comp_hask _loc ~acc exp [ qual; `Guard <:expr< True >> ]
  | `Generator (p, l) :: n when is_irrefut_patt p ->
      let r = list_comp_hask _loc ~acc:(<:expr< __acc__ >>) exp n in
      <:expr< List.fold_right (fun $p$ __acc__ -> $r$) $l$ $acc$ >>
  | `Generator (p, l) :: n ->
      let r = list_comp_hask _loc ~acc:(<:expr< __acc__ >>) exp n in
      <:expr< List.fold_right 
	(fun __arg__ -> fun __acc__ -> 
	  match __arg__ with [ $p$ -> $r$ | _ -> __acc__ ] ) $l$ $acc$ >>
  | `Let_binding (p, e) :: n ->
      let r = list_comp_hask _loc ~acc exp n in
      <:expr< let $p$ = $e$ in $r$ >>
  | `Guard g :: n ->
      let r = list_comp_hask _loc ~acc exp n in
      <:expr< if $g$ then $r$ else $acc$ >>
  | [] ->
      invalid_arg "list_comp_hask"
			     
      
      


EXTEND
  GLOBAL: Pcaml.expr ;

  Pcaml.expr: LEVEL "simple" 
    [ [ "[" ; "+" ; e = Pcaml.expr ; "|" ; 
	data = LIST1 qualifier SEP "|" ; "]" ->
        list_comp_hask _loc e data 
    ] ] ;

  qualifier:
    [ [ p = Pcaml.patt ; "<-" ; v = Pcaml.expr LEVEL "expr1" -> 
       `Generator (p, v)
      | "let" ; lb = Pcaml.let_binding -> 
       `Let_binding lb
      | "when" ; g = Pcaml.expr LEVEL "expr1" ->
       `Guard g ] ] ;

END

(* Local Variables: *)
(* compile-command: "ocamlc -pp camlp4o -c -I +camlp4 pa_compr.ml" *)
(* End: *)