(*
 * pa_optcomp.ml
 * -------------
 * Copyright : (c) 2008, Jeremie Dimino <jeremie@dimino.org>
 * Licence   : BSD3
 *
 * This file is a part of optcomp.
 *)

open Camlp4.Sig
open Camlp4.PreCast

external filter : 'a Gram.not_filtered -> 'a = "%identity"
external not_filtered : 'a -> 'a Gram.not_filtered = "%identity"

(* Subset of supported caml types *)
type typ =
  | Tvar of string
  | Tbool
  | Tint
  | Tchar
  | Tstring
  | Ttuple of typ list

(* Subset of supported caml values *)
type value =
  | Bool of bool
  | Int of int
  | Char of char
  | String of string
  | Tuple of value list

type ident = string
    (* An identifier. It is either a lower or a upper identifier. *)

module Env = Map.Make(struct type t = ident let compare = compare end)

type env = value Env.t

type directive =
  | Dir_let of ident * Ast.expr
  | Dir_default of ident * Ast.expr
  | Dir_if of Ast.expr
  | Dir_else
  | Dir_elif of Ast.expr
  | Dir_endif
  | Dir_include of Ast.expr
  | Dir_error of Ast.expr
  | Dir_warning of Ast.expr
  | Dir_directory of Ast.expr

  (* This one is not part of optcomp but this is one of the directives
     handled by camlp4 we probably want to use. *)
  | Dir_default_quotation of Ast.expr

(* Quotations are evaluated by the token filters, but are expansed
   after. Evaluated quotations are kept in this table, which quotation
   id to to values: *)
let quotations : (int, value) Hashtbl.t = Hashtbl.create 42

let next_quotation_id =
  let r = ref 0 in
  fun _ -> incr r; !r

(* +-------------+
   | Environment |
   +-------------+ *)

let env = ref Env.empty
let define id value = env := Env.add id value !env

let _ =
  define "ocaml_version" (Scanf.sscanf Sys.ocaml_version "%d.%d" (fun major minor -> Tuple [Int major; Int minor]))

let dirs = ref []
let add_include_dir dir = dirs := dir :: !dirs

(* +--------------+
   | Dependencies |
   +--------------+ *)

module String_set = Set.Make(String)

(* All depencies of the file being parsed *)
let dependencies = ref String_set.empty

(* Where to write dependencies *)
let dependency_filename = ref None

(* The file being parsed. This is set when the first (token, location)
   pair is fetched. *)
let source_filename = ref None

let write_depencies () =
  match !dependency_filename, !source_filename with
    | None, _
    | _, None ->
        ()

    | Some dependency_filename, Some source_filename ->
        let oc = open_out dependency_filename in
        if not (String_set.is_empty !dependencies) then begin
          output_string oc "# automatically generated by optcomp\n";
          output_string oc source_filename;
          output_string oc ": ";
          output_string oc (String.concat " " (String_set.elements !dependencies));
          output_char oc '\n'
        end;
        close_out oc

(* +----------------------------------------+
   | Value to expression/pattern conversion |
   +----------------------------------------+ *)

let rec expr_of_value _loc = function
  | Bool true -> <:expr< true >>
  | Bool false -> <:expr< false >>
  | Int x -> <:expr< $int:string_of_int x$ >>
  | Char x -> <:expr< $chr:Char.escaped x$ >>
  | String x -> <:expr< $str:String.escaped x$ >>
  | Tuple [] -> <:expr< () >>
  | Tuple [x] -> expr_of_value _loc x
  | Tuple l -> <:expr< $tup:Ast.exCom_of_list (List.map (expr_of_value _loc) l)$ >>

let rec patt_of_value _loc = function
  | Bool true -> <:patt< true >>
  | Bool false -> <:patt< false >>
  | Int x -> <:patt< $int:string_of_int x$ >>
  | Char x -> <:patt< $chr:Char.escaped x$ >>
  | String x -> <:patt< $str:String.escaped x$ >>
  | Tuple [] -> <:patt< () >>
  | Tuple [x] -> patt_of_value _loc x
  | Tuple l -> <:patt< $tup:Ast.paCom_of_list (List.map (patt_of_value _loc) l)$ >>

(* +-----------------------+
   | Expression evaluation |
   +-----------------------+ *)

let rec type_of_value = function
  | Bool _ -> Tbool
  | Int _ -> Tint
  | Char _ -> Tchar
  | String _ -> Tstring
  | Tuple l -> Ttuple (List.map type_of_value l)

