(*pp camlp4o q_MLast.cmo pa_extend.cmo *)

(* The syntax extension for list comprehension.

   HOW TO COMPILE
   $ ocamlc -c -I +camlp4 -pp "camlp4o q_MLast.cmo pa_extend.cmo -loc loc" listcompr.ml

   HOW TO USE
   # #load "camlp4o.cma";;
   # #load "listcompr.cmo";;
   # [ x + y | x <- [1;2;3] when x > 1; y <- [4;5;6] ];;

   The original code was posted by Anton Moscal to caml-list:
   
     Date: 2000-01-21 (08:43)
     Subject: Re: Q: camlp4 use?
*)


open Stdpp
open Pcaml

let rec is_irrefut_patt = function
  | <:patt< $lid:_$ >> -> true
  | <:patt< () >> -> true
  | <:patt< _ >> -> true
  | <:patt< ($x$ as $_$) >> -> is_irrefut_patt x
  | <: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 compr loc out is =
  let gen_patt_fun p whn expr default =
    match (whn, is_irrefut_patt p) with
    | None, true -> <:expr< fun acc $p$ -> $expr$ >>
    | _, _ -> 
        <:expr< fun acc -> 
        fun [$list:[p,whn,expr; <:patt<_>>,None,<:expr< acc >>]$ ] >>
  in
  let empty_list = <:expr< [] >> in
  let (_, inp, _) = List.hd is in
  let gen_fold fn acc inp = <:expr<List.fold_left $fn$ $acc$ $inp$>> in
  let rec gen_fun = function
    | ((p, inp, whn)::tail) ->
        begin let cons a b = <:expr< [$a$::$b$] >> in
        let acc = <:expr<acc>> in
        match tail with
        | [] -> gen_patt_fun p whn (cons (List.hd out) acc) acc
        | (_, inp', _) :: _ ->
             gen_patt_fun p whn (gen_fold (gen_fun tail) acc inp') acc
        end

    | _ ->
        failwith "Syn.gen_fun"
  in
  let rev_res = <:expr< $gen_fold (gen_fun is) empty_list inp$ >> in
  <:expr< List.rev $rev_res$ >>

let mklistexp loc last =
  let rec loop top =
    function
      | [] ->
       (match last with
        | Some e -> e
        | None -> <:expr< [] >> )
      | e1 :: el ->
        <:expr< [$e1$ :: $loop false el$] >>
  in loop true

let is_str_type =
  Grammar.Entry.of_parser gram "operator" (fun strm ->
    match Stream.peek strm with
    | Some ("", "type")
    | Some ("", "open")
    | Some ("", "class")
    | Some ("", "external")
    | Some ("", "exception") -> ()
    | _ -> raise Stream.Failure
          )

let _ = Gramext.warning_verbose := false
EXTEND
   GLOBAL: expr str_item patt let_binding ctyp;

   whn: [["when"; cond = expr LEVEL "expr1" -> cond]];

  item: [[p = patt; "<-" ; inp = expr LEVEL "expr1"; whn = OPT whn -> 
            (p, inp, whn)
        ]];

  expr1_semi_list:
    [[ e = expr LEVEL "expr1"; ";"; el = expr1_semi_list -> e :: el
     | e = expr LEVEL "expr1"; ";" -> [e]
     | e = expr LEVEL "expr1" -> [e] ] ]
  ;

  expr: LEVEL "simple"
    [[ "["; "]" -> <:expr< [] >>
     | "["; el = expr1_semi_list; inp = ["]" -> None | "|"; 
       is = LIST1 item SEP ";"; "]" -> Some is] -> (
       match inp with
       | None -> <:expr< $mklistexp loc None el$ >>
       | Some is -> compr loc el is) ]];

  expr: LEVEL "expr1"
  [["let"; is_str_type; s = str_item; "in"; e = expr LEVEL "top" ->
      <:expr< let module M_M_temp = struct $s$; value _res = $e$; end in
              M_M_temp._res >>
  ]];

END;;
let _ = Gramext.warning_verbose := true