Source file ThunkLuaUtils.ml

module LuaObserver =
  ThunkParsers.Results.MakeObserverWithDiagnoseErrors
    (Diagnose.Diagnose.ConsolePlainStyle)

let lua_lex map buf = Luascanner.token buf map

let render_lua_error ?downgrade_errors_into_warnings ?origin ~code ~line ~col
    ~script_content ~msg () =
  let range =
    ThunkRanges.range_exact_linecol_in_string ~line ~col script_content
  in
  ( range,
    ThunkResults.single_error ~code ~msg
      ~brief_instruction:
        "This is the location of the final error, but it may be later than the \
         line and column of the root cause error."
      (module LuaObserver)
      (ThunkResults.State.create_with_source ?origin
         ?downgrade_errors_into_warnings script_content)
      (ThunkRanges.raw_range range) )

let lua_error_to_semantic ?downgrade_errors_into_warnings ?origin ~code ~line
    ~col ~script_content ~msg () =
  let range, rendered =
    render_lua_error ?downgrade_errors_into_warnings ?origin ~code ~line ~col
      ~script_content ~msg ()
  in
  let semantic =
    ThunkResults.Semantic.create_rendered ~brief:msg range rendered
  in
  semantic

module IoUtils
    (I : Luainterp.S)
    (Parser : Luaparser.S with type chunk = I.Ast.chunk) =
struct
  let lex_into_chunks map buf =
    try
      let chunks = Parser.chunks (lua_lex map) buf in
      Ok chunks
    with Parser.Error ->
      let file, line, col = Luasrcmap.last map in
      let errmsg =
        Format.asprintf "Lua syntax error at %a"
          (ThunkRanges.pp_pos_1based_linecol (Some file))
          (line, col)
      in
      (* let errmsg =
        let bt = Printexc.raw_backtrace_to_string (Printexc.get_callstack 20) in
        Format.asprintf "Lua syntax error at %a from %s"
          (ThunkRanges.pp_pos_1based_linecol (Some file))
          (line, col) bt
      in *)
      Error (errmsg, Some (`LineCol (line, col)))

  let do_chunks ~sourcename:filename buf =
    let map = Luasrcmap.mk () in
    Luasrcmap.sync map 0 (filename, 1, 1);
    match lex_into_chunks map buf with
    | Error (msg, loc_opt) -> Error (msg, loc_opt)
    | Ok chunks -> Ok (chunks, map)

  let format_errmsg ~file ~line ~col start s =
    Printf.sprintf "%s%s:\n%s" start
      (Format.asprintf "%a"
         (ThunkRanges.pp_pos_1based_linecol (Some file))
         (line, col))
      (ThunkStrings.left_pad s)

  (** [compile_chunks map chunks] compiles the chunks. The chunks will be
      compiled as if a ["$debug"] pragma statement was present as the first
      chunk.

      The compiled chunks are a function which can be called later with the unit
      `()` argument to run the chunks. *)
  let compile_chunks g map chunks =
    let debug = true in
    try
      let pgm = I.compile ~srcdbg:(map, debug) chunks g in
      Ok pgm
    with Parser.Error ->
      let file, line, col = Luasrcmap.last map in
      let errmsg = format_errmsg ~file ~line ~col "Lua syntax error at " "" in
      Error (errmsg, Some (`LineCol (line, col)))

  let eval_chunks ~in_what g map chunks =
    try
      let pgm = compile_chunks g map chunks in
      match pgm with
      | Error (msg, loc_opt) -> Error (msg, loc_opt)
      | Ok pgm ->
          let values = pgm () in
          Ok values
    with
    | ThunkLuaErrors.Error s | I.Error s ->
        (* [I.Error] should not be raised if the ThunkLuaUserCode error fallback
           is registered. But keep in here for safety. *)
        let file, line, col = Luasrcmap.last map in
        let errmsg =
          format_errmsg ~file ~line ~col
            ("Lua runtime error " ^ in_what ^ " at ")
            s
        in
        Error (errmsg, Some (`LineCol (line, col)))
    | I.Value.Projection (_v, w) ->
        let file, line, col = Luasrcmap.last map in
        let errmsg =
          format_errmsg ~file ~line ~col
            "Could not project (convert) the Lua value at "
            (Format.asprintf "`%s`" w)
        in
        Error (errmsg, Some (`LineCol (line, col)))

  let compile_lexbuf ~sourcename:filename g buf =
    let map = Luasrcmap.mk () in
    Luasrcmap.sync map 0 (filename, 1, 1);
    match lex_into_chunks map buf with
    | Error (msg, loc_opt) -> Error (msg, loc_opt)
    | Ok chunks -> compile_chunks g map chunks

  let do_lexbuf ~sourcename:filename ~in_what g buf =
    let map = Luasrcmap.mk () in
    Luasrcmap.sync map 0 (filename, 1, 1);
    match lex_into_chunks map buf with
    | Error (msg, loc_opt) -> Error (msg, loc_opt)
    | Ok chunks -> eval_chunks ~in_what g map chunks

  let do_lexbuf_return_all ~sourcename:filename ~in_what g buf =
    let map = Luasrcmap.mk () in
    Luasrcmap.sync map 0 (filename, 1, 1);
    match lex_into_chunks map buf with
    | Error (msg, loc_opt) -> Error (msg, loc_opt)
    | Ok chunks ->
        eval_chunks ~in_what g map chunks
        |> Result.map (fun r -> (chunks, map, r))
