path: root/src/recite.rs

pub mod path;
mod ps;
use crate::{ctask, task};
use core::fmt;
use libc::{waitpid, WNOHANG};
use path::prefresh;
use std::io::{self, Write};
use std::path::Path;
use std::process::{exit, Command, Stdio};
use std::sync::{Arc, Mutex};

/// Describes the ending of a [Verse]
///
/// The ending of a verse determines how the [Stanza] should be interpreted.
/// For instance, a [Stanza] that is piped needs to have it's `STDOUT`
/// captured (rather than printing out to the terminal), and subsequently sent
/// to the next [Verse] in the [Poem].
///
/// # Values
/// * `None` - A shell command with no additional actions
/// * `Couplet` - Pipe the output of this command into the next (`|`)
/// * `Quiet` - Fork the called process into the background (`&`)
/// * `And` - Run the next command only if this one succeeds (`&&`)
/// * `String` - String commands together on a single line (`;`)
#[derive(Debug, PartialEq, Eq)]
enum Meter {
    None,    // No meter
    Couplet, // Pipe the output of this command into the next
    Quiet,   // Fork the command into the background
    And,     // Run the next command only if this succeeds
    String,  // Run the next command, even if this doesn't succeed
}

impl fmt::Display for Meter {
    /// Determine how to print out a [Meter]
    ///
    /// Each [meter's][Meter] symbol corresponds to it's input.
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let meter = match self {
            Meter::None => "",
            Meter::Couplet => "|",
            Meter::Quiet => "&",
            Meter::And => "&&",
            Meter::String => ";",
        };

        write!(f, "{}", meter)
    }
}

impl Meter {
    /// Recite a verse with [Meter::None]
    ///
    /// Call this function on a [Verse] with a meter of type [Meter::None].
    /// This forks into a child process, calls the `verb` (i.e. program)
    /// that was specified in the [Verse], then waits for that program to
    /// complete. If the last [Verse] piped its contents into `out`, it will
    /// be piped into the STDIN of this [Verse]. If all Rust code is called
    /// successfully, return the exit code of the process. Otherwise, return a
    /// [std::io::Error].
    ///
    /// # Arguments
    /// * `verse: &Verse` - The verse to recite
    /// * `out: &mut String` - A string that may have output from the last command
    fn incant_none(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
        let child = task!(verse, out);

        let output = child.wait_with_output()?;

        if !output.status.success() {
            return Ok(output.status.code().unwrap_or(-1));
        }

        Ok(output.status.code().unwrap_or(0))
    }

    /// Recite a verse with [Meter::None]
    ///
    /// Call this function on a [Verse] with a meter of type [Meter::None].
    /// This forks into a child process, calls the `verb` (i.e. program)
    /// that was specified in the [Verse], then waits for that program to
    /// complete. If the last [Verse] piped its contents into `out`, it will
    /// be piped into the STDIN of this [Verse]. Then, the contents of this
    /// processes' STDOUT are stored in `out`. If all Rust code is called
    /// successfully, return the exit code of the process. Otherwise, return a
    /// [std::io::Error].
    ///
    /// # Arguments
    /// * `verse: &Verse` - The verse to recite
    /// * `out: &mut String` - A string that may have output from the last command
    fn incant_couplet(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
        let child = ctask!(verse, out);

        let output = child.wait_with_output()?;

        if !output.status.success() {
            return Ok(output.status.code().unwrap_or(-1));
        }

        out.push_str(
            String::from_utf8_lossy(&output.stdout)
                .into_owned()
                .as_str(),
        );

        Ok(output.status.code().unwrap_or(0))
    }

