Source file ShellCoreLua.ml

exception
  LuaEngineShutdown of {
    trace : MlFront_Thunk.BuildWriters.Standard.backtrace_item list;
    exitcode_posix : int;
    exitcode_windows : int;
    lua_source_file : Wiring.BuildContext.V.file_object;
    cant_do : string;
  }
(** An extension of {!DkZero_Base.Exceptions.EngineShutdown} *)

(** {1 Modname} *)

type modname = {
  global_opt : string option;
  modname : string;
  version_opt : string option;
}

let log_app msgf =
  Wiring.BuildContext.run_isolated_promise @@ Wiring.BuildContext.Log.app msgf

let parse_modname =
  let stop s =
    log_app (fun l ->
        l
          "FATAL: Invalid -l argument `%s`. Use -l modname, -l g=modname or -l \
           g=modname@version."
          s);
    exit 1
  in
  let modname_modver global_opt s =
    match String.split_on_char '@' s with
    | [ modname ] -> { global_opt; modname; version_opt = None }
    | [ modname; version ] ->
        { global_opt; modname; version_opt = Some version }
    | _ -> stop s
  in
  fun s ->
    match String.split_on_char '=' s with
    | [ modname ] -> modname_modver None modname
    | [ g; modname ] -> modname_modver (Some g) modname
    | _ -> stop s

(** {1 Lua Arg} *)

type embedded_details = {
  embedded_rule_modid : MlFront_Core.StandardModuleId.t;
  embedded_rule_semver : MlFront_Thunk.ThunkSemver64.t;
  embedded_rule_file_object : Wiring.BuildContext.Io.file_object;
}

(** The type of a Lua argument in nanopass 1 *)
type lua_arg1 =
  | Interpreter1 of string
  | DoString1 of string
  | Require1 of modname
  | LuaScript1 of string
  | UiRule1 of MlFront_Thunk.ThunkCommand.module_version
  | Arg1 of string

(** The type of a Lua argument in nanopass 2 *)
type lua_arg2 =
  | Interpreter2 of string
  | DoString2 of string
  | Require2 of modname
  | LuaScript2 of string
  | UiRule2 of {
      rule_modver : MlFront_Thunk.ThunkCommand.module_version;
      user_request_json : MlFront_Thunk.ThunkJson.t;
    }
  | FileWithUiRuleEmbedded2 of {
      file : MlFront_Core.FilePath.t;
      embedded_lines : (Fmlib_parse.Position.range * string) list;
      embedded_details : embedded_details;
    }
  | Arg2 of string

(** The type of a Lua argument in nanopass 3 *)
type lua_arg3 =
  | DoString3 of string
  | Require3 of modname
  | LuaScript3 of string
  | UiRule3 of {
      ctx : Wiring.BuildContext.t;
      build_state : Wiring.BuildContext.State.t;
      rule_modver : MlFront_Thunk.ThunkCommand.module_version;
      source_file : Wiring.BuildContext.Io.file_object;
      tasks : (module Wiring.BuildContext.THUNK_TASKS);
      user_request_json : MlFront_Thunk.ThunkJson.t;
    }
  | FileWithUiRuleEmbedded3 of {
      ctx : Wiring.BuildContext.t;
      build_state : Wiring.BuildContext.State.t;
      file : MlFront_Core.FilePath.t;
      embedded_lines : (Fmlib_parse.Position.range * string) list;
      tasks : (module Wiring.BuildContext.THUNK_TASKS);
      embedded_details : embedded_details;
    }

(** {1 Build} *)

open struct
  module AU =
    MlFront_Thunk.ThunkLuaUtils.AstUtils (MlFront_Thunk.ThunkLuaScript.Ast)
end

let quick_error ~command_name msg =
  log_app (fun l ->
      l "FATAL: The `%s` subcommand failed.@.%s" command_name
        (MlFront_Thunk.ThunkStrings.left_pad msg));
  exit 1

