open Import

(* This module interprets [(menhir ...)] stanzas -- that is, it provides build
   rules for Menhir parsers. *)

(* We assume that Menhir supports the commands [--infer-write-query] and
   [--infer-read-reply]. Although we could test at runtime whether this is the
   case, it is simpler to just require it. Dune is shipped with a constraint
   that Menhir (if present) is at version 20180523 or newer. *)

(* In order to perform OCaml type inference, we first let Menhir create a mock
   [.ml] file that contains just the semantic actions. We then use the OCaml
   compiler to perform type inference and create an inferred [.mli] file.
   Finally, we let Menhir read this [.mli] file. *)

(* This three-step process described is skipped in some cases where we know that
   OCaml type inference is not required, e.g., when Menhir is used in certain
   special modes, such as [--only-tokens]. In those cases, we use a
   straightforward one-step process. *)

(* -------------------------------------------------------------------------- *)

(* This signature describes the input of the functor [Run], which follows. *)

type stanza = Menhir_stanza.t

module type PARAMS = sig
  (* [cctx] is the compilation context. *)

  val cctx : Compilation_context.t

  (* [dir] is the directory inside [_build/<context>/...] where the build
     happens. If the [(menhir ...)] stanza appears in [src/dune], then [dir] is
     of the form [_build/<context>/src], e.g., [_build/default/src]. *)
  val dir : Path.Build.t

  (* [stanza] is the [(menhir ...)] stanza, as found in the [dune] file. *)
  val stanza : stanza

  (* [module_path] is the module path from the module group root to this
     module, accounting for `(include_subdirs ..)`. *)
  val module_path : Module_name.t list
end

(* TODO(anmonteiro): support melange *)
let for_ = Compilation_mode.Ocaml

(* -------------------------------------------------------------------------- *)

(* This functor is where [(menhir ...)] stanzas are desugared. *)

let menhir_env =
  let f =
    Env_stanza_db_flags.flags
      ~name:"jsoo-env"
      ~root:(fun _ _ ->
        Menhir_env.map ~f:Action_builder.return Menhir_env.default |> Memo.return)
      ~f:(fun ~parent expander (local : Dune_env.config) ->
        let local = local.menhir in
        let open Memo.O in
        let+ parent = parent in
        let flags =
          Expander.expand_and_eval_set expander local.flags ~standard:parent.flags
        in
        { Menhir_env.flags; explain = Option.first_some local.explain parent.explain })
  in
  fun ~dir ->
    let open Memo.O in
    let* () = Memo.return () in
    (Staged.unstage f) dir
;;