    /// Recite a verse with [Meter::Quiet]
    ///
    /// Call this function on a [Verse] with a meter of type [Meter::Quiet].
    /// This forks a child process into the background. It then registers a
    /// `SIGCHLD` handler, making sure to do so for each PID in the `pids`
    /// Vec. If the last [Verse] piped its contents into `out`, it will be
    /// piped into the STDIN of this [Verse]. If all Rust code is called
    /// successfully, return the exit code of the process. Otherwise, return a
    /// [std::io::Error].
    ///
    /// # Arguments
    /// * `verse: &Verse` - The verse to recite
    /// * `out: &mut String` - A string that may have output from the last command
    /// * `pids: Arc<Mutex<Vec<i32>>>` - A vector that stores the PIDs of all background processes that belong to the shell
    fn incant_quiet(
        verse: &Verse,
        out: &mut String,
        pids: &mut Arc<Mutex<Vec<i32>>>,
    ) -> Result<i32, io::Error> {
        let child = task!(verse, out);
        println!("[&]  {}", child.id());

        pids.lock().unwrap().push(child.id() as i32);
        let stanza = verse.stanza.to_string();
        let pids = Arc::clone(pids);

        unsafe {
            signal_hook::low_level::register(signal_hook::consts::SIGCHLD, move || {
                for pid in pids.lock().unwrap().iter() {
                    let mut pid = *pid;
                    let mut status: i32 = 0;
                    pid = waitpid(pid, &mut status, WNOHANG);
                    if pid > 0 {
                        print!("\n[&]  + done    {}", stanza);
                        io::stdout().flush().unwrap();
                    }
                }
            })
            .unwrap();
        }

        Ok(0)
    }

    /// Alias to [Meter::incant_none]
    fn incant_and(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
        Meter::incant_none(verse, out)
    }

    /// Alias to [Meter::incant_none]
    fn incant_string(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
        Meter::incant_none(verse, out)
    }
}

/// Holds a program to be called
///
/// This is simply the first word in a full command [String], dilineated via
/// whitespace.
type Verb = String;

/// Holds arguments to a program
///
/// This is a list of all the words that come after the [Verb], dilineated via
/// whitespace.
type Clause = Vec<String>;

/// Holds the interpreted elements of a [Verse]
///
/// Each [Stanza] has two parts, a [Verb] and a [Clause]. The [Verb] is the
/// program, or path to the program to call, while the [Clause] contains
/// arguments to pass to that program.
#[derive(Debug)]
struct Stanza {
    verb: Verb,
    clause: Clause,
}

impl fmt::Display for Stanza {
    /// Print out a [Stanza]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {}", self.verb, self.clause.join(" "))
    }
}

impl Stanza {
    /// Create a new [Stanza]
    ///
    /// Returns a new [Stanza] built from a `Vec<String>`. The first element of
    /// the vector becomes the [Verb], while the remainder of the vector
    /// becomes the [Clause].
    ///
    /// # Arguments
    /// `stanza: Vec<String>` - The full command split into individual strings
    ///                         via whitespace
    ///
    /// # Examples
    /// ```
    /// // Input: cargo build --release
    /// let command = vec!["cargo", "build", "--release"]
    ///               .into_iter()
    ///               .map(String::from)
    ///               .collect<Vec<String>>();
    /// let stanza = Stanza::new(command);
    /// println!("{}", stanza.verb);
    /// println!("{:?}", stanza.clause);
    ///
    /// ```
    fn new(stanza: Vec<String>) -> Stanza {
        Stanza {
            verb: stanza[0].to_owned(),
            clause: stanza[1..].to_vec(),
        }
    }

