variable completion

3 files changed, 143 insertions, 59 deletions

diff --git a/src/completion.rs b/src/completion.rs
index 237e956..61dce76 100644
--- a/src/completion.rs
+++ b/src/completion.rs
@@ -1,5 +1,5 @@
 use crate::*;
-use std::fs;
+use std::{env, fs};
 
 pub struct Suggestion {
     pub display: BString,
@@ -55,3 +55,27 @@ pub fn path_completion(prefix: BString) -> 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() {
+        if var.starts_with(&prefix) {
+            out.push(Suggestion {
+                display: var.to_vec(),
+                delta: var[prefix.len()..].to_vec(),
+            });
+        }
+    }
+    drop(se);
+    for var in env::vars_os() {
+        let var = var.0.as_bytes();
+        if var.starts_with(&prefix) {
+            out.push(Suggestion {
+                display: var.to_vec(),
+                delta: var[prefix.len()..].to_vec(),
+            });
+        }
+    }
+    out
+}
diff --git a/src/main.rs b/src/main.rs
index 37d2ed3..6012a86 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -355,37 +355,39 @@ fn event_loop() {
                     &cmd,
                     &mut Executor::new_for_completion(session.clone()),
                 );
-                let mut se = session.lock().unwrap();
-                match comp.kind {
+
+                let suggestions = match comp.kind {
                     parse::CompletionKind::Command => todo!(),
-                    parse::CompletionKind::Argument => {
-                        let suggestions = completion::path_completion(comp.partial);
-                        if suggestions.len() == 0 {
-                            continue;
-                        }
+                    parse::CompletionKind::Argument => completion::path_completion(comp.partial),
+                    parse::CompletionKind::Variable => {
+                        completion::variable_completion(session.clone(), comp.partial)
+                    }
+                    parse::CompletionKind::None => continue,
+                };
 
-                        if suggestions.len() == 1 {
-                            // apply suggestion
-                            se.type_bytes(&suggestions[0].delta);
-                            continue;
-                        }
+                if suggestions.len() == 0 {
+                    continue;
+                }
 
-                        cursor::save();
+                let mut se = session.lock().unwrap();
 
-                        // one line below
-                        print!("\r\n");
-                        for s in suggestions {
-                            io::stdout().lock().write_all(&s.display).unwrap();
-                            println!();
-                        }
+                if suggestions.len() == 1 {
+                    // apply suggestion
+                    se.type_bytes(&suggestions[0].delta);
+                    continue;
+                }
 
-                        cursor::restore();
-                        stdout.lock().flush().unwrap();
-                    }
-                    parse::CompletionKind::None => {
-                        se.type_bytes(b"    ");
-                    }
+                cursor::save();
+
+                // one line below
+                print!("\r\n");
+                for s in suggestions {
+                    io::stdout().lock().write_all(&s.display).unwrap();
+                    println!();
                 }
+
+                cursor::restore();
+                stdout.lock().flush().unwrap();
             }
 
             // Escape sequence
diff --git a/src/parse/mod.rs b/src/parse/mod.rs
index 47e2d87..4f7dac5 100644
--- a/src/parse/mod.rs
+++ b/src/parse/mod.rs
@@ -61,12 +61,14 @@ pub fn pipes<const N: usize>(cmds: [Command<PreExpansion>; N]) -> Ast<PreExpansi
 pub fn estr(x: &[u8]) -> ExpString {
     ExpString {
         parts: vec![StringPart::Boring(x.to_vec())],
+        delim: b' ',
     }
 }
 
 pub fn str<const N: usize>(parts: [StringPart; N]) -> ExpString {
     ExpString {
         parts: parts.to_vec(),
+        delim: b' ',
     }
 }
 
@@ -78,6 +80,7 @@ pub fn var(x: &[u8]) -> StringPart {
     StringPart::Var(Var {
         name: VarName { name: x.to_vec() },
         default: None,
+        already_complete: true,
     })
 }
 
@@ -85,6 +88,7 @@ pub fn var_default(x: &[u8], default: ExpString) -> StringPart {
     StringPart::Var(Var {
         name: VarName { name: x.to_vec() },
         default: Some(default),
+        already_complete: true,
     })
 }
 
