shed/src/main.rs

#![allow(
  clippy::derivable_impls,
  clippy::tabs_in_doc_comments,
  clippy::while_let_on_iterator,
	clippy::result_large_err
)]
pub mod builtin;
pub mod expand;
pub mod getopt;
pub mod jobs;
pub mod libsh;
pub mod parse;
pub mod prelude;
pub mod procio;
pub mod readline;
pub mod shopt;
pub mod signal;
pub mod state;
#[cfg(test)]
pub mod tests;

use std::os::fd::BorrowedFd;
use std::process::ExitCode;
use std::sync::atomic::Ordering;

use nix::errno::Errno;
use nix::poll::{PollFd, PollFlags, PollTimeout, poll};
use nix::unistd::read;

use crate::builtin::keymap::KeyMapMatch;
use crate::builtin::trap::TrapTarget;
use crate::libsh::error::{self, ShErr, ShErrKind, ShResult};
use crate::libsh::sys::TTY_FILENO;
use crate::parse::execute::exec_input;
use crate::prelude::*;
use crate::readline::term::{LineWriter, RawModeGuard, raw_mode};
use crate::readline::{Prompt, ReadlineEvent, ShedVi};
use crate::signal::{GOT_SIGWINCH, JOB_DONE, QUIT_CODE, check_signals, sig_setup, signals_pending};
use crate::state::{read_logic, source_rc, write_jobs, write_meta};
use clap::Parser;
use state::{read_vars, write_vars};

#[derive(Parser, Debug)]
struct ShedArgs {
  script: Option<String>,

  #[arg(short)]
  command: Option<String>,

  #[arg(trailing_var_arg = true)]
  script_args: Vec<String>,

  #[arg(long)]
  version: bool,

  #[arg(short)]
  interactive: bool,

  #[arg(long, short)]
  login_shell: bool,
}

/// Force evaluation of lazily-initialized values early in shell startup.
///
/// In particular, this ensures that the variable table is initialized, which
/// populates environment variables from the system. If this initialization is
/// deferred too long, features like prompt expansion may fail due to missing
/// environment variables.
///
/// This function triggers initialization by calling `read_vars` with a no-op
/// closure, which forces access to the variable table and causes its `LazyLock`
/// constructor to run.
fn kickstart_lazy_evals() {
  read_vars(|_| {});
}

/// We need to make sure that even if we panic, our child processes get sighup
fn setup_panic_handler() {
  let default_panic_hook = std::panic::take_hook();
  std::panic::set_hook(Box::new(move |info| {
    let _ = state::SHED.try_with(|shed| {
      if let Ok(mut jobs) = shed.jobs.try_borrow_mut() {
        jobs.hang_up();
      }
    });

    default_panic_hook(info);
  }));
}

fn main() -> ExitCode {
  yansi::enable();
  env_logger::init();
  kickstart_lazy_evals();
  setup_panic_handler();

  let mut args = ShedArgs::parse();
  if env::args().next().is_some_and(|a| a.starts_with('-')) {
    // first arg is '-shed'
    // meaning we are in a login shell
    args.login_shell = true;
  }
  if args.version {
    println!(
      "shed {} ({} {})",
      env!("CARGO_PKG_VERSION"),
      std::env::consts::ARCH,
      std::env::consts::OS
    );
    return ExitCode::SUCCESS;
  }

  if let Err(e) = if let Some(path) = args.script {
    run_script(path, args.script_args)
  } else if let Some(cmd) = args.command {
    exec_input(cmd, None, false, None)
  } else {
    shed_interactive()
  } {
    eprintln!("shed: {e}");
  };

  if let Some(trap) = read_logic(|l| l.get_trap(TrapTarget::Exit))
	&& let Err(e) = exec_input(trap, None, false, Some("trap".into())) {
    eprintln!("shed: error running EXIT trap: {e}");
  }

  write_jobs(|j| j.hang_up());
  ExitCode::from(QUIT_CODE.load(Ordering::SeqCst) as u8)
}

