Skip to main content

lib/rx/renv.ex

defmodule Rx.Renv do
  @moduledoc false

  defmodule Config do
    @moduledoc false
    @enforce_keys [
      :project,
      :lockfile,
      :lockfile_sha256,
      :backend,
      :r_binary,
      :lib_paths,
      :restore,
      :clean,
      :renv_env,
      :identity_env
    ]
    defstruct [
      :project,
      :lockfile,
      :lockfile_sha256,
      :backend,
      :r_binary,
      :lib_paths,
      :restore,
      :clean,
      :renv_env,
      :identity_env,
      resolved_lib_paths: nil
    ]
  end

  @reserved_env_names ~w(RX_RENV_ENABLED RX_RENV_PROJECT RX_RENV_LOCKFILE RENV_PATHS_LOCKFILE)
  @unsupported_native_message "Rx.renv_init/2 currently supports only the process backend; use backend: :process or restart the BEAM with a non-renv native setup."

  def unsupported_native_message, do: @unsupported_native_message

  @doc false
  def __preflight_source_for_test__, do: preflight_source()

  def resolve!(project_or_lockfile, opts) do
    validate_path_arg!(project_or_lockfile)
    opts = validate_opts!(opts)
    backend = process_backend!(opts[:backend])
    r_binary = non_empty_binary_option!(:r_binary, opts[:r_binary])
    lib_paths = lib_paths!(opts[:lib_paths])
    restore = boolean_option!(:restore, opts[:restore])
    clean = boolean_option!(:clean, opts[:clean])

    if clean and not restore do
      raise ArgumentError, "Rx.renv_init/2 option :clean requires restore: true"
    end

    {project, lockfile} = resolve_project_and_lockfile!(project_or_lockfile, opts[:project])
    validate_lockfile_json!(lockfile)
    renv_env = renv_env!(opts[:renv_env])

    %Config{
      project: project,
      lockfile: lockfile,
      lockfile_sha256: lockfile_sha256!(lockfile),
      backend: backend,
      r_binary: r_binary,
      lib_paths: lib_paths,
      restore: restore,
      clean: clean,
      renv_env: renv_env,
      identity_env: Enum.sort_by(renv_env, fn {name, value} -> {name, value} end)
    }
  end

  def ensure_runtime_can_use_process!(%Config{} = config) do
    if Rx.backend() == :native do
      raise RuntimeError,
            "initialized embedded native R cannot be shut down inside the current BEAM; " <>
              "restart the BEAM or Livebook runtime before selecting another backend"
    end

    config
  end

  def preflight!(%Config{} = config) do
    executable = resolve_rscript!(config.r_binary)
    args = ["--vanilla", "-e", preflight_source(), "--args"] ++ preflight_args(config)

    case System.cmd(executable, args, env: preflight_env(config), stderr_to_stdout: true) do
      {output, 0} ->
        %{config | resolved_lib_paths: decode_preflight_output!(output, config)}

      {output, _status} ->
        raise RuntimeError, preflight_error_message(output, config)
    end
  rescue
    error in ErlangError ->
      raise RuntimeError,
            "could not run renv preflight with #{inspect(config.r_binary)}: #{Exception.message(error)}"
  end

  def to_runtime_config(%Config{resolved_lib_paths: lib_paths} = config)
      when is_list(lib_paths) do
    [
      backend: config.backend,
      r_binary: config.r_binary,
      r_home: nil,
      lib_r_path: nil,
      lib_paths: config.lib_paths,
      renv: %{
        project: config.project,
        lockfile: config.lockfile,
        lockfile_sha256: config.lockfile_sha256,
        restore: config.restore,
        clean: config.clean,
        lib_paths: lib_paths,
        env: config.renv_env,
        identity_env: config.identity_env
      }
    ]
  end

  def process_env(%Config{} = config),
    do: startup_env(config.project, config.lockfile, config.renv_env)

  def process_env(%{project: project, lockfile: lockfile, env: renv_env}) do
    startup_env(project, lockfile, renv_env)
  end

  defp validate_path_arg!(path) when is_binary(path) do
    if String.trim(path) == "" do
      raise ArgumentError, "Rx.renv_init/2 project_or_lockfile must be a non-empty binary"
    end
  end

  defp validate_path_arg!(_path) do
    raise ArgumentError, "Rx.renv_init/2 project_or_lockfile must be a non-empty binary"
  end

  defp validate_opts!(opts) when is_list(opts) do
    if Keyword.keyword?(opts) do
      Keyword.validate!(opts,
        backend: :process,
        r_binary: "Rscript",
        lib_paths: [],
        project: nil,
        restore: false,
        clean: false,
        renv_env: []
      )
    else
      raise ArgumentError, "Rx.renv_init/2 options must be a keyword list"
    end
  end

  defp validate_opts!(_opts),
    do: raise(ArgumentError, "Rx.renv_init/2 options must be a keyword list")

  defp process_backend!(backend) when backend in [:process, :port_arrow, Rx.Backends.PortArrow],
    do: Rx.Backends.PortArrow

  defp process_backend!(backend) when backend in [:native, :embedded_nif, Rx.Backends.Native],
    do: raise(ArgumentError, @unsupported_native_message)

  defp process_backend!(backend) do
    raise ArgumentError,
          "Rx.renv_init/2 currently supports only the process backend, got: #{inspect(backend)}"
  end

  defp non_empty_binary_option!(name, value) when is_binary(value) do
    if String.trim(value) == "" do
      raise ArgumentError, "Rx.renv_init/2 option #{inspect(name)} must be a non-empty binary"
    end

    value
  end

  defp non_empty_binary_option!(name, _value) do
    raise ArgumentError, "Rx.renv_init/2 option #{inspect(name)} must be a non-empty binary"
  end

  defp lib_paths!(paths) when is_list(paths) and paths != [] do
    if Enum.all?(paths, &(is_binary(&1) and String.trim(&1) != "")) do
      Enum.map(paths, &Path.expand/1)
    else
      raise ArgumentError,
            "Rx.renv_init/2 option :lib_paths must be a list of non-empty binary paths"
    end
  end

  defp lib_paths!([]), do: []

  defp lib_paths!(_paths),
    do:
      raise(
        ArgumentError,
        "Rx.renv_init/2 option :lib_paths must be a list of non-empty binary paths"
      )

  defp boolean_option!(_name, value) when is_boolean(value), do: value

  defp boolean_option!(name, _value),
    do: raise(ArgumentError, "Rx.renv_init/2 option #{inspect(name)} must be a boolean")

  defp resolve_project_and_lockfile!(path, project_opt) do
    expanded = Path.expand(path)

    cond do
      File.dir?(expanded) ->
        project = expanded
        validate_project_option_for_directory!(project_opt, project)
        lockfile = Path.join(project, "renv.lock")
        ensure_lockfile!(lockfile)
        {project, lockfile}

      File.regular?(expanded) ->
        lockfile = expanded
        project = project_from_lockfile!(lockfile, project_opt)
        ensure_project!(project)
        {project, lockfile}

      project_opt != nil ->
        lockfile = expanded
        ensure_lockfile!(lockfile)
        project = project_from_lockfile!(lockfile, project_opt)
        ensure_project!(project)
        {project, lockfile}

      true ->
        raise ArgumentError, "renv project directory does not exist: #{expanded}"
    end
  end

  defp validate_project_option_for_directory!(nil, _project), do: :ok

  defp validate_project_option_for_directory!(project_opt, project) when is_binary(project_opt) do
    if Path.expand(project_opt) != project do
      raise ArgumentError,
            "Rx.renv_init/2 option :project must match the directory path when the first argument is a directory"
    end
  end

  defp validate_project_option_for_directory!(_project_opt, _project) do
    raise ArgumentError, "Rx.renv_init/2 option :project must be a non-empty binary path"
  end

  defp project_from_lockfile!(lockfile, nil), do: Path.dirname(lockfile)

  defp project_from_lockfile!(_lockfile, project) when is_binary(project) do
    if String.trim(project) == "" do
      raise ArgumentError, "Rx.renv_init/2 option :project must be a non-empty binary path"
    end

    Path.expand(project)
  end

  defp project_from_lockfile!(_lockfile, _project) do
    raise ArgumentError, "Rx.renv_init/2 option :project must be a non-empty binary path"
  end

  defp ensure_project!(project) do
    unless File.dir?(project) do
      raise ArgumentError, "renv project directory does not exist: #{project}"
    end
  end

  defp ensure_lockfile!(lockfile) do
    unless File.regular?(lockfile) do
      raise ArgumentError, "renv lockfile does not exist or is not a regular file: #{lockfile}"
    end
  end

  defp validate_lockfile_json!(lockfile) do
    case Jason.decode(File.read!(lockfile)) do
      {:ok, %{"Packages" => packages}} when is_map(packages) ->
        :ok

      {:ok, _other} ->
        raise ArgumentError,
              "renv lockfile must include a top-level \"Packages\" object: #{lockfile}"

      {:error, error} ->
        raise ArgumentError,
              "invalid renv lockfile JSON at #{lockfile}: #{Exception.message(error)}"
    end
  end

  defp lockfile_sha256!(lockfile) do
    lockfile
    |> File.read!()
    |> then(&:crypto.hash(:sha256, &1))
    |> Base.encode16(case: :lower)
  end

  defp renv_env!(env) when is_map(env), do: env |> Map.to_list() |> renv_env!()

  defp renv_env!(env) when is_list(env) do
    env
    |> Enum.map(fn
      {name, value} when is_binary(name) and is_binary(value) ->
        if name in @reserved_env_names do
          raise ArgumentError,
                "Rx.renv_init/2 option :renv_env must not set #{name}; Rx owns that startup variable"
        end

        {name, value}

      _other ->
        raise ArgumentError,
              "Rx.renv_init/2 option :renv_env must be a map or a list of binary name-value pairs"
    end)
    |> reject_duplicate_env_names!()
  end

  defp renv_env!(_env) do
    raise ArgumentError,
          "Rx.renv_init/2 option :renv_env must be a map or a list of binary name-value pairs"
  end

  defp reject_duplicate_env_names!(env) do
    duplicate_names =
      env
      |> Enum.frequencies_by(fn {name, _value} -> name end)
      |> Enum.filter(fn {_name, count} -> count > 1 end)
      |> Enum.map(fn {name, _count} -> name end)
      |> Enum.sort()

    case duplicate_names do
      [] ->
        env

      names ->
        raise ArgumentError,
              "Rx.renv_init/2 option :renv_env must not include duplicate names: #{Enum.join(names, ", ")}"
    end
  end

  defp startup_env(project, lockfile, env) do
    env ++
      [
        {"RX_RENV_ENABLED", "1"},
        {"RX_RENV_PROJECT", project},
        {"RX_RENV_LOCKFILE", lockfile},
        {"RENV_PATHS_LOCKFILE", lockfile}
      ]
  end

  defp resolve_rscript!(r_binary) do
    executable = System.find_executable(r_binary) || r_binary

    if File.exists?(executable) do
      executable
    else
      raise RuntimeError, "could not start R backend: #{r_binary} not found"
    end
  end

  defp preflight_args(%Config{} = config) do
    [
      config.project,
      config.lockfile,
      if(config.restore, do: "true", else: "false"),
      if(config.clean, do: "true", else: "false")
      | config.lib_paths
    ]
  end

  defp preflight_env(%Config{} = config), do: process_env(config)

  defp decode_preflight_output!(output, config) do
    output
    |> preflight_json_line()
    |> Jason.decode()
    |> case do
      {:ok, %{"lib_paths" => paths}} when is_list(paths) ->
        if Enum.all?(paths, &is_binary/1) do
          Enum.map(paths, &Path.expand/1)
        else
          raise RuntimeError,
                "renv preflight returned non-binary library paths: #{inspect(paths)}"
        end

      {:ok, other} ->
        raise RuntimeError, "renv preflight returned invalid JSON payload: #{inspect(other)}"

      {:error, error} ->
        raise RuntimeError,
              "renv preflight did not return JSON for #{config.project}: " <>
                "#{Exception.message(error)}. R output: #{trim_output(output)}"
    end
  end

  defp preflight_json_line(output) do
    prefix = "RX_RENV_PREFLIGHT_JSON="

    output
    |> String.split("\n")
    |> Enum.find_value(fn line ->
      if String.starts_with?(line, prefix) do
        String.replace_prefix(line, prefix, "")
      end
    end)
    |> case do
      nil -> output
      json -> json
    end
  end

  defp preflight_error_message(output, config) do
    trimmed = trim_output(output)

    cond do
      trimmed =~ "missing required R package: renv" ->
        "R package renv is unavailable for #{config.r_binary}; install renv or pass lib_paths pointing to it. R output: #{trimmed}"

      trimmed =~ "missing required R package: jsonlite" ->
        "jsonlite is unavailable in the resolved renv environment for #{config.project}. R output: #{trimmed}"

      true ->
        "renv preflight failed for project #{config.project} and lockfile #{config.lockfile}. R output: #{trimmed}"
    end
  end

  defp trim_output(output) do
    output = String.trim(output)

    if byte_size(output) > 32_000 do
      binary_part(output, 0, 32_000) <> "\n... output truncated by Rx ..."
    else
      output
    end
  end

  defp preflight_source do
    ~S"""
    # rx_renv_preflight_main
    args <- commandArgs(trailingOnly = TRUE)
    if (length(args) > 0L && identical(args[[1]], "--args")) args <- args[-1L]
    project <- normalizePath(args[[1]], mustWork = TRUE)
    lockfile <- normalizePath(args[[2]], mustWork = TRUE)
    restore <- identical(args[[3]], "true")
    clean <- identical(args[[4]], "true")
    bootstrap_libs <- if (length(args) > 4L) args[5:length(args)] else character()

    if (length(bootstrap_libs) > 0L) .libPaths(c(bootstrap_libs, .libPaths()))
    Sys.setenv(RENV_PATHS_LOCKFILE = lockfile)

    if (!requireNamespace("renv", quietly = TRUE)) {
      cat("missing required R package: renv\n", file = stderr())
      quit(save = "no", status = 42)
    }

    tryCatch(
      renv::lockfile_read(lockfile, project = project),
      error = function(e) {
        cat("invalid renv lockfile: ", conditionMessage(e), "\n", sep = "", file = stderr())
        quit(save = "no", status = 44)
      }
    )

    library_path <- renv::paths$library(project = project)

    if (restore) {
      renv::restore(project = project, lockfile = lockfile, clean = clean, prompt = FALSE)
    } else if (!dir.exists(library_path)) {
      cat(
        "renv project library does not exist; run Rx.renv_init(..., restore: true) or restore manually\n",
        file = stderr()
      )
      quit(save = "no", status = 45)
    }

    suppressWarnings(suppressMessages(renv::load(project, quiet = TRUE)))

    if (!requireNamespace("jsonlite", quietly = TRUE)) {
      cat("missing required R package: jsonlite\n", file = stderr())
      cat("R library paths: ", paste(.libPaths(), collapse = .Platform$path.sep), "\n", file = stderr())
      quit(save = "no", status = 43)
    }

    cat(
      "\nRX_RENV_PREFLIGHT_JSON=",
      jsonlite::toJSON(list(lib_paths = normalizePath(.libPaths(), mustWork = FALSE)), auto_unbox = TRUE),
      "\n",
      sep = ""
    )
    """
  end
end