variable refactor: completion also for magic vars

3 files changed, 82 insertions, 18 deletions

diff --git a/src/run/builtin.rs b/src/run/builtin.rs
index 9f3ff62..e928d8b 100644
--- a/src/run/builtin.rs
+++ b/src/run/builtin.rs
@@ -460,7 +460,7 @@ impl Builtin for var {
             return Err(Error::Exit(-2));
         }
         let se = session.lock().unwrap();
-        if let Some(val) = se.vars.get(&args[0]) {
+        if let Some(val) = se.vars.lookup(&args[0]) {
             stdout.write_all(val.as_slice())?;
             Ok(())
         } else {
diff --git a/src/run/mod.rs b/src/run/mod.rs
index 5b2d32d..55bdd29 100644
--- a/src/run/mod.rs
+++ b/src/run/mod.rs
@@ -9,6 +9,9 @@ use crate::wait::{ChildWaiter, ThreadWaiter};
 use crate::*;
 
 mod builtin;
+mod var;
+
+pub use var::Vars;
 
 #[derive(Debug)]
 pub enum ExecError {
@@ -320,7 +323,7 @@ impl Executor {
     }
 
     fn execute_var_assign(&mut self, va: parse::VarAssign<PostExpansion>) -> SpawnedCmd {
-        self.se.lock().unwrap().vars.insert(va.var, va.val);
+        self.se.lock().unwrap().vars.set(va.var, va.val);
         SpawnedCmd::Joined(Ok(()))
     }
 
@@ -455,18 +458,11 @@ impl parse::Expander for Executor {
         var: BString,
         default: Option<BString>,
     ) -> Result<BString, Self::Error> {
-        match &var[..] {
-            b"CWD_PRETTY" => return Ok(crate::pretty_cwd().into_bytes()),
-            b"PISH_VERSION" => return Ok(crate::consts::PISH_VERSION.as_bytes().to_vec()),
-            b"PISH_COMMIT" => return Ok(crate::consts::PISH_COMMIT.as_bytes().to_vec()),
-            b"PISH_DIRTY" => return Ok(vec![crate::consts::PISH_DIRTY as u8 + b'0']),
-            _ => {}
-        }
-
         if var.is_empty() {
             return Err(ExecError::UnknownVariable(var));
         }
 
+        // TODO: this should probably also go somewhere in the var module.
         if var[0].is_ascii_digit()
             && let Some(x) = String::from_utf8(var.clone())
                 .ok()
@@ -481,17 +477,15 @@ impl parse::Expander for Executor {
             }
         }
 
-        if let Some(val) = self.se.lock().unwrap().vars.get(&var) {
-            return Ok(val.clone());
+        if let Some(val) = self.se.lock().unwrap().vars.lookup(&var) {
+            return Ok(val);
         }
 
-        match std::env::var_os(OsStr::from_bytes(&var)) {
-            Some(val) => Ok(val.as_bytes().to_vec()),
-            None => match default {
-                Some(d) => Ok(d),
-                None => Err(ExecError::UnknownVariable(var)),
-            },
+        if let Some(dfl) = default {
+            return Ok(dfl);
         }
+
+        Err(ExecError::UnknownVariable(var))
     }
 
     fn expand_cmd(&mut self, ast: Ast<parse::PostExpansion>) -> Result<BString, Self::Error> {
diff --git a/src/run/var.rs b/src/run/var.rs
new file mode 100644
index 0000000..0e64ad7
--- /dev/null
+++ b/src/run/var.rs
@@ -0,0 +1,70 @@
+use std::collections::{HashMap, HashSet};
+
+use crate::BString;
+use crate::bstr;
+
+pub struct Vars {
+    simple: HashMap<BString, BString>,
+    magic: HashMap<BString, fn() -> BString>,
+    all: HashSet<BString>,
+}
+
+impl Vars {
+    fn new(
+        mut simple: HashMap<BString, BString>,
+        magic: HashMap<BString, fn() -> BString>,
+    ) -> Self {
+        for (var, val) in std::env::vars_os() {
+            simple.insert(var.into_encoded_bytes(), val.into_encoded_bytes());
+        }
+        let all = simple
+            .keys()
+            .cloned()
+            .chain(magic.keys().cloned())
+            .collect();
+        Self { simple, magic, all }
+    }
+
+    pub fn set(&mut self, var: BString, val: BString) {
+        self.simple.insert(var.clone(), val);
+        self.all.insert(var);
+    }
+
+    pub fn lookup(&self, var: &bstr) -> Option<BString> {
+        if let Some(fun) = self.magic.get(var) {
+            return Some(fun());
+        }
+
+        if let Some(val) = self.simple.get(var) {
+            return Some(val.clone());
+        }
+
+        None
+    }
+
+    pub fn vars(&self) -> &HashSet<BString> {
+        &self.all
+    }
+}
+
+macro_rules! map {
+    ($($key:expr => $val:expr),* $(,)?) => {{
+        let mut map = HashMap::<BString, _, _>::new();
+        $(map.insert($key.into(), $val);)*
+        map
+    }};
+}
+
+impl Default for Vars {
+    fn default() -> Self {
+        let simple = map! {
+            b"PISH_VERSION" => crate::consts::PISH_VERSION.as_bytes().to_vec(),
+            b"PISH_COMMIT" => crate::consts::PISH_COMMIT.as_bytes().to_vec(),
+            b"PISH_DIRTY" => vec![crate::consts::PISH_DIRTY as u8 + b'0'],
+        };
+        let magic = map! {
+            b"CWD_PRETTY" => crate::pretty_cwd as _
+        };
+        Self::new(simple, magic)
+    }
+}