fn run_script<P: AsRef<Path>>(path: P, args: Vec<String>) -> ShResult<()> {
  let path = path.as_ref();
	let path_raw = path.to_string_lossy().to_string();
  if !path.is_file() {
    eprintln!("shed: Failed to open input file: {}", path.display());
    QUIT_CODE.store(1, Ordering::SeqCst);
    return Err(ShErr::simple(
      ShErrKind::CleanExit(1),
      "input file not found",
    ));
  }
  let Ok(input) = fs::read_to_string(path) else {
    eprintln!("shed: Failed to read input file: {}", path.display());
    QUIT_CODE.store(1, Ordering::SeqCst);
    return Err(ShErr::simple(
      ShErrKind::CleanExit(1),
      "failed to read input file",
    ));
  };

  write_vars(|v| {
    v.cur_scope_mut()
      .bpush_arg(path.to_string_lossy().to_string())
  });
  for arg in args {
    write_vars(|v| v.cur_scope_mut().bpush_arg(arg))
  }

  exec_input(input, None, false, Some(path_raw))
}

fn shed_interactive() -> ShResult<()> {
  let _raw_mode = raw_mode(); // sets raw mode, restores termios on drop
  sig_setup();

  if let Err(e) = source_rc() {
    e.print_error();
  }

  // Create readline instance with initial prompt
  let mut readline = match ShedVi::new(Prompt::new(), *TTY_FILENO) {
    Ok(rl) => rl,
    Err(e) => {
      eprintln!("Failed to initialize readline: {e}");
      QUIT_CODE.store(1, Ordering::SeqCst);
      return Err(ShErr::simple(
        ShErrKind::CleanExit(1),
        "readline initialization failed",
      ));
    }
  };

  // Main poll loop
  loop {
    write_meta(|m| {
      m.try_rehash_commands();
      m.try_rehash_cwd_listing();
    });
		error::clear_color();

    // Handle any pending signals
    while signals_pending() {
      if let Err(e) = check_signals() {
        match e.kind() {
          ShErrKind::ClearReadline => {
            // Ctrl+C - clear current input and redraw
            readline.reset_active_widget(false)?;
          }
          ShErrKind::CleanExit(code) => {
            QUIT_CODE.store(*code, Ordering::SeqCst);
            return Ok(());
          }
          _ => e.print_error(),
        }
      }
    }

    if GOT_SIGWINCH.swap(false, Ordering::SeqCst) {
      log::info!("Window size change detected, updating readline dimensions");
      // Restore cursor to saved row before clearing, since the terminal
      // may have moved it during resize/rewrap
      readline.writer.update_t_cols();
      readline.prompt_mut().refresh()?;
      readline.mark_dirty();
    }

    if JOB_DONE.swap(false, Ordering::SeqCst) {
      // update the prompt so any job count escape sequences update dynamically
      readline.prompt_mut().refresh()?;
    }

    readline.print_line(false)?;

    // Poll for stdin input
    let mut fds = [PollFd::new(
      unsafe { BorrowedFd::borrow_raw(*TTY_FILENO) },
      PollFlags::POLLIN,
    )];

    let timeout = if readline.pending_keymap.is_empty() {
      PollTimeout::MAX
    } else {
      PollTimeout::from(1000u16)
    };

    match poll(&mut fds, timeout) {
      Ok(_) => {}
      Err(Errno::EINTR) => {
        // Interrupted by signal, loop back to handle it
        continue;
      }
      Err(e) => {
        eprintln!("poll error: {e}");
        break;
      }
    }

    // Timeout — resolve pending keymap ambiguity
    if !readline.pending_keymap.is_empty()
      && fds[0].revents().is_none_or(|r| !r.contains(PollFlags::POLLIN))
    {
      log::debug!("[keymap timeout] resolving pending={:?}", readline.pending_keymap);
      let keymap_flags = readline.curr_keymap_flags();
      let matches = read_logic(|l| l.keymaps_filtered(keymap_flags, &readline.pending_keymap));
      // If there's an exact match, fire it; otherwise flush as normal keys
      let exact = matches.iter().find(|km| km.compare(&readline.pending_keymap) == KeyMapMatch::IsExact);
      if let Some(km) = exact {
        log::debug!("[keymap timeout] firing exact match: {:?} -> {:?}", km.keys, km.action);
        let action = km.action_expanded();
        readline.pending_keymap.clear();
        for key in action {
          if let Some(event) = readline.handle_key(key)? {
            match event {
              ReadlineEvent::Line(input) => {
                let start = Instant::now();
                write_meta(|m| m.start_timer());
                if let Err(e) = RawModeGuard::with_cooked_mode(|| exec_input(input, None, true, Some("<stdin>".into()))) {
                  match e.kind() {
                    ShErrKind::CleanExit(code) => {
                      QUIT_CODE.store(*code, Ordering::SeqCst);
                      return Ok(());
                    }
                    _ => e.print_error(),
                  }
                }
                let command_run_time = start.elapsed();
                log::info!("Command executed in {:.2?}", command_run_time);
                write_meta(|m| m.stop_timer());
                readline.fix_column()?;
                readline.writer.flush_write("\n\r")?;
                readline.reset(true)?;
                break;
              }
              ReadlineEvent::Eof => {
                QUIT_CODE.store(0, Ordering::SeqCst);
                return Ok(());
              }
              ReadlineEvent::Pending => {}
            }
          }
        }
      } else {
        log::debug!("[keymap timeout] no exact match, flushing {} keys as normal input", readline.pending_keymap.len());
        let buffered = std::mem::take(&mut readline.pending_keymap);
        for key in buffered {
          if let Some(event) = readline.handle_key(key)? {
            match event {
              ReadlineEvent::Line(input) => {
                let start = Instant::now();
                write_meta(|m| m.start_timer());
                if let Err(e) = RawModeGuard::with_cooked_mode(|| exec_input(input, None, true, Some("<stdin>".into()))) {
                  match e.kind() {
                    ShErrKind::CleanExit(code) => {
                      QUIT_CODE.store(*code, Ordering::SeqCst);
                      return Ok(());
                    }
                    _ => e.print_error(),
                  }
                }
                let command_run_time = start.elapsed();
                log::info!("Command executed in {:.2?}", command_run_time);
                write_meta(|m| m.stop_timer());
                readline.fix_column()?;
                readline.writer.flush_write("\n\r")?;
                readline.reset(true)?;
                break;
              }
              ReadlineEvent::Eof => {
                QUIT_CODE.store(0, Ordering::SeqCst);
                return Ok(());
              }
              ReadlineEvent::Pending => {}
            }
          }
        }
      }
      readline.print_line(false)?;
      continue;
    }

    // Check if stdin has data
    if fds[0]
      .revents()
      .is_some_and(|r| r.contains(PollFlags::POLLIN))
    {
      let mut buffer = [0u8; 1024];
      match read(*TTY_FILENO, &mut buffer) {
        Ok(0) => {
          // EOF
          break;
        }
        Ok(n) => {
          readline.feed_bytes(&buffer[..n]);
        }
        Err(Errno::EINTR) => {
          // Interrupted, continue to handle signals
          continue;
        }
        Err(e) => {
          eprintln!("read error: {e}");
          break;
        }
      }
    }

    // Process any available input
    match readline.process_input() {
      Ok(ReadlineEvent::Line(input)) => {
        let start = Instant::now();
        write_meta(|m| m.start_timer());
        if let Err(e) = RawModeGuard::with_cooked_mode(|| exec_input(input, None, true, Some("<stdin>".into()))) {
          match e.kind() {
            ShErrKind::CleanExit(code) => {
              QUIT_CODE.store(*code, Ordering::SeqCst);
              return Ok(());
            }
            _ => e.print_error(),
          }
        }
        let command_run_time = start.elapsed();
        log::info!("Command executed in {:.2?}", command_run_time);
        write_meta(|m| m.stop_timer());
				readline.fix_column()?;
				readline.writer.flush_write("\n\r")?;

        // Reset for next command with fresh prompt
        readline.reset(true)?;
        let real_end = start.elapsed();
        log::info!("Total round trip time: {:.2?}", real_end);
      }
      Ok(ReadlineEvent::Eof) => {
        // Ctrl+D on empty line
        QUIT_CODE.store(0, Ordering::SeqCst);
        return Ok(());
      }
      Ok(ReadlineEvent::Pending) => {
        // No complete input yet, keep polling
      }
      Err(e) => match e.kind() {
        ShErrKind::CleanExit(code) => {
          QUIT_CODE.store(*code, Ordering::SeqCst);
          return Ok(());
        }
        _ => e.print_error(),
      },
    }
  }

  Ok(())
}