let rec string_of_type = function
  | Tvar v -> "'" ^ v
  | Tbool -> "bool"
  | Tint -> "int"
  | Tchar -> "char"
  | Tstring -> "string"
  | Ttuple l -> "(" ^ String.concat " * " (List.map string_of_type l) ^ ")"

let invalid_type loc expected real =
  Loc.raise loc (Failure
                   (Printf.sprintf "this expression has type %s but is used with type %s"
                      (string_of_type real) (string_of_type expected)))

let type_of_patt patt =
  let rec aux (a, n) = function
    | <:patt< $tup:x$ >> ->
      let l, x = List.fold_left
        (fun (l, x) patt -> let t, x = aux x patt in (t :: l, x))
        ([], (a, n)) (Ast.list_of_patt x []) in
      (Ttuple(List.rev l), x)
    | _ ->
        (Tvar(Printf.sprintf "%c%s"
                (char_of_int (Char.code 'a' + a))
                (if n = 0 then "" else string_of_int n)),
         if a = 25 then (0, n + 1) else (a + 1, n))
  in
  fst (aux (0, 0) patt)

let rec eval env = function

  (* Literals *)
  | <:expr< true >> -> Bool true
  | <:expr< false >> -> Bool false
  | <:expr< $int:x$ >> -> Int(int_of_string x)
  | <:expr< $chr:x$ >> -> Char(Camlp4.Struct.Token.Eval.char x)
  | <:expr< $str:x$ >> -> String(Camlp4.Struct.Token.Eval.string ~strict:() x)

  (* Tuples *)
  | <:expr< $tup:x$ >> -> Tuple(List.map (eval env) (Ast.list_of_expr x []))

  (* Variables *)
  | <:expr@loc< $lid:x$ >>
  | <:expr@loc< $uid:x$ >> ->
    begin try
      Env.find x env
    with
        Not_found ->
          Loc.raise loc (Failure (Printf.sprintf "unbound value %s" x))
    end

  (* Value comparing *)
  | <:expr< $x$ = $y$ >> -> let x, y = eval_same env x y in Bool(x = y)
  | <:expr< $x$ < $y$ >> -> let x, y = eval_same env x y in Bool(x < y)
  | <:expr< $x$ > $y$ >> -> let x, y = eval_same env x y in Bool(x > y)
  | <:expr< $x$ <= $y$ >> -> let x, y = eval_same env x y in Bool(x <= y)
  | <:expr< $x$ >= $y$ >> -> let x, y = eval_same env x y in Bool(x >= y)
  | <:expr< $x$ <> $y$ >> -> let x, y = eval_same env x y in Bool(x <> y)

  (* min and max *)
  | <:expr< min $x$ $y$ >> -> let x, y = eval_same env x y in min x y
  | <:expr< max $x$ $y$ >> -> let x, y = eval_same env x y in max x y

  (* Arithmetic *)
  | <:expr< $x$ + $y$ >> -> Int(eval_int env x + eval_int env y)
  | <:expr< $x$ - $y$ >> -> Int(eval_int env x - eval_int env y)
  | <:expr< $x$ * $y$ >> -> Int(eval_int env x * eval_int env y)
  | <:expr< $x$ / $y$ >> -> Int(eval_int env x / eval_int env y)
  | <:expr< $x$ mod $y$ >> -> Int(eval_int env x mod eval_int env y)

  (* Boolean operations *)
  | <:expr< not $x$ >> -> Bool(not (eval_bool env x))
  | <:expr< $x$ or $y$ >> -> Bool(eval_bool env x or eval_bool env y)
  | <:expr< $x$ || $y$ >> -> Bool(eval_bool env x || eval_bool env y)
  | <:expr< $x$ && $y$ >> -> Bool(eval_bool env x && eval_bool env y)

  (* String operations *)
  | <:expr< $x$ ^ $y$ >> -> String(eval_string env x ^ eval_string env y)

  (* Pair operations *)
  | <:expr< fst $x$ >> -> fst (eval_pair env x)
  | <:expr< snd $x$ >> -> snd (eval_pair env x)

  (* Let-binding *)
  | <:expr< let $p$ = $x$ in $y$ >> ->
    let vx = eval env x in
    let env =
      try
        bind env p vx
      with
          Exit -> invalid_type (Ast.loc_of_expr x) (type_of_patt p) (type_of_value vx)
    in
    eval env y

  | e -> Loc.raise (Ast.loc_of_expr e) (Stream.Error "expression not supported")