    /// Check if the [Verb] exists in the `$PATH`
    ///
    /// First checks if the [Verb] is a relative or full path. If it is, check
    /// whether or not it exists. If it does exist, return true, otherwise see
    /// if the [Verb] is cached in our list of binaries. Search is done in
    /// $PATH order.
    ///
    /// # Examples
    /// ```
    /// let bins = vec!["cargo", "ruby", "cat"]
    ///            .into_iter()
    ///            .map(String::from)
    ///            .collect<Vec<String>>();
    ///
    /// let command_success = vec!["cargo", "build", "--release"]
    ///                       .into_iter()
    ///                       .map(String::from)
    ///                       .collect<Vec<String>>();
    ///
    /// let command_fail = vec!["make", "-j8"]
    ///                    .into_iter()
    ///                    .map(String::from)
    ///                    .collect<Vec<String>>();
    ///
    /// let stanza_success = Stanza::new(command_success);
    /// let stanza_fail = Stanza::new(command_fail);
    ///
    /// stanza_success.spellcheck(bins) // -> true
    /// stanza_fail.spellcheck(bins) // -> false
    /// ```
    fn spellcheck(&self, bins: &Vec<String>) -> bool {
        // An empty verb (i.e. the empty string) cannot be a program, so
        // return false
        // Thanks to the parsing in Poem::read, however, it's
        // unlikely for this to happen
        if self.verb.is_empty() {
            return false;
        }

        // Only search the $PATH if a full or relative path was not given, or
        // if the path given does not exist
        if !Path::new(self.verb.as_str()).exists() {
            // Try to find a binary in our path with the same name as the verb
            // Searches in $PATH order
            match bins
                .iter()
                .find(|bin| bin.split('/').last().unwrap() == self.verb)
            {
                Some(_) => return true,
                None => return false,
            }
        }

        // Return true if the full path or relative path exists
        true
    }
}

/// Holds a [Stanza] and its [Meter]
///
/// In addition to a [Stanza] and a [Meter], [verse's][Verse] also hold a bool
/// value called `couplet`, indicating that it needs to accept input on `STDIN`
/// from the previous [Verse].
#[derive(Debug)]
struct Verse {
    stanza: Stanza,
    meter: Meter,
    couplet: bool,
}

impl fmt::Display for Verse {
    /// Print out a [Verse]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{} {} {}",
            self.verb(),
            self.clause().join(" "),
            self.meter
        )
    }
}

impl Verse {
    /// Create a new [Verse]
    ///
    /// Returns a new [Verse] built from a [Stanza], a [Meter], and a `couplet`
    /// indicator. See [Poem::read] for more details on how these are
    /// constructed.
    fn new(stanza: Stanza, meter: Meter, couplet: bool) -> Verse {
        Verse {
            stanza,
            meter,
            couplet,
        }
    }

    /// Alias to [Stanza::spellcheck]
    fn spellcheck(&self, bins: &Vec<String>) -> bool {
        self.stanza.spellcheck(bins)
    }

    /// Alias to [stanza's][Stanza] `verb`
    fn verb(&self) -> String {
        self.stanza.verb.clone()
    }

    /// Alias to [stanza's][Stanza] `clause`
    fn clause(&self) -> Vec<String> {
        self.stanza.clause.clone()
    }

    /// Check if this verse is piping output
    fn couplet(verse: Option<&Verse>) -> bool {
        match verse {
            Some(verse) => match verse.meter {
                Meter::Couplet | Meter::Quiet | Meter::And | Meter::String => true,
                Meter::None => false,
            },
            None => false,
        }
    }

    /// Check if this verse has a meter
    fn cadence(verse: Option<&Verse>) -> bool {
        match verse {
            Some(verse) => match verse.meter {
                Meter::Couplet | Meter::Quiet | Meter::And | Meter::String => true,
                Meter::None => false,
            },
            None => false,
        }
    }
}

/// An entire shell command parsed into [verse's][Verse]
///
/// A [Poem] is the structure that contains a full shell command/program. It
/// may be composed of one or many [verse's][Verse].
#[derive(Debug)]
pub struct Poem {
    verses: Vec<Verse>,
}

impl Poem {
    /// Create a new [Poem]
    ///
    /// Returns a new [Poem] built from a list of [verse's][Verse].
    fn new(verses: Vec<Verse>) -> Poem {
        Poem { verses }
    }

