WIP add dynamic completions

Cargo.toml

@@ -30,7 +30,7 @@ clap = { version = "4.5.20", features = [
     "wrap_help",
     "string",
 ] }
-clap_complete = "4.5.36"
+clap_complete = { version = "4.5.36", features = ["unstable-dynamic"] }
 clap_complete_nushell = "4.5.4"
 clap-markdown = "0.1.4"
 clap_mangen = "0.2.10"

cli/src/cli_util.rs

@@ -2761,7 +2761,7 @@ pub struct Args {
     pub global_args: GlobalArgs,
 }
 
-#[derive(clap::Args, Clone, Debug)]
+#[derive(clap::Args, Clone, Debug, Default)]
 #[command(next_help_heading = "Global Options")]
 pub struct GlobalArgs {
     /// Path to repository to operate on
@@ -2826,7 +2826,7 @@ pub struct GlobalArgs {
     pub early_args: EarlyArgs,
 }
 
-#[derive(clap::Args, Clone, Debug)]
+#[derive(clap::Args, Clone, Debug, Default)]
 pub struct EarlyArgs {
     /// When to colorize output (always, never, debug, auto)
     #[arg(long, value_name = "WHEN", global = true)]
@@ -3271,11 +3271,10 @@ impl CliRunner {
     }
 
     #[instrument(skip_all)]
-    fn run_internal(
-        self,
-        ui: &mut Ui,
-        mut layered_configs: LayeredConfigs,
-    ) -> Result<(), CommandError> {
+    fn run_internal(self, mut layered_configs: LayeredConfigs) -> Result<(), CommandError> {
+        let mut owned_ui = dyn_completion_state::UI.take().unwrap();
+        let ui = &mut owned_ui;
+
         // `cwd` is canonicalized for consistency with `Workspace::workspace_root()` and
         // to easily compute relative paths between them.
         let cwd = env::current_dir()
@@ -3311,69 +3310,105 @@ impl CliRunner {
         })?;
 
         let string_args = expand_args(ui, &self.app, env::args_os(), &config)?;
+
+        // `command_helper_data` is needed by dynamic completion logic which
+        // must occur before arg parsing. So we instantiate it with a few
+        // default values that must be updated later, once arg parsing is
+        // done. This includes `matches`, `global_args`, `settings` and
+        // `maybe_workspace_loader`.
+        let command_helper_data = CommandHelperData {
+            app: self.app,
+            cwd,
+            string_args,
+            matches: Default::default(),
+            global_args: Default::default(),
+            settings: UserSettings::from_config(config),
+            layered_configs,
+            revset_extensions: self.revset_extensions.into(),
+            commit_template_extensions: self.commit_template_extensions,
+            operation_template_extensions: self.operation_template_extensions,
+            maybe_workspace_loader: maybe_cwd_workspace_loader,
+            store_factories: self.store_factories,
+            working_copy_factories: self.working_copy_factories,
+        };
+
+        // Store state needed by dynamic completion code in global variables.
+        // It will be taken back out afterwards.
+        dyn_completion_state::COMMAND_HELPER.replace(Some(CommandHelper {
+            data: Rc::new(command_helper_data),
+        }));
+        dyn_completion_state::UI.replace(Some(owned_ui));
+
+        // Run dynamic completion code. This will check the `COMPLETE`
+        // environment variable and exit quickly if it is not set. This must be
+        // run before args are parsed, because the args can be invalid at the
+        // time the shell requests completions.
+        clap_complete::CompleteEnv::with_factory(crate::commands::default_app).complete();
+
+        // Store state needed by dynamic completion code in global variables.
+        // It will be taken back out afterwards.
+        let mut owned_ui = dyn_completion_state::UI.take().unwrap();
+        let ui = &mut owned_ui;
+        let mut command_helper_data =
+            Rc::into_inner(dyn_completion_state::COMMAND_HELPER.take().unwrap().data).unwrap();
+
         let (matches, args) = parse_args(
             ui,
-            &self.app,
+            &command_helper_data.app,
             &self.tracing_subscription,
-            &string_args,
-            &mut layered_configs,
+            &command_helper_data.string_args,
+            &mut command_helper_data.layered_configs,
         )
-        .map_err(|err| map_clap_cli_error(err, ui, &layered_configs))?;
+        .map_err(|err| map_clap_cli_error(err, ui, &command_helper_data.layered_configs))?;
         for process_global_args_fn in self.process_global_args_fns {
             process_global_args_fn(ui, &matches)?;
         }
 
-        let maybe_workspace_loader = if let Some(path) = &args.global_args.repository {
+        if let Some(path) = &args.global_args.repository {
             // Invalid -R path is an error. No need to proceed.
             let loader = self
                 .workspace_loader_factory
-                .create(&cwd.join(path))
+                .create(&command_helper_data.cwd.join(path))
                 .map_err(|err| map_workspace_load_error(err, Some(path)))?;
-            layered_configs.read_repo_config(loader.repo_path())?;
-            Ok(loader)
-        } else {
-            maybe_cwd_workspace_loader
+            command_helper_data
+                .layered_configs
+                .read_repo_config(loader.repo_path())?;
+            command_helper_data.maybe_workspace_loader = Ok(loader);
         };
 
         // Apply workspace configs and --config-toml arguments.
-        let config = layered_configs.merge();
+        let config = command_helper_data.layered_configs.merge();
         ui.reset(&config)?;
 
         // If -R is specified, check if the expanded arguments differ. Aliases
         // can also be injected by --config-toml, but that's obviously wrong.
         if args.global_args.repository.is_some() {
-            let new_string_args = expand_args(ui, &self.app, env::args_os(), &config).ok();
-            if new_string_args.as_ref() != Some(&string_args) {
+            let new_string_args =
+                expand_args(ui, &command_helper_data.app, env::args_os(), &config).ok();
+            if new_string_args.as_ref() != Some(&command_helper_data.string_args) {
                 writeln!(
                     ui.warning_default(),
                     "Command aliases cannot be loaded from -R/--repository path"
                 )?;
             }
         }
 
-        let settings = UserSettings::from_config(config);
-        let command_helper_data = CommandHelperData {
-            app: self.app,
-            cwd,
-            string_args,
-            matches,
-            global_args: args.global_args,
-            settings,
-            layered_configs,
-            revset_extensions: self.revset_extensions.into(),
-            commit_template_extensions: self.commit_template_extensions,
-            operation_template_extensions: self.operation_template_extensions,
-            maybe_workspace_loader,
-            store_factories: self.store_factories,
-            working_copy_factories: self.working_copy_factories,
-        };
+        command_helper_data.matches = matches;
+        command_helper_data.global_args = args.global_args;
+        command_helper_data.settings = UserSettings::from_config(config);
+
         let command_helper = CommandHelper {
             data: Rc::new(command_helper_data),
         };
         for start_hook_fn in self.start_hook_fns {
             start_hook_fn(ui, &command_helper)?;
         }
-        (self.dispatch_fn)(ui, &command_helper)
+        (self.dispatch_fn)(ui, &command_helper)?;
+
+        // Put `ui` back into global state, our caller needs it again.
+        dyn_completion_state::UI.set(Some(owned_ui));
+
+        Ok(())
     }
 
     #[must_use]
@@ -3387,15 +3422,41 @@ impl CliRunner {
             .build()
             .unwrap();
         let layered_configs = LayeredConfigs::from_environment(config);
-        let mut ui = Ui::with_config(&layered_configs.merge())
+        let ui = Ui::with_config(&layered_configs.merge())
             .expect("default config should be valid, env vars are stringly typed");
-        let result = self.run_internal(&mut ui, layered_configs);
+
+        // Put `ui` in static variable so it can be access by dynamic completion
+        // logic. This will be triggered in `run_internal`. We can take the `ui`
+        // back out of the static after that.
+        dyn_completion_state::UI.set(Some(ui));
+        let result = self.run_internal(layered_configs);
+        let mut ui = dyn_completion_state::UI.take().unwrap();
+
         let exit_code = handle_command_result(&mut ui, result);
         ui.finalize_pager();
         exit_code
     }
 }
 
+/// This state needs to be accessible for dynamic completion code, which doesn't
+/// accept arguments. We also can't extract the initialization of this into a
+/// separate function (that could be called by dynamic completion code), because
+/// the initialization can depend on user-provided code, most notably Google's
+/// custom storage backend. Global mutable state seems like the only option
+/// left.
+pub mod dyn_completion_state {
+    use std::cell::Cell;
+
+    use super::CommandHelper;
+    use super::Ui;
+
+    // At least command_helper must be thread_local, Ui may not be necessary.
+    thread_local! {
+        pub static UI: Cell<Option<Ui>> = const { Cell::new(None) };
+        pub static COMMAND_HELPER: Cell<Option<CommandHelper>> = const { Cell::new(None) };
+    }
+}
+
 fn map_clap_cli_error(
     mut cmd_err: CommandError,
     ui: &Ui,

cli/src/commands/bookmark/rename.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use clap_complete::ArgValueCompleter;
 use jj_lib::op_store::RefTarget;
 
 use super::has_tracked_remote_bookmarks;
@@ -20,12 +21,15 @@ use crate::command_error::user_error;
 use crate::command_error::CommandError;
 use crate::ui::Ui;
 
+mod complete;
+
 /// Rename `old` bookmark name to `new` bookmark name
 ///
 /// The new bookmark name points at the same commit as the old bookmark name.
 #[derive(clap::Args, Clone, Debug)]
 pub struct BookmarkRenameArgs {
     /// The old name of the bookmark
+    #[arg(add = ArgValueCompleter::new(complete::branch))]
     old: String,
 
     /// The new name of the bookmark

cli/src/commands/bookmark/rename/complete.rs

@@ -0,0 +1,23 @@
+use clap_complete::CompletionCandidate;
+
+use crate::cli_util::dyn_completion_state;
+
+pub fn branch(current: &std::ffi::OsStr) -> Vec<CompletionCandidate> {
+    let Some(current) = current.to_str() else {
+        return Vec::new();
+    };
+
+    let mut ui = dyn_completion_state::UI.take().unwrap();
+    let command = dyn_completion_state::COMMAND_HELPER.take().unwrap();
+    let ui = &mut ui;
+
+    let workspace_command = command.workspace_helper(ui).unwrap();
+    let repo = workspace_command.repo();
+    let view = repo.view();
+
+    view.bookmarks()
+        .map(|(local, _)| local)
+        .filter(|branch| branch.starts_with(current))
+        .map(CompletionCandidate::new)
+        .collect()
+}

docs/install-and-setup.md

@@ -206,6 +206,25 @@ $ jj config set --user user.email "[email protected]"
 
 ## Command-line completion
 
+TODO: update this section if appropriate.
+
+I think we can simply update the output of `jj util completion $SHELL` to be
+- `source <(COMPLETE=bash jj)`'
+- `COMPLETE=fish jj | source`
+etc. instead of the actual script.
+That will honor the recommendation of clap to regenerate and source the hooks on
+every startup of the shell.
+(see the recommendation here: https://docs.rs/clap_complete/latest/clap_complete/env/index.html)
+
+Nushell is not yet supported for dynamic completions, I presume we can simply
+continue to emit the normal, static completions for Nushell.
+
+Another option would be to have some kind of beta-phase for the dynamic
+completions, where we tell users to manually add `COMPLETE=fish jj | source`
+to their shell init files if they want to opt into the beta. If there are no
+issues found, we change `jj util completion` to emit the dynamic completion
+hooks instead.
+
 To set up command-line completion, source the output of
 `jj util completion bash/zsh/fish`. Exactly how to source it
 depends on your shell.

flake.nix

@@ -113,6 +113,7 @@
             $out/bin/jj util mangen > ./jj.1
             installManPage ./jj.1
 
+            # TODO: install dynamic completions differently ?
             installShellCompletion --cmd jj \
               --bash <($out/bin/jj util completion bash) \
               --fish <($out/bin/jj util completion fish) \