refactor: improve repl completer (#226)

2 files changed, 127 insertions, 58 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 22c58bb..3d4df7d 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -413,28 +413,44 @@ impl Config {
             .unwrap_or_default()
     }
 
-    pub fn repl_complete(&self, cmd: &str, args: &str) -> Vec<String> {
-        let possible_values = match cmd {
-            ".role" => self.roles.iter().map(|v| v.name.clone()).collect(),
-            ".model" => list_models(self).into_iter().map(|v| v.id()).collect(),
-            ".session" => self.list_sessions(),
-            ".set" => {
-                vec![
-                    "temperature ".into(),
-                    format!("save {}", !self.save),
-                    format!("highlight {}", !self.highlight),
-                    format!("dry_run {}", !self.dry_run),
-                    format!("auto_copy {}", !self.auto_copy),
+    pub fn repl_complete(&self, cmd: &str, args: &[&str]) -> Vec<String> {
+        let (values, filter) = if args.len() == 1 {
+            let values = match cmd {
+                ".role" => self.roles.iter().map(|v| v.name.clone()).collect(),
+                ".model" => list_models(self).into_iter().map(|v| v.id()).collect(),
+                ".session" => self.list_sessions(),
+                ".set" => vec![
+                    "temperature ",
+                    "save ",
+                    "highlight ",
+                    "dry_run ",
+                    "auto_copy ",
                 ]
-            }
-            _ => vec![],
+                .into_iter()
+                .map(|v| v.to_string())
+                .collect(),
+                _ => vec![],
+            };
+            (values, args[0])
+        } else if args.len() == 2 {
+            let to_vec = |v: bool| vec![v.to_string()];
+            let values = match args[0] {
+                "save" => to_vec(!self.save),
+                "highlight" => to_vec(!self.highlight),
+                "dry_run" => to_vec(!self.dry_run),
+                "auto_copy" => to_vec(!self.auto_copy),
+                _ => vec![],
+            };
+            (values, args[1])
+        } else {
+            return vec![];
         };
-        let mut possible_values: Vec<String> = possible_values
+        let mut values: Vec<String> = values
             .into_iter()
-            .filter(|v| v.starts_with(args))
+            .filter(|v| v.starts_with(filter))
             .collect();
-        possible_values.sort_unstable();
-        possible_values
+        values.sort_unstable();
+        values
     }
 
     pub fn update(&mut self, data: &str) -> Result<()> {
diff --git a/src/repl/completer.rs b/src/repl/completer.rs
index ea24be8..059c7ca 100644
--- a/src/repl/completer.rs
+++ b/src/repl/completer.rs
@@ -1,54 +1,69 @@
-use std::collections::HashMap;
-
-use super::{parse_command, REPL_COMMANDS};
+use super::REPL_COMMANDS;
 
 use crate::config::GlobalConfig;
 
 use reedline::{Completer, Span, Suggestion};
+use std::collections::HashMap;
 
 impl Completer for ReplCompleter {
     fn complete(&mut self, line: &str, pos: usize) -> Vec<Suggestion> {
         let mut suggestions = vec![];
-        if line.len() != pos {
+        let line = &line[0..pos];
+        let mut parts = split_line(line);
+        if parts.is_empty() {
             return suggestions;
         }
-        let line = &line[0..pos];
-        if let Some((cmd, args)) = parse_command(line) {
-            let commands: Vec<_> = self
-                .commands
-                .iter()
-                .filter(|(cmd_name, _)| match args {
-                    Some(args) => cmd_name.starts_with(&format!("{cmd} {args}")),
-                    None => cmd_name.starts_with(cmd),
-                })
-                .collect();
-
-            if args.is_some() || line.ends_with(' ') {
-                let args = args.unwrap_or_default();
-                let start = line.chars().take_while(|c| *c == ' ').count() + cmd.len() + 1;
-                let span = Span::new(start, pos);
-                suggestions.extend(
-                    self.config
-                        .read()
-                        .repl_complete(cmd, args)
-                        .iter()
-                        .map(|name| create_suggestion(name.clone(), None, span)),
-                )
-            }
+        if parts[0].0 == r#"""""# {
+            parts.remove(0);
+        }
 
-            if suggestions.is_empty() {
-                let start = line.chars().take_while(|c| *c == ' ').count();
-                let span = Span::new(start, pos);
-                suggestions.extend(commands.iter().map(|(name, desc)| {
-                    let has_group = self.groups.get(name).map(|v| *v > 1).unwrap_or_default();
-                    let name = if has_group {
-                        name.to_string()
-                    } else {
-                        format!("{name} ")
-                    };
-                    create_suggestion(name, Some(desc.to_string()), span)
-                }))
-            }
+        let parts_len = parts.len();
+        if parts_len == 0 {
+            return suggestions;
+        }
+        let (cmd, cmd_start) = parts[0];
+
+        if !cmd.starts_with('.') {
+            return suggestions;
+        }
+
+        let commands: Vec<_> = self
+            .commands
+            .iter()
+            .filter(|(cmd_name, _)| {
+                let line = parts
+                    .iter()
+                    .take(2)
+                    .map(|(v, _)| *v)
+                    .collect::<Vec<&str>>()
+                    .join(" ");
+                cmd_name.starts_with(&line)
+            })
+            .collect();
+
+        if parts_len > 1 {
+            let span = Span::new(parts[parts_len - 1].1, pos);
+            let args: Vec<&str> = parts.iter().skip(1).map(|(v, _)| *v).collect();
+            suggestions.extend(
+                self.config
+                    .read()
+                    .repl_complete(cmd, &args)
+                    .iter()
+                    .map(|name| create_suggestion(name.clone(), None, span)),
+            )
+        }
+
+        if suggestions.is_empty() {
+            let span = Span::new(cmd_start, pos);
+            suggestions.extend(commands.iter().map(|(name, desc)| {
+                let has_group = self.groups.get(name).map(|v| *v > 1).unwrap_or_default();
+                let name = if has_group {
+                    name.to_string()
+                } else {
+                    format!("{name} ")
+                };
+                create_suggestion(name, Some(desc.to_string()), span)
+            }))
         }
         suggestions
     }
@@ -92,3 +107,41 @@ fn create_suggestion(value: String, description: Option<String>, span: Span) ->
         append_whitespace: false,
     }
 }
+
+fn split_line(line: &str) -> Vec<(&str, usize)> {
+    let mut parts = vec![];
+    let mut part_start = None;
+    for (i, ch) in line.char_indices() {
+        if ch == ' ' {
+            if let Some(s) = part_start {
+                parts.push((&line[s..i], s));
+                part_start = None;
+            }
+        } else if part_start.is_none() {
+            part_start = Some(i)
+        }
+    }
+    if let Some(s) = part_start {
+        parts.push((&line[s..], s));
+    } else {
+        parts.push(("", line.len()))
+    }
+    parts
+}
+
+#[test]
+fn test_split_line() {
+    assert_eq!(split_line(".role coder"), vec![(".role", 0), ("coder", 6)],);
+    assert_eq!(
+        split_line(" .role   coder"),
+        vec![(".role", 1), ("coder", 9)],
+    );
+    assert_eq!(
+        split_line(".set highlight "),
+        vec![(".set", 0), ("highlight", 5), ("", 15)],
+    );
+    assert_eq!(
+        split_line(".set highlight t"),
+        vec![(".set", 0), ("highlight", 5), ("t", 15)],
+    );
+}