open Pcaml

let id_prefix = "__fr_"

type dependent = Behavior | Event

let dependents = Stack.create ()

let cur_context () =
	try Stack.top dependents
	with Stack.Empty ->
		failwith "[@expr] or [%expr] occurs outside of [let @] or [lift]"

let with_new_context f =
	Stack.push (ref []) dependents;
	try let r = f () in ignore (Stack.pop dependents); r
	with e -> ignore (Stack.pop dependents); raise e

(* Create a unique ID for the expression in the current context and add it
   to the list of dependents.  If the expression was already encountered,
   return its previously assigned ID. *)
let gen_id dtype context expr =
	let loc = MLast.loc_of_expr expr in
	let reloc = expr_reloc (fun _ -> loc) Lexing.dummy_pos in
	let reloc_expr = reloc expr in
	try fst (List.find (fun (id, (dtype', expr')) ->
		dtype = dtype' && reloc_expr = reloc expr') !context)
	with Not_found ->
		(* should this use rename_id? *)
		let id = id_prefix ^ string_of_int (List.length !context) in
		context := (id, (dtype, expr)) :: !context; id

(* Add an expression to the list of dependents and return an expression
   referencing its current value. *)
let gen_dependent _loc dtype context expr =
	try <:expr< $lid:gen_id dtype context expr$ >>
	with e -> Stdpp.raise_with_loc _loc e

let extract_dependents dtype dependents =
	List.rev (List.fold_left (fun l (id, (dtype', expr)) ->
		if dtype' = dtype then (id, expr) :: l else l) [] dependents)

(* Create a signal expression from an expression and a dependency map. *)
let gen_behavior _loc (expr, dependents) =
	let merge_dependents a b = match a, b with
		| (Behavior, a), (Behavior, b) ->
			Behavior, <:expr< Fr.zip_b $a$ $b$ >>
		| (Event, a), (Event, b) ->
			Event, <:expr< Fr.zip_e $a$ $b$ >>
		| (Behavior, a), (Event, b) ->
			Event, <:expr< let __fr_b = $b$ in
				Fr.zip_e (Fr.snapshot_b $a$ __fr_b) __fr_b >>
		| (Event, a), (Behavior, b) ->
			Event, <:expr< let __fr_a = $a$ in
				Fr.zip_e __fr_a (Fr.snapshot_b $b$ __fr_a) >>
	in
	match dependents with
	| [] ->
		<:expr< Fr.lift0 $expr$ >>
	| hd :: tl ->
		let bindings =
			List.fold_left
				(fun p (id, d) -> <:patt< ($p$, $lid:id$) >>)
				<:patt< $lid:fst hd$ >> tl
		and signal =
			List.fold_left
				(fun s (id, d) -> merge_dependents s d)
				(snd hd) tl in
		match signal with
		| Event, event ->
			<:expr< Fr.map_e (fun $bindings$ -> $expr$) $event$ >>
		| Behavior, behavior ->
			<:expr< Fr.lift1 (fun $bindings$ -> $expr$) $behavior$ >>

(* Rule to parse an expr in a new dependency context.  Returns a tuple of
   the expression and its dependents' expr -> ID map. *)
let dependent_expr =
	Grammar.Entry.of_parser gram "dependent_expr"
		(fun s ->
			with_new_context (fun () ->
				let e = Grammar.Entry.parse_token expr s in
				e, !(cur_context ())))

let simple_dependent_expr =
	let simple_expr = Grammar.Entry.create gram "simple_expr" in
	EXTEND
		GLOBAL: simple_expr;
		
		simple_expr: [[ e = expr LEVEL "." -> e ]];
	END;
	Grammar.Entry.of_parser gram "simple_dependent_expr"
		(fun s ->
			with_new_context (fun () ->
				let e = Grammar.Entry.parse_token simple_expr s in
				e, !(cur_context ())))

EXTEND
	GLOBAL: expr let_binding;

	expr: LEVEL "expr1" [
		[ "on"; event = expr LEVEL "simple";
			args = LIST0 patt LEVEL "simple"; "->"; body = expr LEVEL "top" ->
			let f =
				List.fold_right (fun arg body -> 
					<:expr< fun $arg$ -> $body$ >>) args body
			in <:expr< Fr.map_e $f$ $event$ >> ]
	];

	expr: LEVEL "apply" [
		[ "lift"; de = simple_dependent_expr -> gen_behavior _loc de ]
	];

	let_binding: [
		[ "@"; p = patt LEVEL "simple"; e = dependent_fun_binding ->
			p, e ]
	];
	
	dependent_fun_binding: [ LEFTA
		[ "="; de = dependent_expr ->
			gen_behavior _loc de ]
		| [ p = patt LEVEL "simple"; e = dependent_fun_binding ->
			<:expr< fun $p$ -> $e$ >> ]
	];

	expr: LEVEL "simple" [
		[ "@"; e = expr LEVEL "simple" ->
			try gen_dependent _loc Behavior (cur_context ()) e
			with exn -> Stdpp.raise_with_loc _loc exn ]
		| [ "%"; e = expr LEVEL "simple" ->
			try gen_dependent _loc Event (cur_context ()) e
			with exn -> Stdpp.raise_with_loc _loc exn ]
	];
END;;