(*
   Code generator for the CST.ml file.

   This produces code similar to what's found in ../../run/lib/Sample.ml
*)

open Printf
open CST_grammar
open Indent.Types

(* All rule names and other names directly defined in grammar.json
   must go through this translation. For example, it turns
   "true" into "true_" because "true" is a reserved keyword in the generated
   code.
*)
let trans = Codegen_util.translate_ident

let mli_contents grammar : string =
  let lang = grammar.name in
  let root_type = grammar.entrypoint in
  sprintf {|
(**
    Functions for parsing %s programs into a CST.

    Generated by ocaml-tree-sitter.
*)

(** Parse a %s program from a string into a typed OCaml CST.
    The resulting CST is [None] if parsing failed completely, otherwise
    some tree is returned even if some parsing errors occurred, in which
    case the error list is not empty.
*)
val string :
  ?src_file:string -> string ->
  (CST.%s, CST.extra) Tree_sitter_run.Parsing_result.t

(** Parse a %s program from a file into a typed OCaml CST.
    See the [string] function above for details. *)
val file :
  string ->
  (CST.%s, CST.extra) Tree_sitter_run.Parsing_result.t

(** Whether to print debugging information. Default: false. *)
val debug : bool ref

(** The original tree-sitter parser. *)
val ts_parser : Tree_sitter_bindings.Tree_sitter_API.ts_parser

(** Parse a program into a tree-sitter CST. *)
val parse_source_string :
   ?src_file:string -> string -> Tree_sitter_run.Tree_sitter_parsing.t

(** Parse a source file into a tree-sitter CST. *)
val parse_source_file : string -> Tree_sitter_run.Tree_sitter_parsing.t

(** Parse a tree-sitter CST into an OCaml typed CST. *)
val parse_input_tree :
  Tree_sitter_run.Tree_sitter_parsing.t ->
  (CST.%s, CST.extra) Tree_sitter_run.Parsing_result.t
|}
    lang
    lang (trans root_type)
    lang (trans root_type)
    (trans root_type)

let constant_header = "\
(* Generated by ocaml-tree-sitter. *)

