variable refactor: completion also for magic vars

5 files changed, 91 insertions, 25 deletions

diff --git a/src/completion.rs b/src/completion.rs
index 0e24d6d..2913dba 100644
--- a/src/completion.rs
+++ b/src/completion.rs
@@ -90,7 +90,7 @@ pub fn path_exe_completion(cc: CompletionContext) -> Vec<Suggestion> {
 pub fn variable_completion(session: Arc<Mutex<Session>>, prefix: BString) -> Vec<Suggestion> {
     let se = session.lock().unwrap();
     let mut out = Vec::new();
-    for var in se.vars.keys() {
+    for var in se.vars.vars() {
         if var.starts_with(&prefix) {
             out.push(Suggestion {
                 display: var.to_vec(),
diff --git a/src/lib.rs b/src/lib.rs
index 700425a..ef62c44 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -97,7 +97,7 @@ pub struct Session {
     history: Vec<HistoryEntry>,
     prev_path: BString,
     builtins: HashMap<BString, &'static dyn run::BuiltinClone>,
-    vars: HashMap<BString, BString>,
+    vars: run::Vars,
     funs: HashMap<BString, Block>,
     aliases: run::Aliases,
     socket_running: Option<export_fun::SocketRunning>,
@@ -125,7 +125,7 @@ impl Session {
             history: Vec::new(),
             prev_path: b".".into(),
             builtins: HashMap::new(),
-            vars: HashMap::new(),
+            vars: run::Vars::default(),
             funs: HashMap::new(),
             aliases: run::Aliases::new(),
             socket_running: None,
@@ -171,13 +171,15 @@ fn pretty_cwd_res() -> io::Result<String> {
     Ok(s)
 }
 
-fn pretty_cwd() -> String {
-    pretty_cwd_res().unwrap_or_else(|_| String::new())
+fn pretty_cwd() -> BString {
+    pretty_cwd_res()
+        .unwrap_or_else(|_| String::new())
+        .into_bytes()
 }
 
 impl Session {
     fn load_unevaled_prompt(&self) -> BString {
-        match self.vars.get(&b"PROMPT"[..]) {
+        match self.vars.lookup(&b"PROMPT"[..]) {
             Some(prompt) => prompt.clone(),
             None => {
                 let mut prompt = BString::new();
@@ -485,7 +487,7 @@ pub fn event_loop() {
         prev_path: vec![b'.'],
         history_visit: 0,
         socket_running: None,
-        vars: HashMap::new(),
+        vars: run::Vars::default(),
         funs: HashMap::new(),
         aliases: run::Aliases::new(),
         path_cache: Default::default(),
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)
+    }
+}