(* goaljobs
 * Copyright (C) 2013 Red Hat Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *)

(* For general information about camlp4, see:
 * http://brion.inria.fr/gallium/index.php/Camlp4
 *
 * For information about quotations, see:
 * http://brion.inria.fr/gallium/index.php/Quotation
 *)

open Printf

open Camlp4.PreCast
open Syntax
open Ast

let locfail _loc msg = Loc.raise _loc (Failure msg)

(* 'expr' is a function expression (RHS of a binding).  It has the
 * form 'fun a -> fun b -> ...'.  Return the parameters and the body of
 * the function.
 *)
let rec function_parameters = function
  (* fun patt -> expr *)
  | ExFun (_loc, McArr (_ploc, param, ExNil _, expr)) ->
    let params, body = function_parameters expr in
    ((_ploc, param) :: params), body
  (* patt when expr -> expr *)
  | ExFun (_loc, McArr (_, _, _, expr)) ->
    locfail _loc "not supported: goal function uses when-clause"
  | body -> [], body

(* Define one or more 'let [rec] goal ... [and ...]' expressions.
 *
 * 'r' is Some _ if the rec keyword was defined.  'lets' is the list
 * of let statements.
 *)
let generate_let_goal _loc (r : rec_flag) (lets : binding) =
  let autopublish = ref [] in

  (* lets might be a single binding, or multiple bindings using BiAnd
   * ('let .. and').  Rewrite each individual goal in the list.
   *)
  let rec rewrite = function
    | BiNil _loc -> BiNil _loc

    (* let goal left = ... and right = ... *)
    | BiAnd (_loc, left, right) ->
      BiAnd (_loc, rewrite left, rewrite right)

    (* let goal name = expr *)
    | BiEq (_loc, PaId (_, (IdLid (_, name))), expr) ->
      (* Rename the function to goal_<name>. *)
      let gname = "goal_" ^ name in

      (* Convert loc to string for goalloc. *)
      let goalloc = Loc.to_string _loc in

      (* Split the function into parameters and body. *)
      let params, body = function_parameters expr in

      if params = [] then
        locfail _loc "goal must have some parameters; you probably want to put '()' here";

      (* Is it a "zero-parameters" automatically published goal?  What
       * this really means is it has exactly one unit parameter.
       *)
      (match params with
      | [ _, PaId (_, IdUid (_, "()")) ] ->
        autopublish := name :: !autopublish
      | _ -> ()
      );

      (* Put a try-clause around the body. *)
      let body = <:expr<
        (* Define a goal name which the body may use. *)
        let goalname = $str:name$ in

        (* Source location. *)
        let goalloc = $str:goalloc$ in

        (* Define onsuccess, onrun, onfail functions that the body may call. *)
        let _on, _call_on =
          let _on fns f = fns := f :: !fns in
          let _call_on fns a = List.iter (fun f -> f a) !fns in
          _on, _call_on
        in
        let onfail, _call_onfails =
          let fns = ref [] in (_on fns), (_call_on fns)
        in
        let onrun, _call_onruns =
          let fns = ref [] in (_on fns), (_call_on fns)
        in
        let onsuccess, _call_onsuccesses =
          let fns = ref [] in (_on fns), (_call_on fns)
        in

        try
          $body$ ;
          _call_onruns ();
          _call_onsuccesses ();

          (* Avoid a compiler warning: *)
          ignore (goalname); ignore (goalloc)
        with
        (* target() within the body may raise Goal_OK meaning that the
         * goal should be short-circuited.  We return as if it's an
         * ordinary function exit.
         *)
        | Goal_result Goal_OK ->
          _call_onsuccesses ();
          ()
        | exn ->
          _call_onfails exn;
          raise exn
      >> in

      (* Recreate the function with parameters. *)
      let expr =
        List.fold_right (
          fun (_ploc, param) rest ->
            ExFun (_loc, McArr (_ploc, param, ExNil _ploc, rest))
        ) params body in

      <:binding< $lid:gname$ = $expr$ >>
    | _ ->
      locfail _loc "cannot parse 'let goal' expression"
  in
  let lets = rewrite lets in

  (* let [rec] ... and ... in () *)
  let stmts = Ast.StVal (_loc, r, lets) in

  (* Auto-published goals. *)
  List.fold_left (
    fun stmt name ->
      let publish_name =
        let gname = "goal_" ^ name in
        <:str_item<
          let () = publish $str:name$ (
            function
            | [] ->
              Goaljobs.require $lid:gname$
            | _ ->
              failwith (Printf.sprintf "goal '%s' does not take any arguments"
                          $str:name$);
          )
        >> in
      StSem (_loc, stmt, publish_name)
  ) stmts !autopublish

(* Rewrite 'require (name args...)' as 'require (fun () -> goal_name args)'.
 * 'expr' is a function call.
 *)
let generate_require _loc expr =
  (* Note that 'f a b c' is parsed as '((f a) b) c' so the actual
   * function name is buried deeply in the tree.  Rewrite the name.
   *)
  let rec rewrite = function
    | ExApp (_loc, ExId (_loc1, IdLid (_loc2, name)), right) ->
      let gname = "goal_" ^ name in
      ExApp (_loc, ExId (_loc1, IdLid (_loc2, gname)), right)
    | ExApp (_loc, (ExApp _ as expr), right) ->
      ExApp (_loc, rewrite expr, right)
    | _ ->
      locfail _loc "require (...) expression must contain a call to a goal"
  in
  let expr = rewrite expr in
  <:expr< Goaljobs.require (fun () -> $expr$) >>

;;

EXTEND Gram
  GLOBAL: expr str_item;

  (* Rewrite 'require (name args...)'. *)
  expr: LEVEL "apply" [
    [ "require"; e = expr ->
      generate_require _loc e ]
  ];

  (* "str_item" is a top level statement in an OCaml program. *)
  str_item: LEVEL "top" [
    [ "let"; r = opt_rec; "goal"; ls = binding ->
    generate_let_goal _loc r ls ]
  ];

END