path: root/src/run/mod.rs

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use crate::parse::{Ast, PostExpansion, PreExpansion};
use crate::*;

mod builtin;

enum ExecError {
    UnknownVariable(BString),
    ExecError(i32),
}

struct Executor {
    se: Arc<Mutex<Session>>,
}

#[derive(Clone)]
struct ArcWriter {
    inner: Arc<Mutex<Vec<u8>>>,
}

impl ArcWriter {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(Vec::new())),
        }
    }
    pub fn into_inner(self) -> Vec<u8> {
        self.inner.lock().unwrap().clone()
    }
}

impl std::io::Write for ArcWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.lock().unwrap().extend_from_slice(buf);
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Executor {
    fn execute_pipeline(
        &mut self,
        pipes: parse::Pipes<parse::PostExpansion>,
        capture: Option<&mut Vec<u8>>,
    ) -> Result<(), ExecError> {
        let mut children = Vec::new();
        let mut threads = Vec::new();
        let mut prev_reader = None;
        let mut spawn_error = false;

        let last_output = ArcWriter::new();
        let mut last_is_command = false;

        let pipelen = pipes.cmds.len();
        for (i, cmd) in pipes.cmds.into_iter().enumerate() {
            let (reader, writer) = if i < pipelen - 1 {
                let (r, w) = io::pipe().unwrap();
                (Some(r), Some(w))
            } else {
                (None, None)
            };

            let dc = self.se.lock().unwrap().dispatch.get(&cmd.cmd[..]);

            match dc {
                CommandKind::Path(path) => {
                    last_is_command = true;

                    let mut command = Command::new(&path);
                    for arg in cmd.args.iter() {
                        command.arg(OsStr::from_bytes(arg));
                    }

                    if let Some(r) = prev_reader.take() {
                        command.stdin(Stdio::from(r));
                    } else if capture.is_some() {
                        command.stdin(Stdio::null());
                    }

                    if let Some(w) = writer {
                        command.stdout(Stdio::from(w));
                    } else if capture.is_some() {
                        command.stdout(Stdio::piped());
                    }

                    let child = match command.spawn() {
                        Ok(c) => c,
                        Err(e) => {
                            let cmd = path.to_string_lossy();
                            let msg = match e.kind() {
                                io::ErrorKind::NotFound => format!("{cmd} not found"),
                                io::ErrorKind::PermissionDenied => {
                                    format!("{cmd} is not executable")
                                }
                                io::ErrorKind::FileTooLarge => format!("{cmd} is too massive"),
                                io::ErrorKind::ResourceBusy | io::ErrorKind::ExecutableFileBusy => {
                                    format!("{cmd} is busy")
                                }
                                io::ErrorKind::TooManyLinks => {
                                    format!("{cmd} could not be resolved")
                                }
                                io::ErrorKind::InvalidFilename => {
                                    format!("{cmd} is not a valid file name")
                                }
                                io::ErrorKind::ArgumentListTooLong => format!("too many arguments"),
                                io::ErrorKind::Interrupted => format!("got interrupted"),
                                io::ErrorKind::Unsupported => format!("{cmd} is not supported"),
                                e => format!("I am surprised you can get this error here: {e:?}"),
                            };
                            println!("pish: {msg}");
                            spawn_error = true;
                            break;
                        }
                    };

                    children.push(child);
                }

                CommandKind::Builtin(builtin) => {
                    last_is_command = false;

                    let mut input: Box<dyn io::Read + Send> = match prev_reader.take() {
                        Some(r) => Box::new(r),
                        None if capture.is_some() => Box::new(io::empty()),
                        None => Box::new(io::stdin()),
                    };

                    let mut output: Box<dyn io::Write + Send> = match writer {
                        Some(w) => Box::new(w),
                        None if capture.is_some() => Box::new(last_output.clone()),
                        None => Box::new(io::stdout()),
                    };

                    // currently only required for `re`, cannot happen in background thread
                    builtin.special(self.se.clone(), &cmd.args);

                    let cloned_session = self.se.clone();
                    let handle = std::thread::spawn(move || {
                        builtin.io(cloned_session, &cmd.args, &mut input, &mut output)
                    });

                    threads.push(handle);
                }
            }

            prev_reader = reader;
        }

        if spawn_error {
            for child in children.iter_mut() {
                if let Err(e) = child.kill() {
                    println!("failed to kill child - {e:?}");
                }
            }
            Err(ExecError::ExecError(127))
        } else {
            let mut code = 0;
            for jh in threads {
                match jh.join() {
                    Ok(Ok(())) => (),
                    Ok(Err(e)) => match e {
                        BuiltinError::IO(_) => code = -1,
                        BuiltinError::Exit(c) => code = c,
                    },
                    Err(_) => code = 127,
                }
            }
            for child in children.iter_mut() {
                match child.wait() {
                    Ok(ec) => {
                        if let Some(c) = ec.code() {
                            code = c;
                        }
                    }
                    Err(e) => {
                        println!("failed to wait for child - {e:?}")
                    }
                }
            }

            if let Some(cap) = capture {
                if last_is_command {
                    let child = children.into_iter().last().unwrap();
                    let out = child.wait_with_output().unwrap();
                    *cap = out.stdout;
                } else {
                    *cap = last_output.into_inner();
                }
            }

            if code == 0 {
                Ok(())
            } else {
                Err(ExecError::ExecError(code))
            }
        }
    }

    fn execute_var_assign(&mut self, va: parse::VarAssign<PostExpansion>) -> Result<(), ExecError> {
        self.se.lock().unwrap().vars.insert(va.var, va.val);
        Ok(())
    }

    fn execute_fun_decl(&mut self, fd: parse::FunDecl<PostExpansion>) -> Result<(), ExecError> {
        self.se.lock().unwrap().funs.insert(fd.name, *fd.body.body);
        Ok(())
    }

    fn execute(
        &mut self,
        ast: Ast<parse::PostExpansion>,
        capture: Option<&mut Vec<u8>>,
    ) -> Result<(), ExecError> {
        match ast {
            Ast::FunDecl(fd) => self.execute_fun_decl(fd),
            Ast::VarAssign(va) => self.execute_var_assign(va),
            Ast::Pipes(pipes) => self.execute_pipeline(pipes, capture),
        }
    }
}