and bind env patt value = match patt with
  | <:patt< $lid:id$ >> ->
    Env.add id value env

  | <:patt< $tup:patts$ >> ->
    let patts = Ast.list_of_patt patts [] in
    begin match value with
      | Tuple values when List.length values = List.length patts ->
          List.fold_left2 bind env patts values
      | _ ->
          raise Exit
    end

  | _ ->
      Loc.raise (Ast.loc_of_patt patt) (Stream.Error "pattern not supported")

and eval_same env ex ey =
  let vx = eval env ex and vy = eval env ey in
  let tx = type_of_value vx and ty = type_of_value vy in
  if tx = ty then
    (vx, vy)
  else
    invalid_type (Ast.loc_of_expr ey) tx ty

and eval_int env e = match eval env e with
  | Int x -> x
  | v -> invalid_type (Ast.loc_of_expr e) Tint (type_of_value v)

and eval_bool env e = match eval env e with
  | Bool x -> x
  | v -> invalid_type (Ast.loc_of_expr e) Tbool (type_of_value v)

and eval_string env e = match eval env e with
  | String x -> x
  | v -> invalid_type (Ast.loc_of_expr e) Tstring (type_of_value v)

and eval_pair env e = match eval env e with
  | Tuple [x; y] -> (x, y)
  | v -> invalid_type (Ast.loc_of_expr e) (Ttuple [Tvar "a"; Tvar "b"]) (type_of_value v)

(* +-----------------------+
   | Parsing of directives |
   +-----------------------+ *)

let rec skip_space stream = match Stream.peek stream with
  | Some((BLANKS _ | COMMENT _), _) ->
      Stream.junk stream;
      skip_space stream
  | _ ->
      ()

let parse_equal stream =
  skip_space stream;
  match Stream.next stream with
    | KEYWORD "=", _ -> ()
    | _, loc -> Loc.raise loc (Stream.Error "'=' expected")

let rec parse_eol stream =
  let tok, loc = Stream.next stream in
  match tok with
    | BLANKS _ | COMMENT _ ->
        parse_eol stream
    | NEWLINE | EOI ->
        ()
    | _ ->
        Loc.raise loc (Stream.Error "end of line expected")

(* Return wether a keyword can be interpreted as an identifier *)
let keyword_is_id str =
  let rec aux i =
    if i = String.length str then
      true
    else
      match str.[i] with
        | 'A' .. 'Z' | 'a' .. 'z' | '0' .. '9' ->
            aux (i + 1)
        | _ ->
            false
  in
  aux 0

let parse_ident stream =
  skip_space stream;
  let tok, loc = Stream.next stream in
  begin match tok with
    | LIDENT id | UIDENT id ->
        id
    | KEYWORD kwd when keyword_is_id kwd ->
        kwd
    | _ ->
        Loc.raise loc (Stream.Error "identifier expected")
  end

let parse_expr stream =
  (* Lists of opened brackets *)
  let opened_brackets = ref [] in

  (* Return the next token of [stream] until all opened parentheses
     have been closed and a newline is reached *)
  let rec next_token _ =
    Some(match Stream.next stream, !opened_brackets with
           | (NEWLINE, loc), [] ->
               EOI, loc

           | (KEYWORD("(" | "[" | "{" as b), _) as x, l ->
               opened_brackets := b :: l;
               x

           | (KEYWORD ")", loc) as x, "(" :: l ->
               opened_brackets := l;
               x

           | (KEYWORD "]", loc) as x, "[" :: l ->
               opened_brackets := l;
               x

           | (KEYWORD "}", loc) as x, "{" :: l ->
               opened_brackets := l;
               x

           | x, _ ->
               x)
  in

  Gram.parse_tokens_before_filter Syntax.expr_eoi
    (not_filtered (Stream.from next_token))