    /// Recite a [Poem] (run the shell command(s)/program)
    ///
    /// This function attempts to call each [Verse] in the [Poem], in the order
    /// that it was inputted/parsed.
    ///
    /// # Arguments
    /// * `path` - A list of directories from the $PATH environment variable
    ///            Needed in case we need to refresh the $PATH
    /// * `bins` - A list of binaries cached from the $PATH, used for searching
    ///            for a program that matches the verb in each [Verse]
    ///
    /// # Returns
    /// * `true` - If the entire [Poem] was recited without fault
    /// * `false` - If any [Verse] of the [Poem] was invalid
    ///
    /// # Examples
    /// ```
    /// let poetry = "ps aux | grep dwvsh".to_string();
    /// let poem = Poem::read(poetry);
    ///
    /// match poem {
    ///     Some(poem) => { poem.recite(path, &mut bins); }
    ///     None => {}
    /// }
    /// ```
    pub fn recite(&self, path: &Vec<&Path>, bins: &mut Vec<String>) -> Result<(), io::Error> {
        // Variable for storing the output of a piped verse
        let mut out: String = String::new();
        let mut pids: Arc<Mutex<Vec<i32>>> = Arc::new(Mutex::new(Vec::new()));

        // Loop through each verse in the poem
        for verse in self.verses.iter() {
            // Check if user wants to exit the shell
            if verse.verb() == "exit" || verse.verb() == "quit" {
                exit(0);
            }

            // Check if the user wants to change directories
            if verse.verb() == "cd" {
                let path: String;
                if verse.clause().is_empty() {
                    path = env!("HOME").to_string();
                } else {
                    path = verse.clause().first().unwrap().to_owned();
                }

                match std::env::set_current_dir(&path) {
                    Ok(_) => continue,
                    Err(_) => {
                        println!("cd: unable to change into {}", path);
                        continue;
                    }
                }
            }

            // Check if the verb exists
            // If it doesn't exist, try refreshing the binary cache, and check
            // again
            // If it still doesn't exist, print an error
            if !verse.spellcheck(bins) {
                *bins = prefresh(path);
                if !verse.spellcheck(bins) {
                    println!("dwvsh: {}: command not found", verse.verb());
                    continue;
                }
            }

            // Incant the verse, based on its meter
            let status = match verse.meter {
                Meter::None => Meter::incant_none(verse, &mut out)?,
                Meter::Couplet => Meter::incant_couplet(verse, &mut out)?,
                Meter::Quiet => Meter::incant_quiet(verse, &mut out, &mut pids)?,
                Meter::And => Meter::incant_and(verse, &mut out)?,
                Meter::String => Meter::incant_string(verse, &mut out)?,
            };

            // Don't continue reciting if there was an error, unless the meter
            // is String (indicating that errors should be ignored)
            if verse.meter != Meter::String && status != 0 {
                break;
            }
        }

        // If we've successfully exited the loop, then all verse's were
        // properly recited
        Ok(())
    }

