open Pcaml;;

type exprT = MLast.expr Grammar.Entry.e
and  pattT = MLast.patt Grammar.Entry.e;;

let let_gen = Grammar.Entry.create gram "let_general";;
let expr_where = Grammar.Entry.create gram "expr_where";;
let str_item_semi = Grammar.Entry.create gram "str_item_semi";;
let implem_pre = Grammar.Entry.create gram "implem_pre";;
let let_lazy_binding = Grammar.Entry.create gram "let_lazy_binding";;

let rec insert ?(word="WHERE") elems lst =
    match lst with [] -> elems
	| (<:str_item< $str:name$ >>,_ as w)::tl when name=word -> elems@(w::tl)
	| hd::tl -> hd::insert ~word elems tl

let is_header = function <:str_item< value $lid:_$ = LazyX . nondef () >>, _ -> true | _ -> false

EXTEND
    implem:
	[[ (l,stopped) = implem_pre -> 
	    let lst1,lst2 = List.partition is_header l in
	    (List.rev (insert ~word:"HEADER" lst1 lst2),stopped) ]];
    implem_pre:
	[[ "WHERE"; (sil,stopped) = SELF ->
			    ((<:str_item< $str:"WHERE"$ >>,loc)::sil, stopped)
	| "HEADER"; (sil,stopped) = SELF ->
			    (insert [(<:str_item< $str:"HEADER"$ >>,loc)] sil, stopped)
	| si = str_item_semi; OPT ";;"; (sil, stopped) = SELF ->
			    (insert si sil, stopped)
	| "#"; n = LIDENT; dp = OPT expr; ";;" ->
	      ([(<:str_item< # $n$ $opt:dp$ >>, loc)], true)
	| EOI -> ([], false) ]];
    str_item_semi:
	[[ si = str_item; OPT ";;" ->
	    match si with 
		<:str_item< LazyX.set $lid:dest$ $e$ >> as si ->
		    [(<:str_item< value $lid:dest$ = LazyX . nondef () >>,loc);(si, loc)];
		| _ -> [(si, loc)] ]];
    str_item: LEVEL "top"
	[[ "let"; "lazy"; lb = let_lazy_binding -> lb
	| "lazy"; lb = let_lazy_binding -> lb]];
    let_lazy_binding:
	[[ lb = let_binding ->
	    let tr (p,e) = match p with
		<:patt< $lid:name$ >> ->
		    <:str_item< LazyX.set $lid:name$ ( lazy $e$ ) >>
		| _ -> raise Stream.Failure in
	    tr lb ]];
    expr: BEFORE "top"
	[[ e = SELF; (ifRec,l) = expr_where -> 
	    <:expr< let $rec:ifRec!=None$ $list:l$ in $e$ >> ]];
    expr_where:
	[["where"; _ = OPT "begin"; (ifRec,l) = let_gen; "end" ->
	    (ifRec,l)
	| "where"; "{"; (ifRec,l) = let_gen; "}" -> (ifRec,l) ]];
    let_gen: [[r = OPT "rec"; l = LIST1 let_binding SEP "and" -> (r,l)]];
    module_expr:
	[[ "struct"; st = LIST0 [ "WHERE" -> `Where
			    | s = str_item; OPT ";;" -> `Item s ]; "end" ->
		    let rec f ((il::ils) as ilS:'a list list) =
			    function [] -> ilS
				    | (`Item i::l) -> f ((i::il)::ils) l
				    | (`Where::l) -> f ([]::ilS) l in
		    let s = <:str_item< $lid:"b"$ = $lid:"a"$ >> in
		    <:module_expr< struct $list:(List.concat (List.map List.rev (f [[]] st)))$ end >>
	]];
END;;