arf-console 0.7.0-beta.1

//! Restarting the process, optionally with a different R version.

#[cfg(test)]
mod tests;

use crate::external::rig::{self, RigError};
use std::ffi::{OsStr, OsString};
use std::process::Command;
#[cfg(not(unix))]
use std::process::Stdio;

use crate::repl::{arf_eprintln, arf_println};

/// Environment variables whose R-specific values must not leak across `:switch`.
///
/// Values inherited at arf startup are restored for all variables except the
/// names in `ALWAYS_REMOVE_ENV_VARS`. Values absent at startup are removed so
/// that the new R process can compute them afresh.
/// The comments record where each value comes from when the user did not
/// supply one, which is what makes it version-specific.
const R_VERSION_ENV_VARS: &[&str] = &[
    // Always removed, whatever its origin; see `ALWAYS_REMOVE_ENV_VARS`.
    "R_HOME",
    // Always removed, whatever its origin; see `ALWAYS_REMOVE_ENV_VARS`.
    "LD_LIBRARY_PATH",
    // Always removed so rig reevaluates project activation on the new startup.
    "RVENV",
    "RVENV_UNSYNCED",
    #[cfg(target_os = "macos")]
    // Restored from the startup snapshot before the next R target is selected.
    "DYLD_LIBRARY_PATH",
    // Read from the old R_HOME/etc/Renviron.
    "R_LIBS_USER",
    "R_LIBS_SITE",
    // Not set by R itself; may be supplied by the user or the session.
    "R_LIBS",
    "R_SYSTEM_ABI",
    // Set from the old R_HOME/bin/R wrapper script.
    "R_DOC_DIR",
    "R_SHARE_DIR",
    "R_INCLUDE_DIR",
];

/// Environment variables removed before a version switch whatever the startup
/// snapshot holds.
///
/// `R_HOME` is removed because arf always sets it from the R version it
/// resolved, so any inherited value is meaningless and restoring it would be
/// wrong. `LD_LIBRARY_PATH` keeps its existing removal policy. Rig's
/// activation markers are removed so the new startup reevaluates the project.
///
/// TODO: Let `LD_LIBRARY_PATH` follow the same restore-or-remove rule as the
/// rest of `R_VERSION_ENV_VARS`, so a value the user set for unrelated
/// libraries survives a switch instead of being dropped. That changes what
/// users observe, so it belongs in a change of its own.
const ALWAYS_REMOVE_ENV_VARS: &[&str] = &["LD_LIBRARY_PATH", "R_HOME", "RVENV", "RVENV_UNSYNCED"];