end

module IoUtils' (I : Lua.INTERP) = IoUtils (I) (I.Parser)

module ValueUtils (V : Lua.VALUE) = struct
  let getloc g' =
    (* adapted from luastdinterp.nw's currentloc_fromstack.
       but doesn't work, even with multiple items in callstack. *)
    let { callstack; currentloc; globals = _; fallbacks = _; startup = _ } :
        V.state =
      g'
    in
    match currentloc with
    | Some where -> Some where
    | None -> List.find_map (fun (_info, where) -> where) callstack

  let setglobal g s v = V.Table.bind g.V.globals ~key:(String s) ~data:v

  let add_to_table_list t v =
    let idx = 1.0 +. float_of_int (V.Luahash.length t) in
    V.Table.bind t ~key:(Number idx) ~data:v

  let set_table_list t vs =
    V.Luahash.clear t;
    List.iteri
      (fun i v ->
        let idx = 1.0 +. float_of_int i in
        V.Table.bind t ~key:(Number idx) ~data:v)
      vs

  let add_to_global_list g n key v =
    let t =
      match V.Table.find g.V.globals ~key:n with
      | V.LuaValueBase.Table t -> t
      | _ ->
          let t = V.Table.create 0 in
          setglobal g key (Table t);
          t
    in
    add_to_table_list t v

  let get_global_list g n =
    match V.Table.find g.V.globals ~key:n with
    | V.LuaValueBase.Table t ->
        let rec collect acc idx =
          match V.Table.find t ~key:(Number idx) with
          | V.LuaValueBase.Nil -> List.rev acc
          | v -> collect (v :: acc) (idx +. 1.0)
        in
        collect [] 1.0
    | _ -> []

  let is_lua_table_numbered_values table =
    let exception NotNumbered in
    let exception Numbered in
    try
      let k1, _v1 =
        try V.Table.first table with Not_found -> raise Numbered
      in
      let rec check k =
        match (k : V.value) with
        | Number _ ->
            let k_next, _v_next =
              try V.Table.next table k
              with Not_found ->
                (* done *)
                raise Numbered
            in
            check k_next
        | _ -> raise NotNumbered
      in
      check k1
    with
    | NotNumbered -> false
    | Numbered -> true

  let lua_table_numbered_values table =
    let item k v =
      match (k : V.value) with Number n -> Some (n, v) | _ -> None
    in
    let exception Done of (float * V.value) list in
    let kvs =
      try
        let k1, v1 =
          try V.Table.first table with Not_found -> raise (Done [])
        in
        let rec collect acc k v =
          let acc' = match item k v with Some iv -> iv :: acc | None -> acc in
          let k_next, v_next =
            try V.Table.next table k with Not_found -> raise (Done acc')
          in
          collect acc' k_next v_next
        in
        collect [] k1 v1
      with Done lst -> lst
    in
    List.sort (fun (n1, _) (n2, _) -> compare n1 n2) kvs |> List.map snd

  let lua_table_luaordered_values table =
    let exception Done of V.value list in
    try
      let k1, v1 =
        try V.Table.first table with Not_found -> raise (Done [])
      in
      let rec collect acc k v =
        let acc' = v :: acc in
        let k_next, v_next =
          try V.Table.next table k
          with Not_found -> raise (Done (List.rev acc'))
        in
        collect acc' k_next v_next
      in
      collect [] k1 v1
    with Done lst -> lst

  let lua_table_luaordered_keyvalues table =
    let exception Done of (V.value * V.value) list in
    try
      let k1, v1 =
        try V.Table.first table with Not_found -> raise (Done [])
      in
      let rec collect acc k v =
        let acc' = (k, v) :: acc in
        let k_next, v_next =
          try V.Table.next table k
          with Not_found -> raise (Done (List.rev acc'))
        in
        collect acc' k_next v_next
      in
      collect [] k1 v1
    with Done lst -> lst
end

module RuleNameMap = Map.Make (String)

module AstUtils (A : Lua.AST) = struct
  let return_range ~script_contents chunks srcmap =
    match
      List.find_map
        (function
          | A.Debug _ | A.Fundef (_, _, _, _) | A.Methdef (_, _, _, _, _) ->
              None
          | A.Statement (A.Stmt' (pos, A.Return _)) ->
              let _file, line, col = Luasrcmap.location srcmap pos in
              Some
                (ThunkRanges.range_exact_linecol_in_string ~line ~col
                   script_contents)
          | A.Statement (A.Stmt' (_, _))
          | A.Statement (A.Assign (_, _))
          | A.Statement (A.WhileDo (_, _))
          | A.Statement (A.RepeatUntil (_, _))
          | A.Statement (A.If (_, _, _, _))
          | A.Statement (A.Callstmt _)
          | A.Statement (A.Local (_, _)) ->
              None
          | A.Statement (A.Return _) ->
              (* super-annoying we don't have the position here. need fmlib_parse! *)
              None)
        (List.rev chunks)
    with
    | Some r -> r
    | None ->
        (* default to the whole script *)
        ThunkRanges.inner_range
          (ThunkRanges.raw_range_of_string script_contents)

  let rule_ranges ~script_contents chunks srcmap =
    let current_freerules_varname = ref None in
    let current_uirules_varname = ref None in
    let exception Local of string in
    try
      let ranges =
        List.fold_left
          (fun acc v ->
            match v with
            | A.Statement
                (A.Stmt'
                   ( _stmt_pos,
                     A.Assign
                       ( [ Lvar rules_varname ],
                         [
                           Call
                             (Funcall
                                ( Index (Var "build", Lit (String "newrules")),
                                  _funcall_args ));
                         ] ) )) ->
                (* Recognize: `rules = build['newrules'](M)`
               The AST is (using ThunkLuaPrint.show_chunk ~structure:true ~compliant:true):
                 Statement(Stmt'(Assign(Lvar(rules) = Call(Funcall(Index(Var(build), Lit('newrules')), (Var(M))))))
                
               This is a _heuristic_ to find rule definition statements.
               It would simply be better if we ripped out the Lua parser and replaced it with Fmlib_parse,
               and propagated the AST location into Value.value (at least for functions). *)
                current_freerules_varname := Some rules_varname;
                acc
            | A.Statement
                (A.Stmt'
                   ( _stmt_pos,
                     A.Assign
                       ( [ Lvar rules_varname; Lvar uirules_varname ],
                         [
                           Call
                             (Funcall
                                ( Index (Var "build", Lit (String "newrules")),
                                  _funcall_args ));
                         ] ) )) ->
                (* Recognize: `rules, uirules = build['newrules'](M)`
                
               This is a _heuristic_ to find rule definition statements.
               It would simply be better if we ripped out the Lua parser and replaced it with Fmlib_parse,
               and propagated the AST location into Value.value (at least for functions). *)
                (* Printf.eprintf
                  "\n\n\
                   ********************* Assign to rules=%s and uirules=%s\n\n\
                   %!"
                  rules_varname uirules_varname; *)
                current_freerules_varname := Some rules_varname;
                current_uirules_varname := Some uirules_varname;
                acc
            | A.Fundef
                ( fundef_pos,
                  Lindex (Var lindex_varname, Lit (String rule_name)),
                  _args,
                  _block ) -> (
                (* Recognize: `function rules.CmaLibrary(command, options, kontinue)`
               Take: rule_name = "CmaLibrary"               
               The AST is (using ThunkLuaPrint.show_chunk ~structure:true ~compliant:true):
                  Fundef(function Lindex(Var(rules).Lit(CmaLibrary))(command, options, kontinue)
             *)
                match
                  (!current_freerules_varname, !current_uirules_varname)
                with
                | Some s, _ when String.equal s lindex_varname -> (
                    (* Printf.eprintf
                      "\n\n\
                       ********************* free Fundef (Lindex \
                       (Var(%s).Lit(%s)))\n\n\
                       %!"
                      lindex_varname rule_name; *)
                    let _file, line, col =
                      Luasrcmap.location srcmap fundef_pos
                    in
                    let fundef_range =
                      ThunkRanges.range_exact_linecol_in_string ~line ~col
                        script_contents
                    in
                    (* add as a free rule *)
                    match RuleNameMap.find_opt rule_name acc with
                    | None ->
                        RuleNameMap.add rule_name (`Free, fundef_range) acc
                    | Some _ ->
                        raise
                          (Local
                             (Printf.sprintf "The rule name `%s` was duplicated"
                                rule_name)))
                | _, Some s when String.equal s lindex_varname -> (
                    (* Printf.eprintf
                      "\n\n\
                       ********************* ui Fundef (Lindex \
                       (Var(%s).Lit(%s)))\n\n\
                       %!"
                      lindex_varname rule_name; *)
                    let _file, line, col =
                      Luasrcmap.location srcmap fundef_pos
                    in
                    let fundef_range =
                      ThunkRanges.range_exact_linecol_in_string ~line ~col
                        script_contents
                    in
                    (* add as a ui rule *)
                    match RuleNameMap.find_opt rule_name acc with
                    | None -> RuleNameMap.add rule_name (`Ui, fundef_range) acc
                    | Some _ ->
                        raise
                          (Local
                             (Printf.sprintf "The rule name `%s` was duplicated"
                                rule_name)))
                | _ -> acc)
            | A.Debug _ | A.Fundef (_, _, _, _) | A.Methdef (_, _, _, _, _) ->
                acc
            | A.Statement (A.Stmt' (_, _))
            | A.Statement (A.Assign (_, _))
            | A.Statement (A.WhileDo (_, _))
            | A.Statement (A.RepeatUntil (_, _))
            | A.Statement (A.If (_, _, _, _))
            | A.Statement (A.Callstmt _)
            | A.Statement (A.Local (_, _))
            | A.Statement (A.Return _) ->
                acc)
          RuleNameMap.empty chunks
      in
      Ok ranges
    with Local msg -> Error msg
end

module CoreUtils (V : Lua.VALUE) (C : Lua.Lib.CORE with module V = V) = struct
  open struct
    module R = ThunkLuaErrors.Raise (V)

    let expensive_read_location loc :
        [ `Origin of string ] * ThunkRanges.t * [ `Contents of string ] =
      match (loc : Luasrcmap.location) with
      | file, line, col ->
          In_channel.with_open_bin file (fun ic ->
              let contents = In_channel.input_all ic in
              let range =
                ThunkRanges.range_exact_linecol_in_string ~line ~col contents
              in
              (`Origin file, ThunkRanges.raw_range range, `Contents contents))
  end

  let parse_module_version ?loc ?downgrade_errors_into_warnings g s =
    let invalid semantic =
      R.error g (ThunkResults.Semantic.error_message semantic)
    in
    match ThunkCommand.InternalUse.parse_moduleversion s with
    | Ok mv -> mv
    | Error msg -> begin
        (* reparse after reading the source location for nicer error *)
        match Option.map expensive_read_location loc with
        | None -> R.error g msg
        | Some (`Origin origin, pos, `Contents src) -> (
            let state' =
              ThunkParsers.Results.State.create_with_source ~origin
                ?downgrade_errors_into_warnings src
            in
            match
              ThunkCommand.InternalUse.parse_moduleversion_full
                ~package_id_opt:None
                (module LuaObserver)
                state' (Some pos) s
            with
            | Ok _ -> R.error g msg
            | Error semantic' -> invalid semantic')
      end
end