module type RUNNER = sig
  type state
  type value
  type chunk
  type analysis

  val mk : unit -> state

  val getchunks_from_file :
    lua_file:string ->
    (chunk list, string * [ `LineCol of int * int ] option) result

  val dostring :
    exit_on_error:bool -> sourcename:string -> state -> string -> value list

  val dostring_analysis : sourcename:string -> state -> string -> analysis list
  val dofile : state -> string -> value list
  val dofile_analysis : state -> string -> analysis list

  val run_interactive :
    ?prompt:string -> sourcename:string -> in_channel -> state -> unit
end

(* Similar to luarun.ml *)
module MakeRunner
    (InWhat : sig
      val in_what : string
      val command_name : string
    end)
    (I :
      MlFront_Lua.Lua.INTERP
        with module Value = MlFront_Thunk.ThunkLuaScript.Value
         and module Ast = MlFront_Thunk.ThunkLuaScript.Ast) :
  RUNNER
    with type state = I.state
     and type value = I.value
     and type chunk = I.Ast.chunk
     and type analysis = I.value MlFront_Thunk.ThunkLuaTypAnalysis.t = struct
  type state = I.state
  type value = I.value
  type chunk = I.Ast.chunk
  type analysis = I.value MlFront_Thunk.ThunkLuaTypAnalysis.t

  module IoUtils = MlFront_Thunk.ThunkLuaUtils.IoUtils' (I)
  module ScriptUtils = MlFront_Thunk.ThunkLuaScript.ScriptUtils (I)

  let quick_error msg = quick_error ~command_name:InWhat.command_name msg
  let mk () = I.mk ()

  let getchunks_from_file ~lua_file =
    (* Partial implementation of Luabasefile.dofile except stop at chunks,
       before compiling them. *)
    let g = I.mk () in
    let map = MlFront_Lua.Luasrcmap.mk () in
    MlFront_Lua.Luasrcmap.sync map 0 (lua_file, 1, 1);
    In_channel.with_open_bin lua_file (fun ic ->
        I.with_stack
          (MlFront_Thunk.ThunkLuaScript.Value.srcloc ~file:lua_file
             ~linedefined:0)
          g
          (IoUtils.lex_into_chunks map)
          (Lexing.from_channel ic))

  let dostring_full ~exit_on_error ~sourcename g s =
    let file =
      Printf.sprintf "dostring(%s)"
        (MlFront_Thunk.ThunkLuaPrint0.literal_string
           (MlFront_Thunk.ThunkStrings.abbreviate s))
    in
    I.with_stack
      (I.Value.srcloc ~file ~linedefined:0)
      g
      (fun lexbuf ->
        match
          IoUtils.do_lexbuf_return_all ~sourcename ~in_what:InWhat.in_what g
            lexbuf
        with
        | Error (msg, _loc_opt) ->
            if exit_on_error then
              quick_error
                (Printf.sprintf "In:\n%s\ngot: %s"
                   (MlFront_Thunk.ThunkStrings.left_pad (String.trim s))
                   msg)
            else
              (* Nil + Error [+ Code] is Lua best practice.
               https://github.com/luarocks/lua-style-guide?tab=readme-ov-file#errors *)
              ( MlFront_Thunk.ThunkRanges.inner_range
                  (MlFront_Thunk.ThunkRanges.raw_range_of_string s),
                MlFront_Thunk.ThunkLuaUtils.RuleNameMap.empty,
                [ I.Value.LuaValueBase.Nil; I.Value.LuaValueBase.String msg ] )
        | Ok (chunks, srcmap, v) -> (
            let return_range =
              AU.return_range ~script_contents:s chunks srcmap
            in
            let ranges_result =
              AU.rule_ranges ~script_contents:s chunks srcmap
            in
            match ranges_result with
            | Error msg ->
                if exit_on_error then quick_error msg
                else
                  ( return_range,
                    MlFront_Thunk.ThunkLuaUtils.RuleNameMap.empty,
                    [
                      I.Value.LuaValueBase.Nil; I.Value.LuaValueBase.String msg;
                    ] )
            | Ok rule_ranges -> (return_range, rule_ranges, v)))
      (Lexing.from_string s)

  let dostring ~exit_on_error ~sourcename g s =
    match dostring_full ~exit_on_error ~sourcename g s with
    | _ret_range, _rule_ranges, v -> v

  let dofile_full g lua_file =
    let script_contents =
      In_channel.with_open_bin lua_file In_channel.input_all
    in
    try
      (* evaluate the script *)
      let file =
        Printf.sprintf "dofile(%s)"
          (MlFront_Thunk.ThunkLuaPrint0.literal_string lua_file)
      in
      I.with_stack
        (MlFront_Thunk.ThunkLuaScript.Value.srcloc ~file ~linedefined:0)
        g
        (fun lexbuf ->
          match
            IoUtils.do_lexbuf_return_all ~sourcename:lua_file
              ~in_what:InWhat.in_what g lexbuf
          with
          | Error (msg, Some (`LineCol (line, col))) ->
              let script_content =
                In_channel.with_open_bin lua_file In_channel.input_all
              in
              let _range, rendered =
                MlFront_Thunk.ThunkLuaUtils.render_lua_error
                  ?downgrade_errors_into_warnings:None ~origin:lua_file
                  ~code:"2fc2a99b" ~line ~col ~script_content ~msg ()
              in
              quick_error rendered
          | Error (msg, None) -> quick_error msg
          | Ok (chunks, srcmap, v) ->
              let return_range =
                AU.return_range ~script_contents chunks srcmap
              in
              let rule_ranges = AU.rule_ranges ~script_contents chunks srcmap in
              (return_range, rule_ranges, v))
        (Lexing.from_string script_contents)
    with Failure msg -> quick_error msg

  let dofile g lua_file =
    let _ret_range, _rule_ranges, v = dofile_full g lua_file in
    v

  let dostring_analysis ~sourcename g s =
    let return_range, rule_ranges, return_values =
      dostring_full ~exit_on_error:true ~sourcename g s
    in
    (* incorporate the return values from the script into the analysis *)
    (match
       ScriptUtils.include_script_outputs_in_analysis ~return_range ~rule_ranges
         return_values g
     with
    | Error msg -> quick_error msg
    | Ok () -> ());
    (* get the analysis from the [_analysis] variable *)
    let getresult = ScriptUtils.get_analysis g in
    match getresult with
    | Error msg -> quick_error msg
    | Ok typedvalues -> typedvalues

  let dofile_analysis g lua_file =
    let return_range, rule_ranges_result, return_values =
      dofile_full g lua_file
    in
    match rule_ranges_result with
    | Error msg -> quick_error msg
    | Ok rule_ranges -> (
        (* incorporate the return values from the script into the analysis *)
        (match
           ScriptUtils.include_script_outputs_in_analysis ~return_range
             ~rule_ranges return_values g
         with
        | Error msg -> quick_error msg
        | Ok () -> ());
        (* get the analysis from the [_analysis] variable *)
        let getresult = ScriptUtils.get_analysis g in
        match getresult with
        | Error msg -> quick_error msg
        | Ok typedvalues -> typedvalues)

  let print_values ~to_string ~label values =
    List.iter (fun v -> Printf.printf "%s%s\n" label (to_string v)) values

  let run_interactive ?prompt ~sourcename infile g =
    let rec loop n pfx =
      (match prompt with
      | Some prompt_str -> Printf.printf "%s%d> %!" prompt_str n
      | None -> ());
      let line = input_line infile in
      if String.length line > 0 && line.[String.length line - 1] = '\\' then
        loop n (pfx ^ String.sub line 0 (String.length line - 1) ^ "\n")
      else begin
        (* tolerate errors in the REPL. *)
        let values =
          dostring ~exit_on_error:false ~sourcename g (pfx ^ line ^ "\n")
        in
        print_values ~label:"(return) "
          ~to_string:MlFront_Thunk.ThunkLuaPrint.show_value values;
        flush stdout;
        flush stderr;
        loop (n + 1) ""
      end
    in
    try loop 1 "" with End_of_file -> ()
