use std::sync::Arc;
use std::time::Instant;
use criterion::BatchSize;
use criterion::BenchmarkGroup;
use criterion::BenchmarkId;
use criterion::measurement::Measurement;
use futures::StreamExt as _;
use jj_lib::revset::SymbolResolver;
use jj_lib::revset::SymbolResolverExtension;
use jj_lib::revset::UserRevsetExpression;
use pollster::FutureExt as _;
use super::CriterionArgs;
use super::new_criterion;
use crate::cli_util::CommandHelper;
use crate::cli_util::RevisionArg;
use crate::cli_util::WorkspaceCommandHelper;
use crate::command_error::CommandError;
use crate::ui::Ui;
#[derive(clap::Args, Clone, Debug)]
#[command(group(clap::ArgGroup::new("revset_source").required(true)))]
pub struct BenchRevsetArgs {
#[arg(group = "revset_source")]
revisions: Vec<RevisionArg>,
#[arg(long, short = 'f', group = "revset_source", value_hint = clap::ValueHint::FilePath)]
file: Option<String>,
#[command(flatten)]
criterion: CriterionArgs,
}
pub async fn cmd_bench_revset(
ui: &mut Ui,
command: &CommandHelper,
args: &BenchRevsetArgs,
) -> Result<(), CommandError> {
let workspace_command = command.workspace_helper(ui)?;
let revsets = if let Some(file_path) = &args.file {
std::fs::read_to_string(command.cwd().join(file_path))?
.lines()
.map(|line| line.trim().to_owned())
.filter(|line| !line.is_empty() && !line.starts_with('#'))
.map(RevisionArg::from)
.collect()
} else {
args.revisions.clone()
};
let mut criterion = new_criterion(ui, &args.criterion);
let mut group = criterion.benchmark_group("revsets");
for revset in &revsets {
bench_revset(ui, command, &workspace_command, &mut group, revset)?;
}
group.finish();
criterion.final_summary();
Ok(())
}
fn bench_revset<M: Measurement>(
ui: &mut Ui,
command: &CommandHelper,
workspace_command: &WorkspaceCommandHelper,
group: &mut BenchmarkGroup<M>,
revset: &RevisionArg,
) -> Result<(), CommandError> {
writeln!(ui.status(), "----------Testing revset: {revset}----------")?;
let expression = workspace_command
.parse_revset(ui, revset)?
.expression()
.clone();
let routine = |workspace_command: &WorkspaceCommandHelper,
expression: Arc<UserRevsetExpression>| {
let repo = workspace_command.repo().as_ref();
let symbol_resolver =
SymbolResolver::new(repo, &([] as [Box<dyn SymbolResolverExtension>; 0]));
let resolved = expression
.resolve_user_expression(repo, &symbol_resolver)
.unwrap();
let revset = resolved.evaluate(repo).unwrap();
revset.stream().count().block_on()
};
let before = Instant::now();
let result = routine(workspace_command, expression.clone());
let after = Instant::now();
writeln!(
ui.status(),
"First run took {:?} and produced {result} commits",
after.duration_since(before),
)?;
group.bench_with_input(
BenchmarkId::from_parameter(revset),
&expression,
|bencher, expression| {
bencher.iter_batched(
|| {
let workspace_command = command.workspace_helper_no_snapshot(ui).unwrap();
workspace_command.repo().readonly_index();
workspace_command
},
|workspace_command| routine(&workspace_command, expression.clone()),
BatchSize::LargeInput,
);
},
);
Ok(())
}