impl parse::Expander for Executor {
    type Error = ExecError;

    fn expand_var(&mut self, var: BString) -> Result<BString, Self::Error> {
        if let Some(val) = self.se.lock().unwrap().vars.get(&var) {
            return Ok(val.clone())
        }

        match std::env::var_os(OsStr::from_bytes(&var)) {
            Some(val) => Ok(val.as_bytes().to_vec()),
            None => Err(ExecError::UnknownVariable(var)),
        }
    }

    fn expand_cmd(&mut self, ast: Ast<parse::PostExpansion>) -> Result<BString, Self::Error> {
        let mut out = Vec::new();
        self.execute(ast, Some(&mut out))?;
        if out.last() == Some(&b'\n') {
            out.pop();
        }
        Ok(out)
    }
}

fn exec(se: Arc<Mutex<Session>>, ast: Ast<PreExpansion>) -> Result<(), ExecError> {
    let mut exec = Executor { se };
    let ast = ast.expand(&mut exec)?;
    exec.execute(ast, None)
}

pub fn run(se: Arc<Mutex<Session>>, cmd: Vec<u8>) {
    let parsed = parse::do_parse(&cmd);
    let parsed = match parsed {
        Ok(p) => p,
        Err(err) => {
            se.lock().unwrap().raw.disable();
            println!("{:?}: {}", err.0, String::from_utf8_lossy(&err.1));
            print!("{}", se.lock().unwrap().prompt());
            std::io::stdout().lock().flush().unwrap();
            se.lock().unwrap().raw.enable();
            return;
        }
    };

    se.lock().unwrap().raw.disable();
    let result = exec(se.clone(), parsed);
    se.lock().unwrap().raw.enable();

    let status_string = match result {
        Ok(_) => String::new(),
        Err(ExecError::UnknownVariable(var)) => {
            format!("unbound variable: {}", String::from_utf8_lossy(&var))
        }
        Err(ExecError::ExecError(i)) => i.to_string(),
    };

    print!("\r{status_string}\r\n{}", se.lock().unwrap().prompt());
    let _ = std::io::stdout().lock().flush();
}

#[derive(Debug)]
pub enum BuiltinError {
    IO(std::io::Error),
    Exit(i32),
}

impl From<std::io::Error> for BuiltinError {
    fn from(value: std::io::Error) -> Self {
        Self::IO(value)
    }
}

type BuiltinResult = Result<(), BuiltinError>;

#[allow(unused_variables)]
pub trait Builtin: Send + Sync {
    fn name(&self) -> &str;

    fn special(&self, session: Arc<Mutex<Session>>, args: &[BString]) {}

    fn io(
        &self,
        session: Arc<Mutex<Session>>,
        args: &[BString],
        stdin: &mut dyn Read,
        stdout: &mut dyn Write,
    ) -> BuiltinResult;
}

const BUILTINS: &[&'static dyn Builtin] = &[
    &builtin::cd,
    &builtin::clear,
    #[cfg(debug_assertions)]
    &builtin::re,
    &builtin::sink("to", false),
    &builtin::sink("into", false),
    &builtin::sink("append", true),
    &builtin::from,
    &builtin::builtins,
    &builtin::_type,
    &builtin::history,
    &builtin::escape,
    &builtin::parse,
];

pub struct CommandDispatch {
    map: HashMap<BString, CommandKind>,
}

impl CommandDispatch {
    pub fn new() -> Self {
        let mut map = HashMap::new();

        // builtins
        for &b in BUILTINS {
            map.insert(b.name().as_bytes().to_vec(), CommandKind::Builtin(b));
        }

        Self { map }
    }

    fn get(&self, cmd: &bstr) -> CommandKind {
        let path_cmd = CommandKind::Path(PathBuf::from(OsStr::from_bytes(cmd)));
        if cmd.contains(&b'/') {
            path_cmd
        } else if let Some(cmd) = self.map.get(cmd) {
            cmd.clone()
        } else {
            path_cmd
        }
    }
}

#[derive(Clone)]
pub enum CommandKind {
    Builtin(&'static dyn Builtin),
    Path(PathBuf),
}