revset: add a SymbolResolverExtension trait to provide custom resolvers

cli/examples/custom-commit-templater/main.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use itertools::Itertools;
 use jj_cli::cli_util::CliRunner;
 use jj_cli::commit_templater::{
     CommitTemplateBuildFnTable, CommitTemplateLanguage, CommitTemplateLanguageExtension,
@@ -24,7 +25,9 @@ use jj_lib::commit::Commit;
 use jj_lib::extensions_map::ExtensionsMap;
 use jj_lib::object_id::ObjectId;
 use jj_lib::repo::Repo;
-use jj_lib::revset::RevsetExpression;
+use jj_lib::revset::{
+    PartialSymbolResolver, RevsetExpression, RevsetResolutionError, SymbolResolverExtension,
+};
 use once_cell::sync::OnceCell;
 
 struct HexCounter;
@@ -68,6 +71,40 @@ impl MostDigitsInId {
     }
 }
 
+#[derive(Default)]
+struct TheDigitestResolver {
+    cache: MostDigitsInId,
+}
+
+impl PartialSymbolResolver for TheDigitestResolver {
+    fn resolve_symbol(
+        &self,
+        repo: &dyn Repo,
+        symbol: &str,
+    ) -> Result<Option<Vec<CommitId>>, RevsetResolutionError> {
+        if symbol != "thedigitest" {
+            return Ok(None);
+        }
+
+        Ok(Some(
+            RevsetExpression::all()
+                .evaluate_programmatic(repo)
+                .map_err(|err| RevsetResolutionError::Other(err.into()))?
+                .iter()
+                .filter(|id| num_digits_in_id(id) == self.cache.count(repo))
+                .collect_vec(),
+        ))
+    }
+}
+
+struct TheDigitest;
+
+impl SymbolResolverExtension for TheDigitest {
+    fn new_resolvers<'a>(&self, _repo: &'a dyn Repo) -> Vec<Box<dyn PartialSymbolResolver + 'a>> {
+        vec![Box::<TheDigitestResolver>::default()]
+    }
+}
+
 impl CommitTemplateLanguageExtension for HexCounter {
     fn build_fn_table<'repo>(&self) -> CommitTemplateBuildFnTable<'repo> {
         type L<'repo> = CommitTemplateLanguage<'repo>;
@@ -126,6 +163,7 @@ impl CommitTemplateLanguageExtension for HexCounter {
 
 fn main() -> std::process::ExitCode {
     CliRunner::init()
+        .add_symbol_resolver_extension(Box::new(TheDigitest))
         .add_commit_template_extension(Box::new(HexCounter))
         .run()
 }

cli/src/cli_util.rs

@@ -54,8 +54,8 @@ use jj_lib::repo::{
 };
 use jj_lib::repo_path::{FsPathParseError, RepoPath, RepoPathBuf};
 use jj_lib::revset::{
-    RevsetAliasesMap, RevsetExpression, RevsetFilterPredicate, RevsetIteratorExt, RevsetModifier,
-    RevsetParseContext, RevsetWorkspaceContext,
+    RevsetAliasesMap, RevsetExpression, RevsetExtensions, RevsetFilterPredicate, RevsetIteratorExt,
+    RevsetModifier, RevsetParseContext, RevsetWorkspaceContext, SymbolResolverExtension,
 };
 use jj_lib::rewrite::restore_tree;
 use jj_lib::settings::{ConfigResultExt as _, UserSettings};
@@ -193,6 +193,7 @@ pub struct CommandHelper {
     global_args: GlobalArgs,
     settings: UserSettings,
     layered_configs: LayeredConfigs,
+    revset_extensions: Arc<RevsetExtensions>,
     commit_template_extensions: Vec<Arc<dyn CommitTemplateLanguageExtension>>,
     operation_template_extensions: Vec<Arc<dyn OperationTemplateLanguageExtension>>,
     maybe_workspace_loader: Result<WorkspaceLoader, CommandError>,
@@ -236,6 +237,10 @@ impl CommandHelper {
         self.layered_configs.resolved_config_values(prefix)
     }
 
+    pub fn revset_extensions(&self) -> &Arc<RevsetExtensions> {
+        &self.revset_extensions
+    }
+
     /// Loads template aliases from the configs.
     ///
     /// For most commands that depend on a loaded repo, you should use
@@ -473,6 +478,7 @@ pub struct WorkspaceCommandHelper {
     settings: UserSettings,
     workspace: Workspace,
     user_repo: ReadonlyUserRepo,
+    revset_extensions: Arc<RevsetExtensions>,
     // TODO: Parsed template can be cached if it doesn't capture 'repo lifetime
     commit_summary_template_text: String,
     commit_template_extensions: Vec<Arc<dyn CommitTemplateLanguageExtension>>,
@@ -505,6 +511,7 @@ impl WorkspaceCommandHelper {
             settings,
             workspace,
             user_repo: ReadonlyUserRepo::new(repo),
+            revset_extensions: command.revset_extensions.clone(),
             commit_summary_template_text,
             commit_template_extensions: command.commit_template_extensions.clone(),
             revset_aliases_map,
@@ -948,6 +955,7 @@ impl WorkspaceCommandHelper {
     ) -> Result<RevsetExpressionEvaluator<'_>, CommandError> {
         Ok(RevsetExpressionEvaluator::new(
             self.repo().as_ref(),
+            self.revset_extensions.clone(),
             self.id_prefix_context()?,
             expression,
         ))
@@ -962,13 +970,14 @@ impl WorkspaceCommandHelper {
         RevsetParseContext {
             aliases_map: &self.revset_aliases_map,
             user_email: self.settings.user_email(),
+            extensions: &self.revset_extensions,
             workspace: Some(workspace_context),
         }
     }
 
     pub fn id_prefix_context(&self) -> Result<&IdPrefixContext, CommandError> {
         self.user_repo.id_prefix_context.get_or_try_init(|| {
-            let mut context: IdPrefixContext = IdPrefixContext::default();
+            let mut context: IdPrefixContext = IdPrefixContext::new(self.revset_extensions.clone());
             let revset_string: String = self
                 .settings
                 .config()
@@ -1551,7 +1560,7 @@ impl WorkspaceCommandTransaction<'_> {
         commit: &Commit,
     ) -> std::io::Result<()> {
         // TODO: Use the disambiguation revset
-        let id_prefix_context = IdPrefixContext::default();
+        let id_prefix_context = IdPrefixContext::new(self.helper.revset_extensions.clone());
         let language = CommitTemplateLanguage::new(
             self.tx.repo(),
             self.helper.workspace_id(),
@@ -2654,6 +2663,7 @@ pub struct CliRunner {
     extra_configs: Vec<config::Config>,
     store_factories: StoreFactories,
     working_copy_factories: WorkingCopyFactories,
+    revset_extensions: RevsetExtensions,
     commit_template_extensions: Vec<Arc<dyn CommitTemplateLanguageExtension>>,
     operation_template_extensions: Vec<Arc<dyn OperationTemplateLanguageExtension>>,
     dispatch_fn: CliDispatchFn,
@@ -2677,6 +2687,7 @@ impl CliRunner {
             extra_configs: vec![],
             store_factories: StoreFactories::default(),
             working_copy_factories: default_working_copy_factories(),
+            revset_extensions: Default::default(),
             commit_template_extensions: vec![],
             operation_template_extensions: vec![],
             dispatch_fn: Box::new(crate::commands::run_command),
@@ -2712,6 +2723,14 @@ impl CliRunner {
         self
     }
 
+    pub fn add_symbol_resolver_extension(
+        mut self,
+        symbol_resolver: Box<dyn SymbolResolverExtension>,
+    ) -> Self {
+        self.revset_extensions.add_symbol_resolver(symbol_resolver);
+        self
+    }
+
     pub fn add_commit_template_extension(
         mut self,
         commit_template_extension: Box<dyn CommitTemplateLanguageExtension>,
@@ -2846,6 +2865,7 @@ impl CliRunner {
             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,

cli/src/commands/bench.rs

@@ -23,7 +23,7 @@ use criterion::measurement::Measurement;
 use criterion::{BatchSize, BenchmarkGroup, BenchmarkId, Criterion};
 use jj_lib::object_id::HexPrefix;
 use jj_lib::repo::Repo;
-use jj_lib::revset::{self, DefaultSymbolResolver, RevsetExpression};
+use jj_lib::revset::{self, DefaultSymbolResolver, RevsetExpression, SymbolResolverExtension};
 
 use crate::cli_util::{CommandHelper, RevisionArg, WorkspaceCommandHelper};
 use crate::command_error::CommandError;
@@ -210,7 +210,8 @@ fn bench_revset<M: Measurement>(
     let routine = |workspace_command: &WorkspaceCommandHelper, expression: Rc<RevsetExpression>| {
         // Evaluate the expression without parsing/evaluating short-prefixes.
         let repo = workspace_command.repo().as_ref();
-        let symbol_resolver = DefaultSymbolResolver::new(repo);
+        let symbol_resolver =
+            DefaultSymbolResolver::new(repo, &([] as [Box<dyn SymbolResolverExtension>; 0]));
         let resolved = expression
             .resolve_user_expression(repo, &symbol_resolver)
             .unwrap();

cli/src/commands/debug.rs

@@ -181,8 +181,11 @@ fn cmd_debug_revset(
     writeln!(ui.stdout(), "{expression:#?}")?;
     writeln!(ui.stdout())?;
 
-    let symbol_resolver =
-        revset_util::default_symbol_resolver(repo, workspace_command.id_prefix_context()?);
+    let symbol_resolver = revset_util::default_symbol_resolver(
+        repo,
+        command.revset_extensions().symbol_resolvers(),
+        workspace_command.id_prefix_context()?,
+    );
     let expression = expression.resolve_user_expression(repo, &symbol_resolver)?;
     writeln!(ui.stdout(), "-- Resolved:")?;
     writeln!(ui.stdout(), "{expression:#?}")?;

cli/src/commit_templater.rs

@@ -678,7 +678,11 @@ fn evaluate_immutable_revset<'repo>(
         .map_err(|err| {
             TemplateParseError::expression("Failed to parse revset", span).with_source(err)
         })?;
-    let symbol_resolver = revset_util::default_symbol_resolver(repo, language.id_prefix_context);
+    let symbol_resolver = revset_util::default_symbol_resolver(
+        repo,
+        language.revset_parse_context.extensions.symbol_resolvers(),
+        language.id_prefix_context,
+    );
     let revset = revset_util::evaluate(repo, &symbol_resolver, expression).map_err(|err| {
         TemplateParseError::expression("Failed to evaluate revset", span).with_source(err)
     })?;

cli/src/revset_util.rs

@@ -15,6 +15,7 @@
 //! Utility for parsing and evaluating user-provided revset expressions.
 
 use std::rc::Rc;
+use std::sync::Arc;
 
 use itertools::Itertools as _;
 use jj_lib::backend::{BackendResult, CommitId};
@@ -23,8 +24,8 @@ use jj_lib::id_prefix::IdPrefixContext;
 use jj_lib::repo::Repo;
 use jj_lib::revset::{
     self, DefaultSymbolResolver, Revset, RevsetAliasesMap, RevsetCommitRef, RevsetEvaluationError,
-    RevsetExpression, RevsetIteratorExt as _, RevsetParseContext, RevsetParseError,
-    RevsetResolutionError,
+    RevsetExpression, RevsetExtensions, RevsetIteratorExt as _, RevsetParseContext,
+    RevsetParseError, RevsetResolutionError, SymbolResolverExtension,
 };
 use jj_lib::settings::ConfigResultExt as _;
 use thiserror::Error;
@@ -48,18 +49,21 @@ pub enum UserRevsetEvaluationError {
 /// Wrapper around `RevsetExpression` to provide convenient methods.
 pub struct RevsetExpressionEvaluator<'repo> {
     repo: &'repo dyn Repo,
+    extensions: Arc<RevsetExtensions>,
     id_prefix_context: &'repo IdPrefixContext,
     expression: Rc<RevsetExpression>,
 }
 
 impl<'repo> RevsetExpressionEvaluator<'repo> {
     pub fn new(
         repo: &'repo dyn Repo,
+        extensions: Arc<RevsetExtensions>,
         id_prefix_context: &'repo IdPrefixContext,
         expression: Rc<RevsetExpression>,
     ) -> Self {
         RevsetExpressionEvaluator {
             repo,
+            extensions,
             id_prefix_context,
             expression,
         }
@@ -77,7 +81,11 @@ impl<'repo> RevsetExpressionEvaluator<'repo> {
 
     /// Evaluates the expression.
     pub fn evaluate(&self) -> Result<Box<dyn Revset + 'repo>, UserRevsetEvaluationError> {
-        let symbol_resolver = default_symbol_resolver(self.repo, self.id_prefix_context);
+        let symbol_resolver = default_symbol_resolver(
+            self.repo,
+            self.extensions.symbol_resolvers(),
+            self.id_prefix_context,
+        );
         evaluate(self.repo, &symbol_resolver, self.expression.clone())
     }
 
@@ -157,9 +165,10 @@ pub fn evaluate<'a>(
 /// `evaluate()`.
 pub fn default_symbol_resolver<'a>(
     repo: &'a dyn Repo,
+    extensions: &[impl AsRef<dyn SymbolResolverExtension>],
     id_prefix_context: &'a IdPrefixContext,
 ) -> DefaultSymbolResolver<'a> {
-    DefaultSymbolResolver::new(repo).with_id_prefix_context(id_prefix_context)
+    DefaultSymbolResolver::new(repo, extensions).with_id_prefix_context(id_prefix_context)
 }
 
 /// Parses user-configured expression defining the immutable set.

lib/src/id_prefix.rs

@@ -17,6 +17,7 @@
 use std::iter;
 use std::marker::PhantomData;
 use std::rc::Rc;
+use std::sync::Arc;
 
 use itertools::Itertools as _;
 use once_cell::unsync::OnceCell;
@@ -25,7 +26,9 @@ use crate::backend::{ChangeId, CommitId};
 use crate::hex_util;
 use crate::object_id::{HexPrefix, ObjectId, PrefixResolution};
 use crate::repo::Repo;
-use crate::revset::{DefaultSymbolResolver, RevsetExpression};
+use crate::revset::{
+    DefaultSymbolResolver, RevsetExpression, RevsetExtensions, SymbolResolverExtension,
+};
 
 struct PrefixDisambiguationError;
 
@@ -41,9 +44,13 @@ struct Indexes {
 }
 
 impl DisambiguationData {
-    fn indexes(&self, repo: &dyn Repo) -> Result<&Indexes, PrefixDisambiguationError> {
+    fn indexes(
+        &self,
+        repo: &dyn Repo,
+        extensions: &[impl AsRef<dyn SymbolResolverExtension>],
+    ) -> Result<&Indexes, PrefixDisambiguationError> {
         self.indexes.get_or_try_init(|| {
-            let symbol_resolver = DefaultSymbolResolver::new(repo);
+            let symbol_resolver = DefaultSymbolResolver::new(repo, extensions);
             let resolved_expression = self
                 .expression
                 .clone()
@@ -95,9 +102,17 @@ impl IdIndexSourceEntry<ChangeId> for &'_ (CommitId, ChangeId) {
 #[derive(Default)]
 pub struct IdPrefixContext {
     disambiguation: Option<DisambiguationData>,
+    extensions: Arc<RevsetExtensions>,
 }
 
 impl IdPrefixContext {
+    pub fn new(extensions: Arc<RevsetExtensions>) -> Self {
+        Self {
+            disambiguation: None,
+            extensions,
+        }
+    }
+
     pub fn disambiguate_within(mut self, expression: Rc<RevsetExpression>) -> Self {
         self.disambiguation = Some(DisambiguationData {
             expression,
@@ -108,9 +123,11 @@ impl IdPrefixContext {
 
     fn disambiguation_indexes(&self, repo: &dyn Repo) -> Option<&Indexes> {
         // TODO: propagate errors instead of treating them as if no revset was specified
-        self.disambiguation
-            .as_ref()
-            .and_then(|disambiguation| disambiguation.indexes(repo).ok())
+        self.disambiguation.as_ref().and_then(|disambiguation| {
+            disambiguation
+                .indexes(repo, self.extensions.symbol_resolvers())
+                .ok()
+        })
     }
 
     /// Resolve an unambiguous commit ID prefix.