#[derive(Default)]
struct EnvChanges {
    restored: Vec<(&'static str, OsString)>,
    removed: Vec<&'static str>,
}

/// Return the environment changes needed before restarting.
///
/// The startup snapshot is injected so this policy can be tested without
/// depending on the process environment. Variables in
/// `ALWAYS_REMOVE_ENV_VARS` are always removed for their distinct policies:
/// `R_HOME` is always set by arf from the resolved R version, while
/// `LD_LIBRARY_PATH` remains removed to preserve current behaviour even though
/// the carrier preserves a user-set startup value across re-execs.
fn env_changes_for_switch<F>(version: Option<&str>, mut startup_value: F) -> EnvChanges
where
    F: FnMut(&str) -> Option<OsString>,
{
    let mut changes = EnvChanges::default();

    if version.is_none() {
        return changes;
    }

    for &var in R_VERSION_ENV_VARS {
        if ALWAYS_REMOVE_ENV_VARS.contains(&var) {
            changes.removed.push(var);
        } else if let Some(value) = startup_value(var) {
            changes.restored.push((var, value));
        } else {
            changes.removed.push(var);
        }
    }

    changes
}

fn arf_supplied_wrapper_paths_to_remove<FStartup, FCurrent>(
    wrapper_values: impl IntoIterator<Item = (&'static str, OsString)>,
    mut startup_value: FStartup,
    mut current_value: FCurrent,
) -> Vec<&'static str>
where
    FStartup: FnMut(&str) -> Option<OsString>,
    FCurrent: FnMut(&str) -> Option<OsString>,
{
    wrapper_values
        .into_iter()
        .filter_map(|(name, expected)| {
            (startup_value(name).is_none() && current_value(name).as_ref() == Some(&expected))
                .then_some(name)
        })
        .collect()
}

/// Build the command used to restart the process with the requested changes.
///
/// Using `Command::env_remove()` (rather than mutating the current process's
/// own environment before spawning/exec'ing) keeps this scoped to the child:
/// on the non-Unix spawn path the parent process stays alive after this call
/// returns, so mutating its environment directly would leak into it.
fn build_restart_command(
    exe: &std::path::Path,
    args: &[OsString],
    changes: &EnvChanges,
    startup_env_carrier: &str,
) -> Command {
    let mut cmd = Command::new(exe);
    cmd.args(args);
    cmd.env(crate::STARTUP_ENV_CARRIER, startup_env_carrier);
    #[cfg(unix)]
    if let Some(fd) = crate::pid_file::restart_fd_carrier() {
        cmd.env(crate::pid_file::RESTART_PID_FD_ENV, fd);
    }
    for (var, value) in &changes.restored {
        cmd.env(var, value);
    }
    for var in &changes.removed {
        cmd.env_remove(var);
    }
    cmd
}

/// Print the environment changes that will affect the restarted process.
fn print_env_changes(changes: &EnvChanges) {
    let removed: Vec<&str> = changes
        .removed
        .iter()
        .copied()
        .filter(|var| std::env::var_os(var).is_some())
        .collect();

    if changes.restored.is_empty() && removed.is_empty() {
        return;
    }

    let mut details = Vec::new();
    if !changes.restored.is_empty() {
        details.push(format!(
            "restored: {}",
            changes
                .restored
                .iter()
                .map(|(var, _)| *var)
                .collect::<Vec<_>>()
                .join(", ")
        ));
    }
    if !removed.is_empty() {
        details.push(format!("removed: {}", removed.join(", ")));
    }

    arf_println!(
        "Environment variables for the R restart: {}",
        details.join("; ")
    );
}

/// Restart the process, optionally with a new R version.
///
/// This function uses exec() to replace the current process with a new instance.
/// If version is specified, it resolves the R_HOME using rig before restarting.
///
/// This function only returns if exec fails.
pub fn restart_process(version: Option<&str>) {
    // If a version is specified, validate it using rig before restarting
    if let Some(ver) = version {
        if let Err(message) = validate_rig_for_switch_with(rig::rig_available) {
            for line in message.lines() {
                arf_eprintln!("{}", line);
            }
            return;
        }

        // Validate the version exists before restarting
        match rig::resolve_version(ver) {
            Ok(resolved) => {
                log::info!(
                    "Switching to R version {} ({})",
                    resolved.version,
                    resolved.r_home
                );
            }
            Err(e) => {
                arf_eprintln!("Error: {}", e);
                return;
            }
        }
    }

    // Get the current executable
    let exe = match std::env::current_exe() {
        Ok(e) => e,
        Err(e) => {
            arf_eprintln!("Error: Failed to get current executable: {}", e);
            return;
        }
    };

    // Get command-line arguments (skip the program name, we'll use current_exe instead)
    let mut args: Vec<OsString> = crate::normalized_args();

    if let Some(v) = &version {
        args = args_with_r_version(args, v);
    }

    let mut changes = env_changes_for_switch(version, crate::startup_env_value);
    if version.is_none()
        && let Some(old_r_home) = std::env::var_os("R_HOME")
    {
        changes.removed.extend(arf_supplied_wrapper_paths_to_remove(
            arf_libr::r_path_vars_from_wrapper(std::path::Path::new(&old_r_home)),
            crate::startup_env_value,
            |name| std::env::var_os(name),
        ));
    }
    print_env_changes(&changes);
    let startup_env_carrier = crate::startup_env_carrier();

    // Build the command
    #[cfg(unix)]
    {
        use std::os::unix::process::CommandExt;

        let _pid_fd_guard = match crate::pid_file::prepare_pid_fd_for_exec() {
            Some(Ok(guard)) => Some(guard),
            Some(Err(error)) => {
                arf_eprintln!("Error: Failed to prepare PID file for restart: {error}");
                return;
            }
            None => None,
        };
        let mut cmd = build_restart_command(&exe, &args, &changes, &startup_env_carrier);

        // exec() replaces the current process - this should not return
        let err = cmd.exec();
        arf_eprintln!("Error: Failed to restart: {}", err);
    }

    #[cfg(not(unix))]
    {
        // A spawned replacement cannot safely share the parent's PID file.
        // Stop serving first, then relinquish the file; on failure the old
        // file is left alone and no child is started.
        crate::ipc::stop_server();
        if let Some(pid_path) = crate::pid_file::initial_pid_file_path()
            && let Err(e) = crate::pid_file::relinquish_pid_file_for_restart(&pid_path)
        {
            arf_eprintln!("Error: {}", e);
            std::process::exit(1);
        }
        // On non-Unix platforms, spawn a new process and wait for it to exit.
        // Unlike Unix exec(), this doesn't replace the current process, so we
        // wait for the child and then exit with its status code. This keeps the
        // parent alive to hold the terminal session, preventing the shell from
        // reclaiming control.
        match build_restart_command(&exe, &args, &changes, &startup_env_carrier)
            .stdin(Stdio::inherit())
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .spawn()
        {
            Ok(mut child) => match child.wait() {
                Ok(status) => std::process::exit(status.code().unwrap_or(1)),
                Err(e) => {
                    arf_eprintln!("Error: Failed to wait for restarted process: {}", e);
                    std::process::exit(1);
                }
            },
            Err(e) => {
                arf_eprintln!("Error: Failed to restart: {}", e);
                std::process::exit(1);
            }
        }
    }
}

fn validate_rig_for_switch_with<FAvailable>(rig_available: FAvailable) -> Result<(), String>
where
    FAvailable: FnOnce() -> Result<(), RigError>,
{
    rig_available().map_err(|error| match error {
        RigError::NotInstalled => {
            "Error: rig is not installed. Cannot switch R versions.\nInstall rig from https://github.com/r-lib/rig".to_string()
        }
        RigError::CommandFailed(reason) => format!(
            "Error: rig is installed but failed while checking availability: {reason}\nCannot switch R versions until rig is working."
        ),
        error => format!(
            "Error: Could not check whether rig is available: {error}\nCannot switch R versions until rig is working."
        ),
    })
}

/// Filter explicit R source selections before adding the new R version.
/// Preserve all arguments after `--` because they belong to the user program.
fn filter_r_version_args(args: Vec<OsString>) -> Vec<OsString> {
    let mut result = Vec::new();
    let mut index = 0;

    while index < args.len() {
        let arg = &args[index];
        if arg == "--" {
            result.extend(args[index..].iter().cloned());
            break;
        }

        if arg == "--with-r-version" || arg == "--r-provider" {
            index = (index + 2).min(args.len());
            continue;
        }

        if arg_has_ascii_prefix(arg, b"--with-r-version=")
            || arg_has_ascii_prefix(arg, b"--r-provider=")
        {
            index += 1;
            continue;
        }

        result.push(arg.clone());
        index += 1;
    }

    result
}

fn args_with_r_version(args: Vec<OsString>, version: &str) -> Vec<OsString> {
    let mut args = filter_r_version_args(args);
    let separator = args
        .iter()
        .position(|arg| arg == "--")
        .unwrap_or(args.len());
    args.splice(
        separator..separator,
        [OsString::from("--with-r-version"), OsString::from(version)],
    );
    args
}

fn arg_has_ascii_prefix(arg: &OsStr, prefix: &[u8]) -> bool {
    #[cfg(unix)]
    {
        use std::os::unix::ffi::OsStrExt;
        arg.as_bytes().starts_with(prefix)
    }
    #[cfg(windows)]
    {
        use std::os::windows::ffi::OsStrExt;
        arg.encode_wide()
            .take(prefix.len())
            .eq(prefix.iter().copied().map(u16::from))
    }
    #[cfg(not(any(unix, windows)))]
    {
        arg.to_str()
            .is_some_and(|arg| arg.as_bytes().starts_with(prefix))
    }
}