From 27aa919f1c5ce9bbee899d492f3b026a90fba27a Mon Sep 17 00:00:00 2001
From: Kyle Smith <askreet@gmail.com>
Date: Mon, 13 Mar 2023 22:54:46 -0400
Subject: [PATCH] Only complete appropriate arguments for typed commands.
 (#5966)

---
 helix-term/src/commands/typed.rs | 238 ++++++++++++++++++++-----------
 1 file changed, 154 insertions(+), 84 deletions(-)

diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index e9a7222586a9..ae4163fa4c46 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -5,7 +5,7 @@ use crate::job::Job;
 
 use super::*;
 
-use helix_core::encoding;
+use helix_core::{encoding, shellwords::Shellwords};
 use helix_view::document::DEFAULT_LANGUAGE_NAME;
 use helix_view::editor::{Action, CloseError, ConfigEvent};
 use serde_json::Value;
@@ -18,7 +18,49 @@ pub struct TypableCommand {
     pub doc: &'static str,
     // params, flags, helper, completer
     pub fun: fn(&mut compositor::Context, &[Cow<str>], PromptEvent) -> anyhow::Result<()>,
-    pub completer: Option<Completer>,
+    /// What completion methods, if any, does this command have?
+    pub signature: CommandSignature,
+}
+
+impl TypableCommand {
+    fn completer_for_argument_number(&self, n: usize) -> &Completer {
+        match self.signature.positional_args.get(n) {
+            Some(completer) => completer,
+            _ => &self.signature.var_args,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct CommandSignature {
+    // Arguments with specific completion methods based on their position.
+    positional_args: &'static [Completer],
+
+    // All remaining arguments will use this completion method, if set.
+    var_args: Completer,
+}
+
+impl CommandSignature {
+    const fn none() -> Self {
+        Self {
+            positional_args: &[],
+            var_args: completers::none,
+        }
+    }
+
+    const fn positional(completers: &'static [Completer]) -> Self {
+        Self {
+            positional_args: completers,
+            var_args: completers::none,
+        }
+    }
+
+    const fn all(completer: Completer) -> Self {
+        Self {
+            positional_args: &[],
+            var_args: completer,
+        }
+    }
 }
 
 fn quit(cx: &mut compositor::Context, args: &[Cow<str>], event: PromptEvent) -> anyhow::Result<()> {
@@ -2113,112 +2155,114 @@ pub const TYPABLE_COMMAND_LIST: &[TypableCommand] = &[
             aliases: &["q"],
             doc: "Close the current view.",
             fun: quit,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "quit!",
             aliases: &["q!"],
             doc: "Force close the current view, ignoring unsaved changes.",
             fun: force_quit,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "open",
             aliases: &["o"],
             doc: "Open a file from disk into the current view.",
             fun: open,
-            completer: Some(completers::filename),
+            signature: CommandSignature::all(completers::filename),
         },
         TypableCommand {
             name: "buffer-close",
             aliases: &["bc", "bclose"],
             doc: "Close the current buffer.",
             fun: buffer_close,
-            completer: Some(completers::buffer),
+            signature: CommandSignature::all(completers::buffer),
         },
         TypableCommand {
             name: "buffer-close!",
             aliases: &["bc!", "bclose!"],
             doc: "Close the current buffer forcefully, ignoring unsaved changes.",
             fun: force_buffer_close,
-            completer: Some(completers::buffer),
+            signature: CommandSignature::all(completers::buffer)
         },
         TypableCommand {
             name: "buffer-close-others",
             aliases: &["bco", "bcloseother"],
             doc: "Close all buffers but the currently focused one.",
             fun: buffer_close_others,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "buffer-close-others!",
             aliases: &["bco!", "bcloseother!"],
             doc: "Force close all buffers but the currently focused one.",
             fun: force_buffer_close_others,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "buffer-close-all",
             aliases: &["bca", "bcloseall"],
             doc: "Close all buffers without quitting.",
             fun: buffer_close_all,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "buffer-close-all!",
             aliases: &["bca!", "bcloseall!"],
             doc: "Force close all buffers ignoring unsaved changes without quitting.",
             fun: force_buffer_close_all,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "buffer-next",
             aliases: &["bn", "bnext"],
             doc: "Goto next buffer.",
             fun: buffer_next,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "buffer-previous",
             aliases: &["bp", "bprev"],
             doc: "Goto previous buffer.",
             fun: buffer_previous,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "write",
             aliases: &["w"],
             doc: "Write changes to disk. Accepts an optional path (:write some/path.txt)",
             fun: write,
-            completer: Some(completers::filename),
+            signature: CommandSignature::positional(&[completers::filename]),
         },
         TypableCommand {
             name: "write!",
             aliases: &["w!"],
             doc: "Force write changes to disk creating necessary subdirectories. Accepts an optional path (:write some/path.txt)",
             fun: force_write,
-            completer: Some(completers::filename),
+            signature: CommandSignature::positional(&[completers::filename]),
         },
         TypableCommand {
             name: "new",
             aliases: &["n"],
             doc: "Create a new scratch buffer.",
             fun: new_file,
-            completer: Some(completers::filename),
+            // TODO: This seems to complete with a filename, but doesn't use that filename to
+            //       set the path of the newly created buffer.
+            signature: CommandSignature::positional(&[completers::filename]),
         },
         TypableCommand {
             name: "format",
             aliases: &["fmt"],
             doc: "Format the file using the LSP formatter.",
             fun: format,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "indent-style",
             aliases: &[],
             doc: "Set the indentation style for editing. ('t' for tabs or 1-8 for number of spaces.)",
             fun: set_indent_style,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "line-ending",
@@ -2228,231 +2272,231 @@ pub const TYPABLE_COMMAND_LIST: &[TypableCommand] = &[
             #[cfg(feature = "unicode-lines")]
             doc: "Set the document's default line ending. Options: crlf, lf, cr, ff, nel.",
             fun: set_line_ending,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "earlier",
             aliases: &["ear"],
             doc: "Jump back to an earlier point in edit history. Accepts a number of steps or a time span.",
             fun: earlier,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "later",
             aliases: &["lat"],
             doc: "Jump to a later point in edit history. Accepts a number of steps or a time span.",
             fun: later,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "write-quit",
             aliases: &["wq", "x"],
             doc: "Write changes to disk and close the current view. Accepts an optional path (:wq some/path.txt)",
             fun: write_quit,
-            completer: Some(completers::filename),
+            signature: CommandSignature::positional(&[completers::filename]),
         },
         TypableCommand {
             name: "write-quit!",
             aliases: &["wq!", "x!"],
             doc: "Write changes to disk and close the current view forcefully. Accepts an optional path (:wq! some/path.txt)",
             fun: force_write_quit,
-            completer: Some(completers::filename),
+            signature: CommandSignature::positional(&[completers::filename]),
         },
         TypableCommand {
             name: "write-all",
             aliases: &["wa"],
             doc: "Write changes from all buffers to disk.",
             fun: write_all,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "write-quit-all",
             aliases: &["wqa", "xa"],
             doc: "Write changes from all buffers to disk and close all views.",
             fun: write_all_quit,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "write-quit-all!",
             aliases: &["wqa!", "xa!"],
             doc: "Write changes from all buffers to disk and close all views forcefully (ignoring unsaved changes).",
             fun: force_write_all_quit,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "quit-all",
             aliases: &["qa"],
             doc: "Close all views.",
             fun: quit_all,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "quit-all!",
             aliases: &["qa!"],
             doc: "Force close all views ignoring unsaved changes.",
             fun: force_quit_all,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "cquit",
             aliases: &["cq"],
             doc: "Quit with exit code (default 1). Accepts an optional integer exit code (:cq 2).",
             fun: cquit,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "cquit!",
             aliases: &["cq!"],
             doc: "Force quit with exit code (default 1) ignoring unsaved changes. Accepts an optional integer exit code (:cq! 2).",
             fun: force_cquit,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "theme",
             aliases: &[],
             doc: "Change the editor theme (show current theme if no name specified).",
             fun: theme,
-            completer: Some(completers::theme),
+            signature: CommandSignature::positional(&[completers::theme]),
         },
         TypableCommand {
             name: "clipboard-yank",
             aliases: &[],
             doc: "Yank main selection into system clipboard.",
             fun: yank_main_selection_to_clipboard,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "clipboard-yank-join",
             aliases: &[],
             doc: "Yank joined selections into system clipboard. A separator can be provided as first argument. Default value is newline.", // FIXME: current UI can't display long doc.
             fun: yank_joined_to_clipboard,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "primary-clipboard-yank",
             aliases: &[],
             doc: "Yank main selection into system primary clipboard.",
             fun: yank_main_selection_to_primary_clipboard,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "primary-clipboard-yank-join",
             aliases: &[],
             doc: "Yank joined selections into system primary clipboard. A separator can be provided as first argument. Default value is newline.", // FIXME: current UI can't display long doc.
             fun: yank_joined_to_primary_clipboard,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "clipboard-paste-after",
             aliases: &[],
             doc: "Paste system clipboard after selections.",
             fun: paste_clipboard_after,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "clipboard-paste-before",
             aliases: &[],
             doc: "Paste system clipboard before selections.",
             fun: paste_clipboard_before,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "clipboard-paste-replace",
             aliases: &[],
             doc: "Replace selections with content of system clipboard.",
             fun: replace_selections_with_clipboard,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "primary-clipboard-paste-after",
             aliases: &[],
             doc: "Paste primary clipboard after selections.",
             fun: paste_primary_clipboard_after,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "primary-clipboard-paste-before",
             aliases: &[],
             doc: "Paste primary clipboard before selections.",
             fun: paste_primary_clipboard_before,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "primary-clipboard-paste-replace",
             aliases: &[],
             doc: "Replace selections with content of system primary clipboard.",
             fun: replace_selections_with_primary_clipboard,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "show-clipboard-provider",
             aliases: &[],
             doc: "Show clipboard provider name in status bar.",
             fun: show_clipboard_provider,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "change-current-directory",
             aliases: &["cd"],
             doc: "Change the current working directory.",
             fun: change_current_directory,
-            completer: Some(completers::directory),
+            signature: CommandSignature::positional(&[completers::directory]),
         },
         TypableCommand {
             name: "show-directory",
             aliases: &["pwd"],
             doc: "Show the current working directory.",
             fun: show_current_directory,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "encoding",
             aliases: &[],
             doc: "Set encoding. Based on `https://encoding.spec.whatwg.org`.",
             fun: set_encoding,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "character-info",
             aliases: &["char"],
             doc: "Get info about the character under the primary cursor.",
             fun: get_character_info,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "reload",
             aliases: &[],
             doc: "Discard changes and reload from the source file.",
             fun: reload,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "reload-all",
             aliases: &[],
             doc: "Discard changes and reload all documents from the source files.",
             fun: reload_all,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "update",
             aliases: &[],
             doc: "Write changes only if the file has been modified.",
             fun: update,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "lsp-workspace-command",
             aliases: &[],
             doc: "Open workspace command picker",
             fun: lsp_workspace_command,
-            completer: Some(completers::lsp_workspace_command),
+            signature: CommandSignature::positional(&[completers::lsp_workspace_command]),
         },
         TypableCommand {
             name: "lsp-restart",
             aliases: &[],
             doc: "Restarts the Language Server that is in use by the current doc",
             fun: lsp_restart,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "lsp-stop",
@@ -2466,182 +2510,183 @@ pub const TYPABLE_COMMAND_LIST: &[TypableCommand] = &[
             aliases: &[],
             doc: "Display tree sitter scopes, primarily for theming and development.",
             fun: tree_sitter_scopes,
-            completer: None,
-       },
+            signature: CommandSignature::none(),
+        },
         TypableCommand {
             name: "debug-start",
             aliases: &["dbg"],
             doc: "Start a debug session from a given template with given parameters.",
             fun: debug_start,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "debug-remote",
             aliases: &["dbg-tcp"],
             doc: "Connect to a debug adapter by TCP address and start a debugging session from a given template with given parameters.",
             fun: debug_remote,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "debug-eval",
             aliases: &[],
             doc: "Evaluate expression in current debug context.",
             fun: debug_eval,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "vsplit",
             aliases: &["vs"],
             doc: "Open the file in a vertical split.",
             fun: vsplit,
-            completer: Some(completers::filename),
+            signature: CommandSignature::all(completers::filename)
         },
         TypableCommand {
             name: "vsplit-new",
             aliases: &["vnew"],
             doc: "Open a scratch buffer in a vertical split.",
             fun: vsplit_new,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "hsplit",
             aliases: &["hs", "sp"],
             doc: "Open the file in a horizontal split.",
             fun: hsplit,
-            completer: Some(completers::filename),
+            signature: CommandSignature::all(completers::filename)
         },
         TypableCommand {
             name: "hsplit-new",
             aliases: &["hnew"],
             doc: "Open a scratch buffer in a horizontal split.",
             fun: hsplit_new,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "tutor",
             aliases: &[],
             doc: "Open the tutorial.",
             fun: tutor,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "goto",
             aliases: &["g"],
             doc: "Goto line number.",
             fun: goto_line_number,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "set-language",
             aliases: &["lang"],
             doc: "Set the language of current buffer (show current language if no value specified).",
             fun: language,
-            completer: Some(completers::language),
+            signature: CommandSignature::positional(&[completers::language]),
         },
         TypableCommand {
             name: "set-option",
             aliases: &["set"],
             doc: "Set a config option at runtime.\nFor example to disable smart case search, use `:set search.smart-case false`.",
             fun: set_option,
-            completer: Some(completers::setting),
+            // TODO: Add support for completion of the options value(s), when appropriate.
+            signature: CommandSignature::positional(&[completers::setting]),
         },
         TypableCommand {
             name: "toggle-option",
             aliases: &["toggle"],
             doc: "Toggle a boolean config option at runtime.\nFor example to toggle smart case search, use `:toggle search.smart-case`.",
             fun: toggle_option,
-            completer: Some(completers::setting),
+            signature: CommandSignature::positional(&[completers::setting]),
         },
         TypableCommand {
             name: "get-option",
             aliases: &["get"],
             doc: "Get the current value of a config option.",
             fun: get_option,
-            completer: Some(completers::setting),
+            signature: CommandSignature::positional(&[completers::setting]),
         },
         TypableCommand {
             name: "sort",
             aliases: &[],
             doc: "Sort ranges in selection.",
             fun: sort,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "rsort",
             aliases: &[],
             doc: "Sort ranges in selection in reverse order.",
             fun: sort_reverse,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "reflow",
             aliases: &[],
             doc: "Hard-wrap the current selection of lines to a given width.",
             fun: reflow,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "tree-sitter-subtree",
             aliases: &["ts-subtree"],
             doc: "Display tree sitter subtree under cursor, primarily for debugging queries.",
             fun: tree_sitter_subtree,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "config-reload",
             aliases: &[],
             doc: "Refresh user config.",
             fun: refresh_config,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "config-open",
             aliases: &[],
             doc: "Open the user config.toml file.",
             fun: open_config,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "log-open",
             aliases: &[],
             doc: "Open the helix log file.",
             fun: open_log,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "insert-output",
             aliases: &[],
             doc: "Run shell command, inserting output before each selection.",
             fun: insert_output,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "append-output",
             aliases: &[],
             doc: "Run shell command, appending output after each selection.",
             fun: append_output,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "pipe",
             aliases: &[],
             doc: "Pipe each selection to the shell command.",
             fun: pipe,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "pipe-to",
             aliases: &[],
             doc: "Pipe each selection to the shell command, ignoring output.",
             fun: pipe_to,
-            completer: None,
+            signature: CommandSignature::none(),
         },
         TypableCommand {
             name: "run-shell-command",
             aliases: &["sh"],
             doc: "Run a shell command",
             fun: run_shell_command,
-            completer: Some(completers::filename),
+            signature: CommandSignature::all(completers::filename)
         },
        TypableCommand {
             name: "reset-diff-change",
@@ -2665,8 +2710,6 @@ pub static TYPABLE_COMMAND_MAP: Lazy<HashMap<&'static str, &'static TypableComma
 
 #[allow(clippy::unnecessary_unwrap)]
 pub(super) fn command_mode(cx: &mut Context) {
-    use shellwords::Shellwords;
-
     let mut prompt = Prompt::new(
         ":".into(),
         Some(':'),
@@ -2705,10 +2748,11 @@ pub(super) fn command_mode(cx: &mut Context) {
                     )
                 };
 
-                if let Some(typed::TypableCommand {
-                    completer: Some(completer),
-                    ..
-                }) = typed::TYPABLE_COMMAND_MAP.get(&words[0] as &str)
+                let argument_number = argument_number_of(&shellwords);
+
+                if let Some(completer) = TYPABLE_COMMAND_MAP
+                    .get(&words[0] as &str)
+                    .map(|tc| tc.completer_for_argument_number(argument_number))
                 {
                     completer(editor, part)
                         .into_iter()
@@ -2773,3 +2817,29 @@ pub(super) fn command_mode(cx: &mut Context) {
     prompt.recalculate_completion(cx.editor);
     cx.push_layer(Box::new(prompt));
 }
+
+fn argument_number_of(shellwords: &Shellwords) -> usize {
+    if shellwords.ends_with_whitespace() {
+        shellwords.words().len().saturating_sub(1)
+    } else {
+        shellwords.words().len().saturating_sub(2)
+    }
+}
+
+#[test]
+fn test_argument_number_of() {
+    let cases = vec![
+        ("set-option", 0),
+        ("set-option ", 0),
+        ("set-option a", 0),
+        ("set-option asdf", 0),
+        ("set-option asdf ", 1),
+        ("set-option asdf xyz", 1),
+        ("set-option asdf xyz abc", 2),
+        ("set-option asdf xyz abc ", 3),
+    ];
+
+    for case in cases {
+        assert_eq!(case.1, argument_number_of(&Shellwords::from(case.0)));
+    }
+}