@@ -338,6 +342,20 @@ pub enum StringPart {
 pub struct Var {
     name: VarName,
     default: Option<ExpString>,
+    /// if pressing tab right after the parsed variable should not try to complete the variable
+    ///
+    /// i.e. `${HOM}` -> true, `$HOM` -> false, `${HOM` -> false
+    already_complete: bool,
+}
+
+impl Var {
+    pub fn new(name: VarName) -> Self {
+        Self {
+            name,
+            default: None,
+            already_complete: false,
+        }
+    }
 }
 
 impl StringPart {
@@ -356,6 +374,7 @@ impl StringPart {
 /// `"hi ${var} $(cmd) "` gets mapped to `[Boring("hi "), Var("var"), String(" "), Cmd(...), Boring(" ")]`
 pub struct ExpString {
     parts: Vec<StringPart>,
+    delim: u8,
 }
 
 impl ExpString {
@@ -456,7 +475,7 @@ impl Parse for ExpString {
                 b.adv();
             } else if is_symbol(delim) && delim != b'$' && delim != b'\\' {
                 return if already_parsed {
-                    Ok(Self { parts })
+                    Ok(Self { parts, delim })
                 } else {
                     Err(ParseError::NotAString)
                 };
@@ -507,7 +526,7 @@ impl Parse for ExpString {
                     if x == b'\'' || x == b'"' {
                         break;
                     } else {
-                        return Ok(Self { parts });
+                        return Ok(Self { parts, delim });
                     }
                 }
 
@@ -540,23 +559,19 @@ impl Parse for ExpString {
 
                     if x == b'?' || x == b'!' {
                         b.adv();
-                        p.push(StringPart::Var(Var {
-                            name: VarName { name: vec![x] },
-                            default: None,
-                        }))
+                        p.push(StringPart::Var(Var::new(VarName { name: vec![x] })))
                     } else if is_var_begin(x) {
                         let v = VarName::parse(b)?;
-                        p.push(StringPart::Var(Var {
-                            name: v,
-                            default: None,
-                        }));
+                        p.push(StringPart::Var(Var::new(v)));
                     } else if x == b'{' {
                         b.adv();
                         let v = VarName::parse(b)?;
                         let mut default = None;
 
                         if !b.has() {
-                            return Err(ParseError::Eof);
+                            if !b.is_completion() {
+                                return Err(ParseError::Eof);
+                            }
                         } else if b.peek() == b':' {
                             b.adv();
                             if !b.has() {
@@ -571,13 +586,24 @@ impl Parse for ExpString {
                         }
 
                         if !b.has() {
-                            return Err(ParseError::Eof);
+                            if !b.is_completion() {
+                                return Err(ParseError::Eof);
+                            }
                         } else if b.peek() != b'}' {
-                            return Err(ParseError::Incomplete);
+                            return Err(ParseError::Expected('}'));
                         }
 
-                        b.adv();
-                        p.push(StringPart::Var(Var { name: v, default }));
+                        let already_complete = b.has();
+
+                        if already_complete {
+                            b.adv();
+                        }
+
+                        p.push(StringPart::Var(Var {
+                            name: v,
+                            default,
+                            already_complete,
+                        }));
                     } else if x == b'(' {
                         b.adv();
                         let cmd = Ast::parse(b)?;
@@ -608,6 +634,7 @@ impl Parse for ExpString {
                             name: b"HOME".to_vec(),
                         },
                         default: None,
+                        already_complete: true,
                     }));
                 } else {
                     add_char(p, x);
@@ -626,7 +653,7 @@ impl Parse for ExpString {
         }
 
         if b.is_completion() || already_parsed {
-            Ok(Self { parts })
+            Ok(Self { parts, delim: b' ' })
         } else {
             Err(ParseError::Eof)
         }
@@ -687,6 +714,7 @@ pub fn do_parse(x: &[u8]) -> Res<Ast<PreExpansion>, (ParseError, &[u8])> {
 pub enum CompletionKind {
     Command,
     Argument,
+    Variable,
     None,
 }
 
@@ -704,30 +732,60 @@ impl CompletionContext {
     }
 }
 
-fn expstr_cc<E: Expander>(s: &ExpString, kind: CompletionKind, e: &mut E) -> CompletionContext {
-    match s.clone().expand(e) {
-        Ok(buf) => CompletionContext { kind, partial: buf },
-        Err(_) => CompletionContext::none(),
+impl Ast<PreExpansion> {
+    fn completion<E: Expander>(&self, e: &mut E) -> CompletionContext {
+        match self {
+            Ast::FunDecl(fd) => fd.body.body.completion(e),
+            Ast::VarAssign(va) => va.val.completion(e, CompletionKind::Argument),
+            Ast::Pipes(p) => p.completion(e),
+        }
+    }
+}
+
+impl ExpString {
+    fn completion<E: Expander>(&self, e: &mut E, kind: CompletionKind) -> CompletionContext {
+        if let Some(StringPart::Var(var)) = self.parts.last()
+            && !var.already_complete
+        {
+            CompletionContext {
+                kind: CompletionKind::Variable,
+                partial: var.name.name.clone(),
+            }
+        } else if let Ok(s) = self.clone().expand(e) {
+            CompletionContext { kind, partial: s }
+        } else {
+            CompletionContext::none()
+        }
+    }
+}
+
+impl Pipes<PreExpansion> {
+    fn completion<E: Expander>(&self, e: &mut E) -> CompletionContext {
+        let Some(cmd) = self.cmds.last() else {
+            return CompletionContext::none();
+        };
+
+        if let Some(arg) = cmd.args.last() {
+            arg.completion(e, CompletionKind::Argument)
+        } else {
+            cmd.cmd.completion(e, CompletionKind::Command)
+        }
     }
 }
 
 pub fn completion_context<'a, E: Expander>(x: &'a [u8], e: &mut E) -> CompletionContext {
     let mut cursor = Cursor::new(x, ParseMode::Completion);
     let ast = Ast::parse(&mut cursor);
-    match ast {
-        Ok(Ast::Pipes(pipes)) if cursor.spaced == false => {
-            if let Some(cmd) = pipes.cmds.last() {
-                if cmd.args.is_empty() {
-                    expstr_cc(&cmd.cmd, CompletionKind::Command, e)
-                } else {
-                    expstr_cc(&cmd.args[cmd.args.len() - 1], CompletionKind::Argument, e)
-                }
-            } else {
-                CompletionContext::none()
-            }
-        }
-        _ => CompletionContext::none(),
+
+    let Ok(ast) = ast else {
+        return CompletionContext::none();
+    };
+
+    if cursor.spaced {
+        return CompletionContext::none();
     }
+
+    ast.completion(e)
 }
 
 trait Parse: Sized {