cli: revive jj bench command

This just backs out commit 18c0b97d9d without making any changes,
except for resolving conflicts.

I want a way to benchmark different revsets on e.g. the Git Core repo
or the Linux repo.

Cargo.lock

@@ -785,6 +785,7 @@ dependencies = [
  "clap_complete",
  "clap_mangen",
  "config",
+ "criterion",
  "crossterm",
  "dirs",
  "esl01-renderdag",

Cargo.toml

@@ -35,6 +35,7 @@ clap = { version = "4.1.9", features = ["derive", "deprecated"] }
 clap_complete = "4.1.5"
 clap_mangen = "0.2.10"
 config = { version = "0.13.3", default-features = false, features = ["toml"] }
+criterion = "0.4.0"
 crossterm = { version = "0.26", default-features = false }
 dirs = "4.0.0"
 esl01-renderdag = "0.3.0"

src/commands/mod.rs

@@ -21,16 +21,19 @@ use std::fmt::Debug;
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Instant;
 use std::{fs, io};
 
 use clap::builder::NonEmptyStringValueParser;
 use clap::{ArgGroup, ArgMatches, Command, CommandFactory, FromArgMatches, Subcommand};
+use criterion::Criterion;
 use indexmap::{IndexMap, IndexSet};
 use itertools::Itertools;
 use jujutsu_lib::backend::{CommitId, ObjectId, TreeValue};
 use jujutsu_lib::commit::Commit;
 use jujutsu_lib::dag_walk::topo_order_reverse;
 use jujutsu_lib::default_index_store::{DefaultIndexStore, ReadonlyIndexWrapper};
+use jujutsu_lib::index::HexPrefix;
 use jujutsu_lib::matchers::EverythingMatcher;
 use jujutsu_lib::op_store::{RefTarget, WorkspaceId};
 use jujutsu_lib::repo::{ReadonlyRepo, Repo};
@@ -65,6 +68,8 @@ enum Commands {
     Abandon(AbandonArgs),
     Backout(BackoutArgs),
     #[command(subcommand)]
+    Bench(BenchCommands),
+    #[command(subcommand)]
     Branch(branch::BranchSubcommand),
     #[command(alias = "print")]
     Cat(CatArgs),
@@ -923,6 +928,47 @@ struct SupportMangenArgs {}
 #[derive(clap::Args, Clone, Debug)]
 struct SupportConfigSchemaArgs {}
 
+/// Commands for benchmarking internal operations
+#[derive(Subcommand, Clone, Debug)]
+enum BenchCommands {
+    #[command(name = "commonancestors")]
+    CommonAncestors(BenchCommonAncestorsArgs),
+    #[command(name = "isancestor")]
+    IsAncestor(BenchIsAncestorArgs),
+    #[command(name = "walkrevs")]
+    WalkRevs(BenchWalkRevsArgs),
+    #[command(name = "resolveprefix")]
+    ResolvePrefix(BenchResolvePrefixArgs),
+}
+
+/// Find the common ancestor(s) of a set of commits
+#[derive(clap::Args, Clone, Debug)]
+struct BenchCommonAncestorsArgs {
+    revision1: String,
+    revision2: String,
+}
+
+/// Checks if the first commit is an ancestor of the second commit
+#[derive(clap::Args, Clone, Debug)]
+struct BenchIsAncestorArgs {
+    ancestor: String,
+    descendant: String,
+}
+
+/// Walk revisions that are ancestors of the second argument but not ancestors
+/// of the first
+#[derive(clap::Args, Clone, Debug)]
+struct BenchWalkRevsArgs {
+    unwanted: String,
+    wanted: String,
+}
+
+/// Resolve a commit ID prefix
+#[derive(clap::Args, Clone, Debug)]
+struct BenchResolvePrefixArgs {
+    prefix: String,
+}
+
 /// Low-level commands not intended for users
 #[derive(Subcommand, Clone, Debug)]
 #[command(hide = true)]
@@ -3081,6 +3127,100 @@ fn cmd_support(
     Ok(())
 }
 
+fn run_bench<R, O>(ui: &mut Ui, id: &str, mut routine: R) -> io::Result<()>
+where
+    R: (FnMut() -> O) + Copy,
+    O: Debug,
+{
+    let mut criterion = Criterion::default();
+    let before = Instant::now();
+    let result = routine();
+    let after = Instant::now();
+    writeln!(
+        ui,
+        "First run took {:?} and produced: {:?}",
+        after.duration_since(before),
+        result
+    )?;
+    criterion.bench_function(id, |bencher: &mut criterion::Bencher| {
+        bencher.iter(routine);
+    });
+    Ok(())
+}
+
+fn cmd_bench(
+    ui: &mut Ui,
+    command: &CommandHelper,
+    subcommand: &BenchCommands,
+) -> Result<(), CommandError> {
+    match subcommand {
+        BenchCommands::CommonAncestors(command_matches) => {
+            let workspace_command = command.workspace_helper(ui)?;
+            let commit1 = workspace_command.resolve_single_rev(&command_matches.revision1)?;
+            let commit2 = workspace_command.resolve_single_rev(&command_matches.revision2)?;
+            let index = workspace_command.repo().index();
+            let routine =
+                || index.common_ancestors(&[commit1.id().clone()], &[commit2.id().clone()]);
+            run_bench(
+                ui,
+                &format!(
+                    "commonancestors-{}-{}",
+                    &command_matches.revision1, &command_matches.revision2
+                ),
+                routine,
+            )?;
+        }
+        BenchCommands::IsAncestor(command_matches) => {
+            let workspace_command = command.workspace_helper(ui)?;
+            let ancestor_commit =
+                workspace_command.resolve_single_rev(&command_matches.ancestor)?;
+            let descendant_commit =
+                workspace_command.resolve_single_rev(&command_matches.descendant)?;
+            let index = workspace_command.repo().index();
+            let routine = || index.is_ancestor(ancestor_commit.id(), descendant_commit.id());
+            run_bench(
+                ui,
+                &format!(
+                    "isancestor-{}-{}",
+                    &command_matches.ancestor, &command_matches.descendant
+                ),
+                routine,
+            )?;
+        }
+        BenchCommands::WalkRevs(command_matches) => {
+            let workspace_command = command.workspace_helper(ui)?;
+            let unwanted_commit =
+                workspace_command.resolve_single_rev(&command_matches.unwanted)?;
+            let wanted_commit = workspace_command.resolve_single_rev(&command_matches.wanted)?;
+            let index = workspace_command.repo().index();
+            let routine = || {
+                index
+                    .walk_revs(
+                        &[wanted_commit.id().clone()],
+                        &[unwanted_commit.id().clone()],
+                    )
+                    .count()
+            };
+            run_bench(
+                ui,
+                &format!(
+                    "walkrevs-{}-{}",
+                    &command_matches.unwanted, &command_matches.wanted
+                ),
+                routine,
+            )?;
+        }
+        BenchCommands::ResolvePrefix(command_matches) => {
+            let workspace_command = command.workspace_helper(ui)?;
+            let prefix = HexPrefix::new(&command_matches.prefix).unwrap();
+            let index = workspace_command.repo().index();
+            let routine = || index.resolve_prefix(&prefix);
+            run_bench(ui, &format!("resolveprefix-{}", prefix.hex()), routine)?;
+        }
+    }
+    Ok(())
+}
+
 fn cmd_debug(
     ui: &mut Ui,
     command: &CommandHelper,
@@ -3464,6 +3604,7 @@ pub fn run_command(
         Commands::Sparse(sub_args) => cmd_sparse(ui, command_helper, sub_args),
         Commands::Git(sub_args) => git::cmd_git(ui, command_helper, sub_args),
         Commands::Support(sub_args) => cmd_support(ui, command_helper, sub_args),
+        Commands::Bench(sub_args) => cmd_bench(ui, command_helper, sub_args),
         Commands::Debug(sub_args) => cmd_debug(ui, command_helper, sub_args),
     }
 }