yggdrasil/userspace/shell/src/env.rs

use std::{
    collections::{hash_map::Entry, HashMap},
    fmt,
    path::PathBuf,
};

use crate::syntax::{
    lex::{Fragment, Word},
    parse::{ICommand, IPipelineElement, IRedirects},
};

pub trait Expand {
    type Output;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error>;
}

#[derive(Clone)]
pub enum Variable {
    String(String),
    Array(Vec<Variable>),
}

pub struct Env {
    vars: HashMap<String, Variable>,
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Undefined variable {0:?}")]
    UndefinedVariable(String),
}

impl Env {
    pub fn new() -> Self {
        Self {
            vars: HashMap::new(),
        }
    }

    pub fn setup_builtin(&mut self) {
        let bin_name = std::env::args().next().unwrap_or_default();
        self.vars.insert("SHELL".into(), Variable::String(bin_name));
    }

    pub fn lookup(&self, name: &str) -> Option<Variable> {
        self.vars
            .get(name)
            .cloned()
            .or_else(|| std::env::var(name).ok().map(Variable::String))
    }

    pub fn put_var(&mut self, name: &str, value: Variable) -> bool {
        match self.vars.entry(name.into()) {
            Entry::Vacant(entry) => {
                entry.insert(value);
                true
            }
            Entry::Occupied(_) => false,
        }
    }
}

impl fmt::Display for Variable {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::String(text) => f.write_str(text),
            Self::Array(items) => {
                for (i, item) in items.iter().enumerate() {
                    if i != 0 {
                        f.write_str(" ")?;
                    }
                    fmt::Display::fmt(item, f)?;
                }
                Ok(())
            }
        }
    }
}

impl From<&str> for Variable {
    fn from(value: &str) -> Self {
        Self::String(value.into())
    }
}
impl From<String> for Variable {
    fn from(value: String) -> Self {
        Self::String(value)
    }
}

impl<T: Expand> Expand for Option<T> {
    type Output = Option<T::Output>;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        self.as_ref().map(|e| e.expand(env)).transpose()
    }
}

impl<T: Expand> Expand for Vec<T> {
    type Output = Vec<T::Output>;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        self.iter().map(|e| e.expand(env)).collect()
    }
}

impl<T: Expand, U: Expand> Expand for (T, U) {
    type Output = (T::Output, U::Output);

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        Ok((self.0.expand(env)?, self.1.expand(env)?))
    }
}

pub struct Command {
    pub pipeline: Vec<PipelineElement>,
    pub redirects: Redirects,
}

pub struct PipelineElement {
    pub envs: Vec<(String, Option<String>)>,
    pub words: Vec<String>,
}
pub struct Redirects {
    pub stdin: Option<PathBuf>,
    pub stdout: Option<PathBuf>,
    pub stderr: Option<PathBuf>,
}

impl Expand for ICommand<'_> {
    type Output = Command;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        let pipeline = self.pipeline.expand(env)?;
        let redirects = self.redirects.expand(env)?;
        Ok(Command {
            pipeline,
            redirects,
        })
    }
}

impl Expand for IRedirects<'_> {
    type Output = Redirects;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        let stdin = self.stdin.expand(env)?.map(PathBuf::from);
        let stdout = self.stdout.expand(env)?.map(PathBuf::from);
        let stderr = self.stderr.expand(env)?.map(PathBuf::from);
        Ok(Redirects {
            stdin,
            stdout,
            stderr,
        })
    }
}

impl Expand for IPipelineElement<'_> {
    type Output = PipelineElement;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        let words = self.words.expand(env)?;
        let envs = self.envs.expand(env)?;
        Ok(PipelineElement { envs, words })
    }
}

impl Expand for Word<'_> {
    type Output = String;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        let mut output = String::new();
        for fragment in self.0.iter() {
            output.push_str(&fragment.expand(env)?);
        }
        Ok(output)
    }
}

impl Expand for Fragment<'_> {
    type Output = String;

    fn expand(&self, env: &Env) -> Result<Self::Output, Error> {
        match *self {
            Self::Literal(text) => Ok(text.into()),
            Self::Variable(name) => {
                let text = env.lookup(name).map(|v| v.to_string()).unwrap_or_default();
                Ok(text)
            }
            Self::QuotedLiteral(text) => Ok(text.into()),
        }
    }
}

impl fmt::Display for Command {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, element) in self.pipeline.iter().enumerate() {
            if i != 0 {
                f.write_str(" | ")?;
            }
            fmt::Display::fmt(element, f)?;
        }
        if !self.pipeline.is_empty() {
            f.write_str(" ")?;
        }
        fmt::Display::fmt(&self.redirects, f)
    }
}

impl fmt::Display for Redirects {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(path) = self.stdin.as_ref() {
            write!(f, "<{}", path.display())?;
        }
        if let Some(path) = self.stdout.as_ref() {
            if self.stdin.is_some() {
                f.write_str(" ")?;
            }
            write!(f, ">{}", path.display())?;
        }
        Ok(())
    }
}

impl fmt::Display for PipelineElement {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, word) in self.words.iter().enumerate() {
            if i != 0 {
                f.write_str(" ")?;
            }
            f.write_str(word)?;
        }
        Ok(())
    }
}