let parse_directive stream = match Stream.peek stream with
  | Some(KEYWORD "#", loc) ->
      Stream.junk stream;

      (* Move the location to the beginning of the line *)
      let (file_name,
           start_line, start_bol, start_off,
           stop_line,  stop_bol,  stop_off,
           ghost) = Loc.to_tuple loc in
      let loc = Loc.of_tuple (file_name,
                              start_line, start_bol, start_bol,
                              start_line, start_bol, start_bol,
                              ghost) in

      begin match parse_ident stream with

        | "let" ->
            let id = parse_ident stream in
            parse_equal stream;
            let expr = parse_expr stream in
            Some(Dir_let(id, expr), loc)

        | "let_default" ->
            let id = parse_ident stream in
            parse_equal stream;
            let expr = parse_expr stream in
            Some(Dir_default(id, expr), loc)

        (* For compatibility *)
        | "define" ->
            let id = parse_ident stream in
            let expr = parse_expr stream in
            Some(Dir_let(id, expr), loc)

        (* For compatibility *)
        | "default" ->
            let id = parse_ident stream in
            let expr = parse_expr stream in
            Some(Dir_default(id, expr), loc)

        | "if" ->
            Some(Dir_if(parse_expr stream), loc)

        | "else" ->
            parse_eol stream;
            Some(Dir_else, loc)

        | "elif" ->
            Some(Dir_elif(parse_expr stream), loc)

        | "endif" ->
            parse_eol stream;
            Some(Dir_endif, loc)

        | "include" ->
            Some(Dir_include(parse_expr stream), loc)

        | "directory" ->
            Some(Dir_directory(parse_expr stream), loc)

        | "error" ->
            Some(Dir_error(parse_expr stream), loc)

        | "warning" ->
            Some(Dir_warning(parse_expr stream), loc)

        | "default_quotation" ->
            Some(Dir_default_quotation(parse_expr stream), loc)

        | dir ->
            Loc.raise loc (Stream.Error (Printf.sprintf "bad directive ``%s''" dir))
      end

  | _ ->
      None

let parse_command_line_define str =
  match Gram.parse_string Syntax.expr (Loc.mk "<command line>") str with
    | <:expr< $lid:id$ = $e$ >>
    | <:expr< $uid:id$ = $e$ >> -> define id (eval !env e)
    | _ -> invalid_arg str

(* +----------------+
   | BLock skipping |
   +----------------+ *)

let rec skip_line stream =
  match Stream.next stream with
    | NEWLINE, _ -> ()
    | EOI, loc -> Loc.raise loc (Stream.Error "#endif missing")
    | _ -> skip_line stream

let rec next_directive stream = match parse_directive stream with
  | Some dir -> dir
  | None -> skip_line stream; next_directive stream

let rec next_endif stream =
  let dir, loc = next_directive stream in
  match dir with
    | Dir_if _ -> skip_if stream; next_endif stream
    | Dir_else
    | Dir_elif _
    | Dir_endif -> dir
    | _ -> next_endif stream

and skip_if stream =
  let dir, loc = next_directive stream in
  match dir with
    | Dir_if _ ->
        skip_if stream;
        skip_if stream

    | Dir_else ->
        skip_else stream

    | Dir_elif _ ->
        skip_if stream

    | Dir_endif ->
        ()

    | _ -> skip_if stream

and skip_else stream =
  let dir, loc = next_directive stream in
  match dir with
    | Dir_if _ ->
        skip_if stream;
        skip_else stream

    | Dir_else ->
        Loc.raise loc (Stream.Error "#else without #if")

    | Dir_elif _ ->
        Loc.raise loc (Stream.Error "#elif without #if")

    | Dir_endif ->
        ()

    | _ ->
        skip_else stream

(* +-----------------+
   | Token filtering |
   +-----------------+ *)

type context = Ctx_if | Ctx_else

(* State of the token filter *)
type state = {
  stream : (Gram.Token.t * Loc.t) Stream.t;
  (* Input stream *)

  mutable bol : bool;
  (* Wether we are at the beginning of a line *)

  mutable stack : context list;
  (* Nested contexts *)

  on_eoi : Gram.Token.t * Loc.t -> Gram.Token.t * Loc.t;
  (* Eoi handler, it is used to restore the previous sate on #include
     directives *)
}