module Run (P : PARAMS) = struct
  open P

  (* [sctx] is the "super context", while [context] is the "context". Both store
     information about the current build context. *)

  let sctx = Compilation_context.super_context cctx

  (* [build_dir] is the base directory of the context; we run menhir from this
     directory to we get correct error paths. *)
  let build_dir = Super_context.context sctx |> Context.build_dir
  let expander = Super_context.expander ~dir sctx
  let env = menhir_env ~dir

  let sandbox =
    let scope = Compilation_context.scope cctx in
    let project = Scope.project scope in
    if Dune_project.dune_version project < (3, 5)
    then Sandbox_config.default
    else Sandbox_config.needs_sandboxing
  ;;

  (* ------------------------------------------------------------------------ *)

  (* Naming conventions. *)

  (* If [m] is a (short) module name, such as "myparser", then [source m] is the
     corresponding source file, and [targets m] is the list of targets that
     Menhir must build. *)

  let targets m ~cmly =
    let base = [ m ^ ".ml"; m ^ ".mli" ] in
    List.map ~f:(Path.Build.relative dir) (if cmly then (m ^ ".cmly") :: base else base)
  ;;

  (* The following definitions control where the mock [.ml] file and the
     inferred [.mli] file are created and how they are named. *)

  (* We change the module's base name, and use dummy extensions, so as to
     minimize the risk of confusing the build system (and the user). *)

  let mock m = m ^ "__mock"
  let mock_ml m : Path.Build.t = Path.Build.relative dir (mock m ^ ".ml.mock")
  let inferred_mli m : Path.Build.t = Path.Build.relative dir (mock m ^ ".mli.inferred")

  (* ------------------------------------------------------------------------ *)

  (* Rule generation. *)

  let menhir_binary =
    Super_context.resolve_program
      sctx
      ~dir
      ~where:Original_path
      "menhir"
      ~loc:None
      ~hint:"opam install menhir"
  ;;

  (* Reminder (from command.mli):

     [Deps] is for command line arguments that are dependencies. [As] is for
     command line arguments that are neither dependencies nor targets.
     [Hidden_targets] is for targets that are *not* command line arguments. *)

  type 'a args = 'a Command.Args.t list

  (* [menhir args] generates a Menhir command line (a build action). *)

  let menhir (args : 'a args) : Action.Full.t Action_builder.With_targets.t =
    Command.run_dyn_prog ~sandbox ~dir:(Path.build build_dir) menhir_binary args
  ;;

  let rule ~mode : Action.Full.t Action_builder.With_targets.t -> unit Memo.t =
    Super_context.add_rule sctx ~dir ~mode ~loc:stanza.loc
  ;;

  let explain_flags base stanza =
    let open Memo.O in
    let* expander = expander
    and* env = env in
    let+ explain =
      match Option.first_some stanza.Menhir_stanza.explain env.Menhir_env.explain with
      | None -> Memo.return (stanza.menhir_syntax >= Dune_lang.Menhir.explain_since)
      | Some explain -> Expander.eval_blang expander explain
    in
    if explain
    then
      [ Command.Args.A "--explain"
      ; Hidden_targets [ Path.Build.relative dir (base ^ ".conflicts") ]
      ]
    else []
  ;;

  let expand_flags flags =
    let open Memo.O in
    let+ env = env
    and+ expander = expander in
    Action_builder.memoize
      ~cutoff:(List.equal String.equal)
      "menhir flags"
      (Expander.expand_and_eval_set expander flags ~standard:env.Menhir_env.flags)
  ;;

  (* ------------------------------------------------------------------------ *)

  (* If there is no [base] clause, then a stanza that mentions several modules
     is equivalent to a list of stanzas, each of which mentions one module, so
     Menhir must be invoked once per module, separately. If there is a [base]
     clause, then the stanza describes a multi-module parser, so Menhir must be
     invoked once. In either case, we are able to reformulate the input in the
     form of a list of stanzas, each of which has a [base] clause. *)

  (* The current concrete name for [base] clauses is [merge_into], but I would
     like to change it in the future. *)

  let stanzas : (stanza * Path.Set.t) list Memo.Lazy.t =
    Memo.lazy_ (fun () ->
      let open Memo.O in
      let+ { Ml_sources.Parser_generators.deps; targets = _ } =
        let sctx = Compilation_context.super_context cctx in
        Dir_contents.get sctx ~dir
        >>= Dir_contents.ml ~for_
        >>| Ml_sources.Parser_generators.modules ~for_:(Menhir stanza.loc)
      in
      match stanza.merge_into with
      | None ->
        Path.Set.fold deps ~init:[] ~f:(fun p acc ->
          let merge_into = Filename.remove_extension (Path.basename p) in
          ({ stanza with merge_into = Some merge_into }, Path.Set.singleton p) :: acc)
      | Some _ -> [ stanza, deps ])
  ;;

  (* ------------------------------------------------------------------------ *)

  (* The [--infer-*] commands should not be passed by the user; we take care of
     using these commands appropriately. Fail if they are present. *)

  let check =
    Memo.lazy_ (fun () ->
      let open Memo.O in
      Memo.Lazy.force stanzas
      >>| List.iter ~f:(fun ((stanza, _) : stanza * Path.Set.t) ->
        Ordered_set_lang.Unexpanded.fold_strings
          stanza.flags
          ~init:()
          ~f:(fun _pos sw () ->
            match String_with_vars.text_only sw with
            | None -> ()
            | Some "--explain" ->
              if stanza.menhir_syntax >= Dune_lang.Menhir.explain_since
              then
                User_error.raise
                  ~loc:(String_with_vars.loc sw)
                  [ Pp.textf
                      "The Menhir '.conflicts' file is generated by default, so \
                       '--explain' should not be explicitly added to the list of Menhir \
                       flags."
                  ]
            | Some text ->
              if
                List.mem
                  ~equal:String.equal
                  [ "--depend"
                  ; "--raw-depend"
                  ; "--infer"
                  ; "--infer-write-query"
                  ; "--infer-read-reply"
                  ]
                  text
              then
                User_error.raise
                  ~loc:(String_with_vars.loc sw)
                  [ Pp.textf "The flag %s must not be used in a menhir stanza." text ])))
  ;;

  (* ------------------------------------------------------------------------ *)

  (* [process3 stanza] converts a Menhir stanza into a set of build rules. This
     is the three-step process where Menhir is invoked twice and OCaml type
     inference is performed in between. *)

  let process3 base ~cmly ((stanza, deps) : stanza * Path.Set.t) : unit Memo.t =
    let open Memo.O in
    let* expanded_flags = expand_flags stanza.flags in
    (* 1. A first invocation of Menhir creates a mock [.ml] file. *)
    let source_deps = Path.Set.to_list deps in
    let* () =
      menhir
        [ Command.Args.dyn expanded_flags
        ; Deps source_deps
        ; A "--base"
        ; Path (Path.relative (Path.build dir) base)
        ; A "--infer-write-query"
        ; Target (mock_ml base)
        ]
      |> rule ~mode:Standard
    in
    (* 2. The OCaml compiler performs type inference. *)
    let name =
      Module_name.of_string_allow_invalid (stanza.loc, mock base)
      |> Module_name.Unchecked.allow_invalid
    in
    let mock_module : Module.t =
      let source =
        let impl = Module.File.make Dialect.ocaml (Path.build (mock_ml base)) in
        Module.Source.make
          ~impl:(Some impl)
          ~intf:None
          Nonempty_list.(module_path @ [ name ])
      in
      Module.of_source ~visibility:Public ~kind:Impl source
    in
    let* mock_module =
      Pp_spec.pp_module_as
        (Compilation_context.preprocessing cctx)
        name
        mock_module
        ~lint:false
    in
    let cctx =
      Compilation_context.set_sandbox cctx Sandbox_config.needs_sandboxing
      |> Compilation_context.without_bin_annot
    in
    let* deps =
      let obj_dir = Compilation_context.obj_dir cctx in
      let modules = Compilation_context.modules cctx in
      let impl = Compilation_context.implements cctx in
      let dir = Obj_dir.dir obj_dir in
      Dep_rules.for_module ~obj_dir ~modules ~sandbox ~impl ~dir ~sctx mock_module ~for_
    in
    let* () =
      Module_compilation.ocamlc_i ~deps cctx mock_module ~output:(inferred_mli base)
    in
    let* explain_flags = explain_flags base stanza
    and* mode =
      let sctx = Compilation_context.super_context cctx in
      let* expander = Super_context.expander sctx ~dir in
      Rule_mode_expand.expand_path ~expander ~dir stanza.mode
    in
    (* 3. A second invocation of Menhir reads the inferred [.mli] file. *)
    menhir
      [ Command.Args.dyn expanded_flags
      ; S explain_flags
      ; Deps source_deps
      ; A "--base"
      ; Path (Path.relative (Path.build dir) base)
      ; A "--infer-read-reply"
      ; Dep (Path.build (inferred_mli base))
      ; Hidden_targets (targets base ~cmly)
      ]
    |> rule ~mode
  ;;

  (* ------------------------------------------------------------------------ *)

  (* [process3 stanza] converts a Menhir stanza into a set of build rules. This
     is a simpler one-step process where Menhir is invoked directly. *)

  let process1 base ~cmly ((stanza, deps) : stanza * Path.Set.t) : unit Memo.t =
    let open Memo.O in
    let* expanded_flags = expand_flags stanza.flags
    and* explain_flags = explain_flags base stanza
    and* mode =
      let sctx = Compilation_context.super_context cctx in
      let* expander = Super_context.expander sctx ~dir in
      Rule_mode_expand.expand_path ~expander ~dir stanza.mode
    in
    menhir
      [ Command.Args.dyn expanded_flags
      ; S explain_flags
      ; Deps (Path.Set.to_list deps)
      ; A "--base"
      ; Path (Path.relative (Path.build dir) base)
      ; Hidden_targets (targets base ~cmly)
      ]
    |> rule ~mode
  ;;

  (* ------------------------------------------------------------------------ *)

  (* [process stanza] converts a Menhir stanza into a set of build rules, using
     either [process3] or [process1], as appropriate. *)

  (* Because Menhir processes [--only-tokens] before the [--infer-*] commands,
     when [--only-tokens] is present, no [--infer-*] command should be used. *)

  let process ((stanza, sources) : stanza * Path.Set.t) : unit Memo.t =
    let base = Option.value_exn stanza.merge_into in
    let ocaml_type_inference_disabled, cmly =
      Ordered_set_lang.Unexpanded.fold_strings
        stanza.flags
        ~init:(false, false)
        ~f:(fun pos sw ((only_tokens, cmly) as acc) ->
          match pos with
          | Neg -> acc
          | Pos ->
            (match String_with_vars.text_only sw with
             | Some "--only-tokens" -> true, cmly
             | Some "--cmly" -> only_tokens, true
             | Some _ | None -> acc))
    in
    if ocaml_type_inference_disabled || not stanza.infer
    then process1 base (stanza, sources) ~cmly
    else process3 base (stanza, sources) ~cmly
  ;;

  (* ------------------------------------------------------------------------ *)
  let gen_rules () =
    let open Memo.O in
    let* () = Memo.Lazy.force check
    and* stanzas = Memo.Lazy.force stanzas in
    Memo.sequential_iter stanzas ~f:process
  ;;
end

(* -------------------------------------------------------------------------- *)

(* The final glue. *)

let gen_rules ~dir ~module_path cctx stanza =
  let module R = Run (struct
      let cctx = cctx
      let dir = dir
      let stanza = stanza
      let module_path = module_path
    end)
  in
  R.gen_rules ()
;;