end

open struct
  let start_phase3 ctx ~traces ~lua_valuescan_origin ~prefetch2_lua_scripts
      module_or : Wiring.ShellCore.phase3 =
    (* Load state *)
    let initiator =
      DkZero_Base.BuildRequest.UserInitiated
        { agent = "dk0 lua"; request_slot = None }
    in
    let state, tasks, prefetch_keys =
      Wiring.BuildEngine.load_state_and_tasks_gracefully ~lua_valuescan_origin
        ~lua_allow_fetch:true ctx ~traces module_or
    in

    { ctx; initiator; state; tasks; prefetch_keys; prefetch2_lua_scripts }
end

let mk_source cmdline =
  let contents =
    List.map
      MlFront_Thunk.ThunkLexers.ValueShellLexer.Token.quote_literal_if_needed
      cmdline
    |> String.concat " "
  in
  let file =
    Wiring.BuildContext.Io.inmemory_file
      ~origin:
        (MlFront_Core.FilePath.append_exn
           MlFront_Thunk.ThunkIo.shmdir_for_inmem_filesystem "argv")
      contents
  in
  match
    Wiring.BuildContext.run_isolated_promise
      (Wiring.BuildContext.Io.checksum_file ~algo:`SHA256 file)
  with
  | `Error msg -> ShellCore.quick_error msg
  | `Checksum sha256 -> (file, sha256)

let start_config ~baseconfig ~install ~random_seed ~verbosity ~progress
    ~nobuiltininc ~nosysinc ~noworkspaceinc ~sysincludedirs ~userincludedirs
    ~cells ~local_packages ~build_number ~long_ids ~import debugmodes module_or
    =
  (* Start phase 1 *)
  let ({ preconfig } : Wiring.ShellCore.phase1) =
    ShellVSL.start_phase1 ~baseconfig ~random_seed ~cells ~install ()
  in

  (* Start phase 2 *)
  let config =
    ShellVSL.start_phase2 ~preconfig ~autofix:false ~verbosity ~progress
      ~nobuiltininc ~nosysinc ~noworkspaceinc ~sysincludedirs ~userincludedirs
      ~local_packages ~build_number ~long_ids ~import debugmodes module_or
  in
  (preconfig, config)

let start_build ~preconfig ctx ~command_name ~lua_valuescan_origin
    ~prefetch2_lua_scripts ~wait_trace_store ~invalidations parent_txn =
  let baseconfig = DkZero_Base.Config.preconfig_baseconfig preconfig in
  let newgen, state_after_run, lua_interpreter, shell =
    (* Start transaction with trace file *)
    Txn.with_child_txn ~mode:`Create ~wait:wait_trace_store parent_txn
      baseconfig (fun txn ->
        let tracefd = Txn.tracefd txn in

        (* Load traces. *)
        let traces =
          Invalidations.load_traces_gracefully ~preconfig ~reader_generation:0
            ~invalidations tracefd
        in

        (* Start phase 3 (our own) *)
        let shell =
          start_phase3 ctx ~traces ~lua_valuescan_origin ~prefetch2_lua_scripts
            ()
        in

        (* Start phase 4 *)
        let shell =
          let state2 = ShellVSL.start_phase4 ~shell ~baseconfig () in
          { shell with state = state2 }
        in

        (* Run modules (the lua scripts; we need them compiled and registered for Lua [require]) *)
        let shell =
          let lua_script_keys =
            List.map
              (fun (modid, version) ->
                Wiring.BuildContext.K.create_for_scriptmodule
                  ~debug_reference:None ~module_id:modid ~module_semver:version
                  ())
              prefetch2_lua_scripts
          in
          let kont =
            Wiring.BuildTaskUnresolved.run_multiple ctx ~tasks:shell.tasks
              lua_script_keys
          in
          let state' =
            Wiring.BuildTaskUnresolved.run_unit_continuation kont shell.state
          in
          { shell with state = state' }
        in

        (* Create Lua runner from build system Lua interpreter that can
             fetch into the build system *)
        let lua_interpreter =
          Wiring.BuildContext.State.lua_interpreter shell.state
        in

        (* Finish phase 1 *)
        let state_after_run = shell.state in

        let (newgen : int) =
          ShellVSL.finish_phase1 shell.ctx state_after_run tracefd
        in
        (newgen, state_after_run, lua_interpreter, shell))
  in

  (* Get Lua state and Lua accessors *)
  let module ForBuild = struct
    let in_what = "in the build"
    let command_name = command_name
  end in
  let module I = (val lua_interpreter) in
  let module Runner = MakeRunner (ForBuild) (I) in
  let g = Wiring.BuildContext.State.lua_state state_after_run in
  let doscript lua_file =
    let (_ : Runner.value list) = Runner.dofile g lua_file in
    ()
  in
  ( g,
    Runner.dostring,
    doscript,
    Runner.run_interactive,
    shell.tasks,
    shell.state,
    newgen )

(** {1 Arg table} *)

(*** Set the 'arg' global table. https://www.lua.org/manual/5.4/manual.html#7 *)
let create_argtable_statement lua_args =
  (* Lua 5.4 has negative indices. *)
  let lua_argv = Array.of_list (lua_args : lua_arg2 list) in
  let script_index =
    let rec aux j =
      if j >= Array.length lua_argv then None
      else
        match lua_argv.(j) with
        | LuaScript2 _ | FileWithUiRuleEmbedded2 _ -> Some j
        | _ -> aux (j + 1)
    in
    aux 0
  in
  let offset = match script_index with Some idx -> 0 - idx | None -> 0 in
  let i = ref 0 in

  let buf = Buffer.create 256 in
  Buffer.add_string buf "arg = {}\n";
  let l s =
    Buffer.add_string buf (MlFront_Thunk.ThunkLuaPrint0.literal_string s)
  in
  let idx () =
    Buffer.add_string buf (Printf.sprintf "arg[%d] = " (!i + offset))
  in
  let optopt name =
    idx ();
    l name;
    Buffer.add_string buf "\n"
  in
  let optval value =
    idx ();
    l value;
    Buffer.add_string buf "\n"
  in
  while !i < List.length lua_args do
    (match lua_argv.(!i) with
    | Interpreter2 s -> optopt s
    | DoString2 s ->
        optopt "-e";
        i := !i + 1;
        optval s
    | Require2 { global_opt; modname; version_opt } -> (
        optopt "-l";
        i := !i + 1;
        let modnamver =
          match version_opt with Some v -> modname ^ "@" ^ v | None -> modname
        in
        match global_opt with
        | Some g -> optval (g ^ "=" ^ modnamver)
        | None -> optval modnamver)
    | LuaScript2 s -> optval s
    | UiRule2 { rule_modver; user_request_json = _ } ->
        optval (MlFront_Thunk.ThunkCommand.show_module_version rule_modver)
    | FileWithUiRuleEmbedded2 { file; embedded_lines = _; embedded_details = _ }
      ->
        optval (MlFront_Core.FilePath.to_string file)
    | Arg2 s -> optval s);
    i := !i + 1
  done;
  Buffer.contents buf

(** {1 Evaluation} *)

(** Evaluate scripts and libraries. Raises {!LuaEngineShutdown} on errors. *)
let evaluate_args ~command_name ~dostring ~doscript g lua_args =
  List.fold_left
    (fun acc_state ->
      (function
      | DoString3 s ->
          let (_ : MlFront_Thunk.ThunkLuaScript.Value.value list) =
            dostring ~exit_on_error:true ~sourcename:"<cli>" g s
          in
          acc_state
      | Require3 { global_opt; modname; version_opt } ->
          let lua_require_statement =
            match version_opt with
            | None ->
                Printf.sprintf "%srequire(%s)"
                  (match global_opt with Some g -> g ^ " = " | None -> "")
                  (MlFront_Thunk.ThunkLuaPrint0.literal_string modname)
            | Some version ->
                Printf.sprintf
                  "local _build_imported = require(%s)\n\
                   %s_build_imported.at(%s)"
                  (MlFront_Thunk.ThunkLuaPrint0.literal_string modname)
                  (match global_opt with Some g -> g ^ " = " | None -> "")
                  (MlFront_Thunk.ThunkLuaPrint0.literal_string version)
          in
          let (_ : MlFront_Thunk.ThunkLuaScript.Value.value list) =
            dostring ~exit_on_error:true ~sourcename:"<cli>" g
              lua_require_statement
          in
          acc_state
      | LuaScript3 lua_file ->
          if Sys.file_exists lua_file then (
            doscript lua_file;
            acc_state)
          else
            let msg =
              Printf.sprintf "Lua script file `%s` not found." lua_file
            in
            quick_error ~command_name msg
      | UiRule3
          {
            ctx;
            build_state;
            rule_modver;
            source_file;
            tasks;
            user_request_json;
          } -> begin
          let build_state = Option.value ~default:build_state acc_state in
          let latest_cant_do = ref "run UI rule" in
          try
            let k_kont =
              let open Wiring.BuildContext.Syntax in
              (* Instantiate rule *)
              let* apply_aliases =
                Wiring.BuildTraceStore.get_apply_aliases ctx
              in
              Wiring.BuildTaskUnresolved.instantiate_rule ctx ~tasks
                ~ui_rights:`AnyRuleTypeAllowed
                ~apply_aliases:(Some apply_aliases) ~debug_reference:None
                ~module_id:(MlFront_Thunk.ThunkCommand.modver_id rule_modver)
                ~module_semver:
                  (MlFront_Thunk.ThunkCommand.modver_semver rule_modver)
                ~hostscript_filepath:None ~user_request_json
                ~user_request_slot:None ()
            in
            let key, state' =
              Wiring.BuildTaskUnresolved.run_continuation k_kont build_state
            in
            (* Fetch/run the rule *)
            let kont = Wiring.BuildTaskUnresolved.run_single ctx ~tasks key in
            let state'' : Wiring.BuildContext.State.t =
              Wiring.BuildTaskUnresolved.run_unit_continuation kont state'
            in
            Some state''
          with
          | DkZero_Base.Exceptions.EngineShutdown
              { trace; exitcode_posix; exitcode_windows }
          ->
            (* re-raise with richer information *)
            let bt = Printexc.get_raw_backtrace () in
            Printexc.raise_with_backtrace
              (LuaEngineShutdown
                 {
                   trace;
                   exitcode_posix;
                   exitcode_windows;
                   lua_source_file = source_file;
                   cant_do = !latest_cant_do;
                 })
              bt
        end
      | FileWithUiRuleEmbedded3
          {
            ctx;
            build_state;
            file = hostscript_filepath;
            embedded_lines = _;
            tasks;
            embedded_details =
              {
                embedded_rule_modid;
                embedded_rule_semver;
                embedded_rule_file_object;
              };
          } -> begin
          let build_state = Option.value ~default:build_state acc_state in
          let latest_cant_do = ref "run embedded UI rule" in
          try
            let k_kont =
              let open Wiring.BuildContext.Syntax in
              (* Instantiate rule *)
              let* apply_aliases =
                Wiring.BuildTraceStore.get_apply_aliases ctx
              in
              Wiring.BuildTaskUnresolved.instantiate_rule ctx ~tasks
                ~ui_rights:`AnyRuleTypeAllowed
                ~apply_aliases:(Some apply_aliases) ~debug_reference:None
                ~module_id:embedded_rule_modid
                ~module_semver:embedded_rule_semver
                ~hostscript_filepath:(Some hostscript_filepath)
                ~user_request_json:(`Assoc []) ~user_request_slot:None ()
            in
            let key, state' =
              Wiring.BuildTaskUnresolved.run_continuation k_kont build_state
            in
            (* Fetch/run the rule *)
            let kont = Wiring.BuildTaskUnresolved.run_single ctx ~tasks key in
            let state'' : Wiring.BuildContext.State.t =
              Wiring.BuildTaskUnresolved.run_unit_continuation kont state'
            in
            Some state''
          with
          | DkZero_Base.Exceptions.EngineShutdown
              { trace; exitcode_posix; exitcode_windows }
          ->
            (* re-raise with richer information *)
            let bt = Printexc.get_raw_backtrace () in
            Printexc.raise_with_backtrace
              (LuaEngineShutdown
                 {
                   trace;
                   exitcode_posix;
                   exitcode_windows;
                   lua_source_file = embedded_rule_file_object;
                   cant_do = !latest_cant_do;
                 })
              bt
        end))
    None lua_args