(*
   Disable warning 42:
     \"this use of Foo relies on type-directed disambiguation,
     it will not compile with OCaml 4.00 or earlier.\"
*)
[@@@warning \"-42\"]

(* Disable warnings against unused variables. *)
[@@@warning \"-26-27-32\"]

open Tree_sitter_bindings
open Tree_sitter_run

let debug = ref false

type mt = Run.matcher_token
"

let declare_externals lang = sprintf "\
external create_parser :
  unit -> Tree_sitter_API.ts_parser = \"octs_create_parser_%s\"

let ts_parser = create_parser ()

let parse_source_string ?src_file contents =
  Tree_sitter_parsing.parse_source_string ?src_file ts_parser contents

let parse_source_file src_file =
  Tree_sitter_parsing.parse_source_file ts_parser src_file
"
    lang

let gen_extras grammar =
  let items = List.map (fun name ->
    Line (sprintf "%S;" name)
  ) grammar.extras
  in
  [
    Line "let extras = [";
    Block items;
    Line "]";
  ]

let preamble grammar =
  [
    Line constant_header;
    Line (declare_externals grammar.name);
    Inline (gen_extras grammar);
    Line "";
  ]

let rec fmt_regexp (x : rule_body) =
  match x with
  | Symbol name
  | Alias (name, _) -> [Line (sprintf "Token (Name %S);" name)]
  | Token token -> [Line (sprintf "Token (Literal %S);" token.name)]
  | Blank -> [Line "Nothing;"]
  | Repeat x ->
      [
        Line "Repeat (";
        Block (fmt_regexp x);
        Line ");";
      ]
  | Repeat1 x ->
      [
        Line "Repeat1 (";
        Block (fmt_regexp x);
        Line ");";
      ]
  | Optional x ->
      [
        Line "Opt (";
        Block (fmt_regexp x);
        Line ");";
      ]

  | Choice cases ->
      [
        Line "Alt [|";
        Block (List.map (fun (_name, x) -> Inline (fmt_regexp x)) cases);
        Line "|];"
      ]

  | Seq l ->
      [
        Line "Seq [";
        Block (List.map (fun x -> Inline (fmt_regexp x)) l);
        Line "];";
      ]

let gen_regexps grammar =
  let list_elements =
    List.flatten grammar.rules
    |> List.map (fun rule ->
      let opt_regexp =
        match rule.body with
        | Token _ ->
            [Line "None;"]
        | body ->
            [
              Line "Some (";
              Block (fmt_regexp body);
              Line ");";
            ]
      in
      Group [
        Line (sprintf "%S," rule.name);
        Space;
        Inline opt_regexp;
      ]
    )
  in
  [
    Line "let children_regexps : (string * Run.exp option) list = [";
    Block list_elements;
    Line "]";
  ]

let gen_trans_token () =
  [
    Line "match body with";
    Line "| Leaf v -> v";
    Line "| Children _ -> assert false";
  ]

(*
   Mini DSL that allows simplifications of the generated code such as

     trans_expression (Run.matcher_token v2)

   instead of

     (fun v -> trans_expression (Run.matcher_token v)) v2
*)
type exp =
  | Code of Indent.t
  (* opaque code *)
  | Fun of (string -> exp)
  (* parametrized function body *)
  | App of exp * string
  (* application of a function to a precomputed value,
     such as a variable *)

let rec compile_exp exp : Indent.t =
  match exp with
  | Code code -> code
  | Fun f ->
      [
        Group [
          Line "(fun v ->";
          Space;
          Block (compile_exp (f "v"));
          Line ")"
        ]
      ]
  | App (Fun f, arg) ->
      compile_exp (f arg)
  | App (e, arg) ->
      [
        Group [
          Inline (compile_exp e);
          Space;
          Block [Line arg]
        ]
      ]

let gen_seq_pat l =
  List.mapi (fun i _ -> sprintf "v%i" i) l
  |> String.concat "; "

let wrap_seq elts =
  let n = List.length elts in
  [
    Group [
      Line "(";
      Block (
        List.mapi (fun i x ->
          if i = n - 1 then x
          else
            Inline [
              Group [
                x;
                Line ","
              ];
              Space
            ]
        ) elts
      );
      Line ")"
    ]
  ]

let rec gen_trans_capture rule_body : exp =
  match rule_body with
  | Symbol name
  | Alias (name, _) ->
      Fun (fun arg ->
        Code [
          Line (sprintf "trans_%s (Run.matcher_token %s)"
                  (trans name) arg);
        ]
      )
  | Token _ ->
      Fun (fun arg ->
        Code [
          Line (sprintf "Run.trans_token (Run.matcher_token %s)" arg);
        ]
      )
  | Blank -> Code [Line "Run.nothing"]
  | Repeat x ->
      Fun (fun arg ->
        Code [
          Line "Run.repeat";
          Block [
            Inline (compile_exp (gen_trans_capture x));
            Line arg;
          ];
        ]
      )
  | Repeat1 x ->
      Fun (fun arg ->
        Code [
          Line "Run.repeat1";
          Block [
            Inline (compile_exp (gen_trans_capture x));
            Line arg;
          ];
        ]
      )
  | Optional x ->
      Fun (fun arg ->
        Code [
          Line "Run.opt";
          Block [
            Inline (compile_exp (gen_trans_capture x));
            Line arg;
          ];
        ]
      )
  | Choice cases ->
      Fun (fun arg ->
        Code [
          Line (sprintf "(match %s with" arg);
          Inline (List.mapi gen_case cases);
          Line "| _ -> assert false";
          Line ")";
        ]
      )
  | Seq l ->
      Fun (fun arg ->
        Code [
          Line (sprintf "(match %s with" arg);
          Line (sprintf "| Seq [%s] ->" (gen_seq_pat l));
          Block [Block (wrap_seq (List.mapi gen_seq_elt l))];
          Line "| _ -> assert false";
          Line ")"
        ]
      )

and gen_case i (cons, x) =
  Inline [
    Line (sprintf "| Alt (%i, v) ->" i);
    Block [
      Block [
        Line (sprintf "`%s (" cons);
        Block (compile_exp (App (gen_trans_capture x, "v")));
        Line ")";
      ]
    ]
  ]

and gen_seq_elt i x =
  Inline (
    App (gen_trans_capture x, sprintf "v%i" i)
    |> compile_exp
  )

let gen_trans_body rule_body =
  let res =
    App (gen_trans_capture rule_body, "v")
    |> compile_exp
  in
  [
    Line "match body with";
    Line "| Children v ->";
    Block [Block res];
    Line "| Leaf _ -> assert false";
  ]

let gen_translator ~cst_module_name (rule : rule) =
  let name = rule.name in
  let body =
    match rule.body with
    | Token _token -> gen_trans_token ()
    | body -> gen_trans_body body
  in
  [
    Line (sprintf "trans_%s ((kind, body) : mt) : %s.%s ="
            (trans name) cst_module_name (trans name));
    Block body;
  ]

let gen_translators ~cst_module_name grammar =
  List.map (fun rule_group ->
    let is_rec =
      match rule_group with
      | [x] -> x.is_rec
      | _ -> true
    in
    let bindings =
      List.map (fun rule ->
        gen_translator ~cst_module_name rule
      ) rule_group
    in
    Codegen_util.format_bindings ~is_rec ~is_local:false bindings
  ) grammar.rules
  |> Codegen_util.interleave [Line ""]
  |> List.flatten

let trans_tree =
  {|(*
   Costly operation that translates a whole tree or subtree.

   The first pass translates it into a generic tree structure suitable
   to guess which node corresponds to each grammar rule.
   The second pass is a translation into a typed tree where each grammar
   node has its own type.

   This function is called:
   - once on the root of the program after removing extras
     (comments and other nodes that occur anywhere independently from
     the grammar);
   - once of each extra node, resulting in its own independent tree of type
     'extra'.
*)
let translate_tree src node trans_x =
  let matched_tree = Run.match_tree children_regexps src node in
  Option.map trans_x matched_tree
|}

let gen_trans_extras grammar =
  let body =
    match grammar.extras with
    | [] -> [ Line "None" ]
    | extras ->
        let cases =
          (extras
           |> List.map (fun name ->
             Inline [
               Line (sprintf "| %S ->" name);
               Block [
                 Block [
                   Line (sprintf
                           "(match translate_tree src node trans_%s with"
                           (trans name));
                   Line "| None -> None";
                   Line (sprintf "| Some x -> Some (`%s (Run.get_loc node, x)))"
                           (Codegen_util.translate_ident_uppercase name));
                 ];
               ]
             ]
           ))
          @ [ Line "| _ -> None" ]
        in
        [
          Line "match node.type_ with";
          Inline cases;
        ]
  in
  [
    Line "let translate_extra src (node : Tree_sitter_output_t.node) \
          : CST.extra option =";
    Block body
  ]

let gen ~cst_module_name grammar =
  let regexps = gen_regexps grammar in
  let translators = gen_translators ~cst_module_name grammar in
  let trans_extras = gen_trans_extras grammar in
  [
    Inline regexps;
    Line "";
    Inline translators;
    Line "";
    Line trans_tree;
    Line "";
    Inline trans_extras;
  ]

let ml_trailer grammar = sprintf {|
let translate_root src root_node =
  translate_tree src root_node trans_%s

let parse_input_tree input_tree =
  let orig_root_node = Tree_sitter_parsing.root input_tree in
  let src = Tree_sitter_parsing.src input_tree in
  let errors = Run.extract_errors src orig_root_node in
  let opt_program, extras =
     Run.translate
       ~extras
       ~translate_root:(translate_root src)
       ~translate_extra:(translate_extra src)
       orig_root_node
  in
  Parsing_result.create src opt_program extras errors

let string ?src_file contents =
  let input_tree = parse_source_string ?src_file contents in
  parse_input_tree input_tree

let file src_file =
  let input_tree = parse_source_file src_file in
  parse_input_tree input_tree
|}
    (trans grammar.entrypoint)

(*
   Generate a parsing function for each rule in the tree-sitter grammar
   ('grammar.json' file).
*)
let generate ~cst_module_name grammar =
  let grammar =
    let rules =
      List.map (
        List.filter (fun rule -> not rule.is_inlined_rule)
      ) grammar.rules
    in
    { grammar with rules }
  in
  let tree = gen ~cst_module_name grammar in
  let ml_contents = Indent.to_string [
    Inline (preamble grammar);
    Inline tree;
    Line (ml_trailer grammar);
  ] in
  mli_contents grammar, ml_contents