    /// Parse a [Poem] from a raw [String] input
    ///
    /// Takes a shell command/program and converts it to a machine-runnable
    /// [Poem]. If there is a parse error, [Poem::read] may [Option]ally return
    /// `None`. As of now, there is no support for multiline programs, unless
    /// newlines (`\n`) were to be swapped out for semicolons (`;`) before
    /// calling this function. See [Poem::recite] for how each [Verse] in a
    /// [Poem] is called.
    ///
    /// # Examples
    /// ```
    /// let poetry = "ps aux | grep dwvsh".to_string();
    /// let poem = Poem::read(poetry);
    /// ```
    pub fn read(poetry: String) -> Option<Poem> {
        // Need to loop through each char in the input string, since some
        // characters aren't whitespace dilineated (`;`, `&`, etc.)
        //
        // Need to keep track of the previous verse, since it might haver
        // a Meter of Couplet, meaning that we need to set couplet on the
        // current verse
        let mut chars = poetry.chars();
        let mut verses: Vec<Verse> = Vec::new(); // Accumulate verses
        let mut stanza: Vec<String> = Vec::new(); // Stack for each stanza
        let mut word: Vec<char> = Vec::new(); // Stack for each word
        let mut prev: Option<&Verse> = None; // The previous verse
        let mut i: usize = 0; // Keep track of our index into chars

        // Parse from left to right
        loop {
            // Get the next character in the input string
            let char = chars.next();

            // Do something depending on what the character is
            match char {
                // Print an error, and return None if a Meter was used without
                // a Stanza before it
                Some(meter)
                    if ((meter == '|' || meter == '&')
                        && Verse::cadence(prev)
                        && stanza.is_empty())
                        || ((meter == '|' || meter == '&') && i == 0) =>
                {
                    eprintln!(
                        "dwvsh: parse error: verse must have a stanza: rune {} at column {}",
                        meter, i
                    );
                    return None;
                }

                // The character represents the Couplet Meter
                Some(meter) if meter == '|' => {
                    // If there are chars on the word stack, push that word
                    // onto the stanza
                    if !word.is_empty() {
                        stanza.push(word.iter().collect());
                    }

                    // A meter indicates the end of a verse
                    verses.push(Verse::new(
                        Stanza::new(stanza.clone()),
                        Meter::Couplet,
                        Verse::couplet(prev),
                    ));

                    // Clear the stacks
                    stanza.clear();
                    word.clear();
                }

                // The character represents the Quiet (or And) Meter
                Some(meter) if meter == '&' => {
                    // If there are chars on the word stack, push that word
                    // onto the stanza
                    if !word.is_empty() {
                        stanza.push(word.iter().collect());
                    }

                    // Need to peek at the next character, since '&' can mean
                    // Meter::Quiet, or '&&' can mean Meter::And
                    match chars.clone().peekable().peek() {
                        // Indicated Meter::And
                        Some(c) if c == &'&' => {
                            // Pop the next character from the input string
                            // (since we know what it is)
                            chars.next();

                            // A meter indicates the end of a verse
                            verses.push(Verse::new(
                                Stanza::new(stanza.clone()),
                                Meter::And,
                                Verse::couplet(prev),
                            ));
                        }

                        // Indicates Meter::Quiet
                        Some(_) => {
                            // A meter indicates the end of a verse
                            verses.push(Verse::new(
                                Stanza::new(stanza.clone()),
                                Meter::Quiet,
                                Verse::couplet(prev),
                            ));
                        }

                        // Indicated the end of the input
                        None => {
                            // The end of input also indicates the end of a
                            // verse
                            verses.push(Verse::new(
                                Stanza::new(stanza.clone()),
                                Meter::Quiet,
                                Verse::couplet(prev),
                            ));

                            // We can break out of the loop here, since it's
                            // the end of the raw input
                            break;
                        }
                    }

                    // Clear the stacks
                    stanza.clear();
                    word.clear();
                }

                // The character represents the String Meter
                Some(meter) if meter == ';' => {
                    // If there are chars on the word stack, push that word
                    // onto the stanza
                    if !word.is_empty() {
                        stanza.push(word.iter().collect());
                    }

                    // A meter indicates the end of a verse
                    if !stanza.is_empty() {
                        verses.push(Verse::new(
                            Stanza::new(stanza.clone()),
                            Meter::String,
                            Verse::couplet(prev),
                        ));
                    }

                    // Clear the stacks
                    stanza.clear();
                    word.clear();
                }

                // The character is whitespace
                Some(char) if char == ' ' || char == '\t' => {
                    // If there are chars on the word stack, push that word
                    // onto the stanza
                    if !word.is_empty() {
                        stanza.push(word.iter().collect());
                        word.clear();
                    }
                }

                // The character is any other utf8 glyph
                Some(char) => {
                    // Add the character onto the current word stack
                    word.push(char);
                }

                // Indicates the end of the list of characters
                None => {
                    // Always push the last word onto the stanza
                    if !word.is_empty() {
                        stanza.push(word.iter().collect());
                    }

                    // Only push the stanza into a verse if it contains
                    // any words
                    if !stanza.is_empty() {
                        verses.push(Verse::new(
                            Stanza::new(stanza.clone()),
                            Meter::None,
                            Verse::couplet(prev),
                        ));
                    }

                    // Break from the loop, since we are out of chars
                    break;
                }
            }

            // Set previous verse to the verse that was just pushed at the end
            // of each loop
            prev = match verses.last() {
                Some(verse) => Some(verse),
                None => None,
            };

            // Increment the index
            i += 1;
        }

        // Return the (parsed) poem
        Some(Poem::new(verses))
    }
}