completion: pad empty argument at $_CLAP_COMPLETE_INDEX

This is a workaround for command name completion. On zsh, the complete position
is specified by $_CLAP_COMPLETE_INDEX, and an empty argument isn't padded. So,
for "jj <TAB>", { args = ["jj"], index = 1 } is provided, and our expand_args()
helpfully fills in the default command "log". Since args[index] is now "log",
no other commands are listed.

cli/src/cli_util.rs

@@ -24,6 +24,7 @@ use std::fmt::Debug;
 use std::fs;
 use std::io;
 use std::io::Write as _;
+use std::iter;
 use std::mem;
 use std::path::Path;
 use std::path::PathBuf;
@@ -3248,12 +3249,27 @@ fn handle_shell_completion(
     // without any changes. They are usually "jj --".
     args.extend(env::args_os().take(2));
 
-    if env::args_os().nth(2).is_some() {
-        // Make sure aliases are expanded before passing them to
-        // clap_complete. We skip the first two args ("jj" and "--") for
-        // alias resolution, then we stitch the args back together, like
-        // clap_complete expects them.
-        let resolved_aliases = expand_args(ui, app, env::args_os().skip(2), config)?;
+    // Make sure aliases are expanded before passing them to clap_complete. We
+    // skip the first two args ("jj" and "--") for alias resolution, then we
+    // stitch the args back together, like clap_complete expects them.
+    let orig_args = env::args_os().skip(2);
+    if orig_args.len() > 0 {
+        let arg_index: Option<usize> = env::var("_CLAP_COMPLETE_INDEX")
+            .ok()
+            .and_then(|s| s.parse().ok());
+        let resolved_aliases = if let Some(index) = arg_index {
+            // As of clap_complete 4.5.38, zsh completion script doesn't pad an
+            // empty arg at the complete position. If the args doesn't include a
+            // command name, the default command would be expanded at that
+            // position. Therefore, no other command names would be suggested.
+            // TODO: Maybe we should instead expand args[..index] + [""], adjust
+            // the index accordingly, strip the last "", and append remainder?
+            let pad_len = usize::saturating_sub(index + 1, orig_args.len());
+            let padded_args = orig_args.chain(iter::repeat(OsString::new()).take(pad_len));
+            expand_args(ui, app, padded_args, config)?
+        } else {
+            expand_args(ui, app, orig_args, config)?
+        };
         args.extend(resolved_aliases.into_iter().map(OsString::from));
     }
     let ran_completion = clap_complete::CompleteEnv::with_factory(|| {

cli/tests/test_completion.rs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use itertools::Itertools as _;
+
+use crate::common::get_stdout_string;
 use crate::common::TestEnvironment;
 
 #[test]
@@ -207,6 +210,39 @@ fn test_completions_are_generated() {
     assert!(stdout.starts_with("complete --keep-order --exclusive --command jj --arguments"));
 }
 
+#[test]
+fn test_zsh_completion() {
+    let mut test_env = TestEnvironment::default();
+    test_env.add_env_var("COMPLETE", "zsh");
+
+    // ["--", "jj"]
+    //        ^^^^ index = 0
+    let complete_at = |index: usize, args: &[&str]| {
+        let assert = test_env
+            .jj_cmd(test_env.env_root(), args)
+            .env("_CLAP_COMPLETE_INDEX", index.to_string())
+            .assert()
+            .success();
+        get_stdout_string(&assert)
+    };
+
+    // Command names should be suggested. If the default command were expanded,
+    // only "log" would be listed.
+    let stdout = complete_at(1, &["--", "jj"]);
+    insta::assert_snapshot!(stdout.lines().take(2).join("\n"), @r"
+    abandon:Abandon a revision
+    absorb:Move changes from a revision into the stack of mutable revisions
+    ");
+    let stdout = complete_at(2, &["--", "jj", "--no-pager"]);
+    insta::assert_snapshot!(stdout.lines().take(2).join("\n"), @r"
+    abandon:Abandon a revision
+    absorb:Move changes from a revision into the stack of mutable revisions
+    ");
+
+    let stdout = complete_at(1, &["--", "jj", "b"]);
+    insta::assert_snapshot!(stdout, @"bookmark:Manage bookmarks [default alias: b]");
+}
+
 #[test]
 fn test_remote_names() {
     let mut test_env = TestEnvironment::default();
@@ -344,6 +380,16 @@ fn test_revisions() {
     k	(no description set)
     r	mutable
     ");
+
+    // complete args of the default command
+    test_env.add_config("ui.default-command = 'log'");
+    let stdout = test_env.jj_cmd_success(&repo_path, &["--", "jj", "-r", ""]);
+    insta::assert_snapshot!(stdout, @r"
+    k	(no description set)
+    r	mutable
+    q	immutable
+    z	(no description set)
+    ");
 }
 
 #[test]