open Names
open Location
open Heptagon
open Hept_mapfold

type error =
  | Etoo_much_async

let message loc kind =
  begin match kind with
    | Etoo_much_async ->
        Format.eprintf "%aInvalid async nesting.@."
          print_location loc
  end;
  raise Errors.Error

(* Compute the set of nodes that use at least one async call. *)
let exp_callers funs (callers, current) e =
  let e, (callers, current) = Hept_mapfold.exp funs (callers, current) e in
  match e.e_desc with
    | Eapp({ a_op = Easync _ }, _, _) ->
      e, (QualSet.add (Option.get current) callers, current)
    (* TODO(Arduino): Eiterator *)
    | _ -> e, (callers, current)

let node_dec_callers funs (callers, _) n =
  Hept_mapfold.node_dec funs (callers, Some n.n_name) n

let funs_callers =
  { Hept_mapfold.defaults with
      node_dec = node_dec_callers;
      exp = exp_callers }

(* Ensure that no node using an async call is called more than once. *)
let exp_async funs (callers, calls) e =
  let e, (callers, calls) = Hept_mapfold.exp funs (callers, calls) e in
  match e.e_desc with
    | Eapp({ a_op = Easync (name, _) }, _, _) ->
      let caller = QualSet.mem name callers in
      let exists = QualSet.mem name calls in
      if caller && exists then
        message e.e_loc Etoo_much_async
      else
        e, (callers, QualSet.add name calls)
    (* TODO(Arduino): Eiterator *)
    | _ -> e, (callers, calls)

let funs_async =
  { Hept_mapfold.defaults with
      exp = exp_async }

let program p =
  let _, (callers, _) =
    Hept_mapfold.program_it funs_callers (QualSet.empty, None) p
  in
  let _ = Hept_mapfold.program_it funs_async (callers, QualSet.empty) p in
  p