refactor: improve repl completer (#199)

3 files changed, 147 insertions, 59 deletions

diff --git a/src/repl/completer.rs b/src/repl/completer.rs
new file mode 100644
index 0000000..ea24be8
--- /dev/null
+++ b/src/repl/completer.rs
@@ -0,0 +1,94 @@
+use std::collections::HashMap;
+
+use super::{parse_command, REPL_COMMANDS};
+
+use crate::config::GlobalConfig;
+
+use reedline::{Completer, Span, Suggestion};
+
+impl Completer for ReplCompleter {
+    fn complete(&mut self, line: &str, pos: usize) -> Vec<Suggestion> {
+        let mut suggestions = vec![];
+        if line.len() != pos {
+            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 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)
+                }))
+            }
+        }
+        suggestions
+    }
+}
+
+pub struct ReplCompleter {
+    config: GlobalConfig,
+    commands: Vec<(&'static str, &'static str)>,
+    groups: HashMap<&'static str, usize>,
+}
+
+impl ReplCompleter {
+    pub fn new(config: &GlobalConfig) -> Self {
+        let mut groups = HashMap::new();
+
+        let mut commands = REPL_COMMANDS.to_vec();
+        commands.sort_by(|(a, _), (b, _)| a.cmp(b));
+
+        for (name, _) in REPL_COMMANDS.iter() {
+            if let Some(count) = groups.get(name) {
+                groups.insert(*name, count + 1);
+            } else {
+                groups.insert(*name, 1);
+            }
+        }
+
+        Self {
+            config: config.clone(),
+            commands,
+            groups,
+        }
+    }
+}
+
+fn create_suggestion(value: String, description: Option<String>, span: Span) -> Suggestion {
+    Suggestion {
+        value,
+        description,
+        extra: None,
+        span,
+        append_whitespace: false,
+    }
+}
diff --git a/src/repl/highlighter.rs b/src/repl/highlighter.rs
index 37df659..7d3c197 100644
--- a/src/repl/highlighter.rs
+++ b/src/repl/highlighter.rs
@@ -1,18 +1,18 @@
+use super::REPL_COMMANDS;
+
 use crate::config::GlobalConfig;
 
 use nu_ansi_term::{Color, Style};
 use reedline::{Highlighter, StyledText};
 
 pub struct ReplHighlighter {
-    external_commands: Vec<String>,
     config: GlobalConfig,
 }
 
 impl ReplHighlighter {
-    pub fn new(external_commands: Vec<String>, config: GlobalConfig) -> Self {
+    pub fn new(config: &GlobalConfig) -> Self {
         Self {
-            external_commands,
-            config,
+            config: config.clone(),
         }
     }
 }
@@ -28,17 +28,11 @@ impl Highlighter for ReplHighlighter {
 
         let mut styled_text = StyledText::new();
 
-        if self
-            .external_commands
-            .clone()
-            .iter()
-            .any(|x| line.contains(x))
-        {
-            let matches: Vec<&str> = self
-                .external_commands
+        if REPL_COMMANDS.iter().any(|(cmd, _)| line.contains(cmd)) {
+            let matches: Vec<&str> = REPL_COMMANDS
                 .iter()
-                .filter(|c| line.contains(*c))
-                .map(std::ops::Deref::deref)
+                .filter(|(cmd, _)| line.contains(*cmd))
+                .map(|(cmd, _)| *cmd)
                 .collect();
             let longest_match = matches.iter().fold(String::new(), |acc, &item| {
                 if item.len() > acc.len() {
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index c19058b..6dba38e 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -1,7 +1,9 @@
+mod completer;
 mod highlighter;
 mod prompt;
 mod validator;
 
+use self::completer::ReplCompleter;
 use self::highlighter::ReplHighlighter;
 use self::prompt::ReplPrompt;
 use self::validator::ReplValidator;
@@ -19,8 +21,8 @@ use lazy_static::lazy_static;
 use reedline::Signal;
 use reedline::{
     default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
-    ColumnarMenu, DefaultCompleter, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
-    ReedlineEvent, ReedlineMenu, Vi,
+    ColumnarMenu, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline, ReedlineEvent,
+    ReedlineMenu, Vi,
 };
 use std::cell::RefCell;
 use std::io::Read;
@@ -45,7 +47,7 @@ const REPL_COMMANDS: [(&str, &str); 14] = [
 ];
 
 lazy_static! {
-    static ref COMMAND_RE: Regex = Regex::new(r"^\s*(\.\S+)\s*").unwrap();
+    static ref COMMAND_RE: Regex = Regex::new(r"^\s*(\.\S*)\s*").unwrap();
     static ref EDIT_RE: Regex = Regex::new(r"^\s*\.edit\s*").unwrap();
 }
 
@@ -59,36 +61,7 @@ pub struct Repl {
 
 impl Repl {
     pub fn init(config: GlobalConfig) -> Result<Self> {
-        let commands: Vec<String> = REPL_COMMANDS
-            .into_iter()
-            .map(|(v, _)| v.to_string())
-            .collect();
-
-        let completer = Self::create_completer(&config, &commands);
-        let highlighter = ReplHighlighter::new(commands, config.clone());
-        let menu = Self::create_menu();
-        let edit_mode: Box<dyn EditMode> = if config.read().keybindings.is_vi() {
-            let mut normal_keybindings = default_vi_normal_keybindings();
-            let mut insert_keybindings = default_vi_insert_keybindings();
-            Self::extra_keybindings(&mut normal_keybindings);
-            Self::extra_keybindings(&mut insert_keybindings);
-            Box::new(Vi::new(insert_keybindings, normal_keybindings))
-        } else {
-            let mut keybindings = default_emacs_keybindings();
-            Self::extra_keybindings(&mut keybindings);
-            Box::new(Emacs::new(keybindings))
-        };
-        let mut editor = Reedline::create()
-            .with_completer(Box::new(completer))
-            .with_highlighter(Box::new(highlighter))
-            .with_menu(menu)
-            .with_edit_mode(edit_mode)
-            .with_quick_completions(true)
-            .with_partial_completions(true)
-            .with_validator(Box::new(ReplValidator))
-            .with_ansi_colors(true);
-
-        editor.enable_bracketed_paste()?;
+        let editor = Self::create_editor(&config)?;
 
         let prompt = ReplPrompt::new(config.clone());
 
@@ -281,13 +254,24 @@ Type ".help" for more information.
         )
     }
 
-    fn create_completer(config: &GlobalConfig, commands: &[String]) -> DefaultCompleter {
-        let mut completion = commands.to_vec();
-        completion.extend(config.read().repl_completions());
-        let mut completer =
-            DefaultCompleter::with_inclusions(&['.', '-', '_', ':']).set_min_word_len(2);
-        completer.insert(completion.clone());
-        completer
+    fn create_editor(config: &GlobalConfig) -> Result<Reedline> {
+        let completer = ReplCompleter::new(config);
+        let highlighter = ReplHighlighter::new(config);
+        let menu = Self::create_menu();
+        let edit_mode = Self::create_edit_mode(config);
+        let mut editor = Reedline::create()
+            .with_completer(Box::new(completer))
+            .with_highlighter(Box::new(highlighter))
+            .with_menu(menu)
+            .with_edit_mode(edit_mode)
+            .with_quick_completions(true)
+            .with_partial_completions(true)
+            .with_validator(Box::new(ReplValidator))
+            .with_ansi_colors(true);
+
+        editor.enable_bracketed_paste()?;
+
+        Ok(editor)
     }
 
     fn extra_keybindings(keybindings: &mut Keybindings) {
@@ -306,6 +290,21 @@ Type ".help" for more information.
         );
     }
 
+    fn create_edit_mode(config: &GlobalConfig) -> Box<dyn EditMode> {
+        let edit_mode: Box<dyn EditMode> = if config.read().keybindings.is_vi() {
+            let mut normal_keybindings = default_vi_normal_keybindings();
+            let mut insert_keybindings = default_vi_insert_keybindings();
+            Self::extra_keybindings(&mut normal_keybindings);
+            Self::extra_keybindings(&mut insert_keybindings);
+            Box::new(Vi::new(insert_keybindings, normal_keybindings))
+        } else {
+            let mut keybindings = default_emacs_keybindings();
+            Self::extra_keybindings(&mut keybindings);
+            Box::new(Emacs::new(keybindings))
+        };
+        edit_mode
+    }
+
     fn create_menu() -> ReedlineMenu {
         let completion_menu = ColumnarMenu::default().with_name(MENU_NAME);
         ReedlineMenu::EngineCompleter(Box::new(completion_menu))
@@ -343,15 +342,15 @@ Press Ctrl+C to abort readline, Ctrl+D to exit the REPL"###,
 }
 
 fn parse_command(line: &str) -> Option<(&str, Option<&str>)> {
-    if let Ok(Some(captures)) = COMMAND_RE.captures(line) {
-        if let Some(cmd) = captures.get(1) {
-            let cmd = cmd.as_str();
+    match COMMAND_RE.captures(line) {
+        Ok(Some(captures)) => {
+            let cmd = captures.get(1)?.as_str();
             let args = line[captures[0].len()..].trim();
             let args = if args.is_empty() { None } else { Some(args) };
-            return Some((cmd, args));
+            Some((cmd, args))
         }
+        _ => None,
     }
-    None
 }
 #[cfg(test)]
 mod tests {
@@ -359,6 +358,7 @@ mod tests {
 
     #[test]
     fn test_process_command_line() {
+        assert_eq!(parse_command(" ."), Some((".", None)));
         assert_eq!(parse_command(" .role"), Some((".role", None)));
         assert_eq!(parse_command(" .role  "), Some((".role", None)));
         assert_eq!(