(* Read and return one token *)
let really_read state =
  let tok, loc = Stream.next state.stream in
  state.bol <- tok = NEWLINE;
  match tok with
    | QUOTATION ({ q_name = "optcomp" } as quot) ->
        let id = next_quotation_id () in
        Hashtbl.add quotations id (eval !env (Gram.parse_string
                                                Syntax.expr_eoi
                                                (Loc.move `start quot.q_shift loc)
                                                quot.q_contents));

        (* Replace the quotation by its id *)
        (QUOTATION { quot with q_contents = string_of_int id }, loc)

    | EOI ->
        (* If end of input is reached, we call the eoi handler. It may
           continue if we were parsing an included file *)
        if state.stack <> [] then
          Loc.raise loc (Stream.Error "#endif missing");
        state.on_eoi (tok, loc)

    | _ ->
        (tok, loc)

(* Return the next token from a stream, interpreting directives. *)
let rec next_token state_ref =
  let state = !state_ref in
  if state.bol then
    match parse_directive state.stream, state.stack with
      | Some(Dir_if e, _), _ ->
          let rec aux e =
            if eval_bool !env e then begin
              state.stack <- Ctx_if :: state.stack;
              next_token state_ref
            end else
              match next_endif state.stream with
                | Dir_else ->
                    state.stack <- Ctx_else :: state.stack;
                    next_token state_ref

                | Dir_elif e ->
                    aux e

                | Dir_endif ->
                    next_token state_ref

                | _ ->
                    assert false
          in
          aux e

      | Some(Dir_else, loc), ([] | Ctx_else :: _) ->
          Loc.raise loc (Stream.Error "#else without #if")

      | Some(Dir_elif _, loc), ([] | Ctx_else :: _) ->
          Loc.raise loc (Stream.Error "#elif without #if")

      | Some(Dir_endif, loc), [] ->
          Loc.raise loc (Stream.Error "#endif without #if")

      | Some(Dir_else, loc), Ctx_if :: l ->
          skip_else state.stream;
          state.stack <- l;
          next_token state_ref

      | Some(Dir_elif _, loc), Ctx_if :: l ->
          skip_if state.stream;
          state.stack <- l;
          next_token state_ref

      | Some(Dir_endif, loc), _ :: l ->
          state.stack <- l;
          next_token state_ref

      | Some(Dir_let(id, e), _), _ ->
          define id (eval !env e);
          next_token state_ref

      | Some(Dir_default(id, e), _), _ ->
          if not (Env.mem id !env) then
            define id (eval !env e);
          next_token state_ref

      | Some(Dir_include e, _), _ ->
          let fname = eval_string !env e in
          (* Try to looks up in all include directories *)
          let fname =
            try
              List.find (fun dir -> Sys.file_exists (Filename.concat dir fname)) !dirs
            with
                (* Just try in the current directory *)
                Not_found -> fname
          in
          dependencies := String_set.add fname !dependencies;
          let ic = open_in fname in
          let nested_state = {
            stream = Gram.Token.Filter.filter (Gram.get_filter ()) (filter (Gram.lex (Loc.mk fname) (Stream.of_channel ic)));
            bol = true;
            stack = [];
            on_eoi = (fun _ ->
                        (* Restore previous state and close channel on
                           eoi *)
                        state_ref := state;
                        close_in ic;
                        next_token state_ref)
          } in
          (* Replace current state with the new one *)
          state_ref := nested_state;
          next_token state_ref

      | Some(Dir_directory e, loc), _ ->
          let dir = eval_string !env e in
          add_include_dir dir;
          next_token state_ref

      | Some(Dir_error e, loc), _ ->
          Loc.raise loc (Failure (eval_string !env e))

      | Some(Dir_warning e, loc), _ ->
          Syntax.print_warning loc (eval_string !env e);
          next_token state_ref

      | Some(Dir_default_quotation e, loc), _ ->
          Syntax.Quotation.default := eval_string !env e;
          next_token state_ref

      | None, _ ->
          really_read state

  else
    really_read state

let stream_filter filter stream =
  (* Set the source filename *)
  begin match !source_filename with
    | Some _ ->
        ()
    | None ->
        match Stream.peek stream with
          | None ->
              ()
          | Some(tok, loc) ->
              source_filename := Some(Loc.file_name loc)
  end;
  let state_ref = ref { stream = stream;
                        bol = true;
                        stack = [];
                        on_eoi = (fun x -> x) } in
  filter (Stream.from (fun _ -> Some(next_token state_ref)))

(* +----------------------+
   | Quotations expansion |
   +----------------------+ *)

let expand f loc _ contents =
  try
    f loc (Hashtbl.find quotations (int_of_string contents))
  with
      exn -> Loc.raise loc (Failure "fatal error in optcomp!")

(* +--------------+
   | Registration |
   +--------------+ *)

let _ =
  Camlp4.Options.add "-let" (Arg.String parse_command_line_define)
    "<string> Binding for a #let directive.";
  Camlp4.Options.add "-depend"
    (Arg.String (fun filename -> dependency_filename := Some filename))
    "<file> Write dependencies to <file>.";

  Pervasives.at_exit write_depencies;

  Syntax.Quotation.add "optcomp" Syntax.Quotation.DynAst.expr_tag (expand expr_of_value);
  Syntax.Quotation.add "optcomp" Syntax.Quotation.DynAst.patt_tag (expand patt_of_value);

  Gram.Token.Filter.define_filter (Gram.get_filter ()) stream_filter