revset: support extensions of symbol resolution

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;
@@ -73,6 +76,45 @@ impl MostDigitsInId {
     }
 }
 
+struct TheDigitestResolver<'a> {
+    repo: &'a dyn Repo,
+    cache: MostDigitsInId,
+}
+
+impl<'a> TheDigitestResolver<'a> {
+    fn new(repo: &'a dyn Repo) -> Self {
+        Self {
+            repo,
+            cache: MostDigitsInId::new(),
+        }
+    }
+}
+
+impl<'a> PartialSymbolResolver for TheDigitestResolver<'a> {
+    fn resolve_symbol(&self, symbol: &str) -> Result<Option<Vec<CommitId>>, RevsetResolutionError> {
+        if symbol != "thedigitest" {
+            return Ok(None);
+        }
+
+        Ok(Some(
+            RevsetExpression::all()
+                .evaluate_programmatic(self.repo)
+                .map_err(|err| RevsetResolutionError::Other(err.into()))?
+                .iter()
+                .filter(|id| num_digits_in_id(id) == self.cache.count(self.repo))
+                .collect_vec(),
+        ))
+    }
+}
+
+struct TheDigitest;
+
+impl SymbolResolverExtension for TheDigitest {
+    fn get_symbol_resolver<'a>(&self, repo: &'a dyn Repo) -> Box<dyn PartialSymbolResolver + 'a> {
+        Box::new(TheDigitestResolver::new(repo))
+    }
+}
+
 impl CommitTemplateLanguageExtension for HexCounter {
     fn build_fn_table<'repo>(&self) -> CommitTemplateBuildFnTable<'repo> {
         type L<'repo> = CommitTemplateLanguage<'repo>;
@@ -132,5 +174,6 @@ impl CommitTemplateLanguageExtension for HexCounter {
 fn main() -> std::process::ExitCode {
     CliRunner::init()
         .add_commit_template_extension(Box::new(HexCounter))
+        .add_symbol_resolver(Box::new(TheDigitest))
         .run()
 }

cli/src/cli_util.rs

@@ -55,7 +55,7 @@ use jj_lib::repo::{
 use jj_lib::repo_path::{FsPathParseError, RepoPath, RepoPathBuf};
 use jj_lib::revset::{
     RevsetAliasesMap, RevsetExpression, RevsetFilterPredicate, RevsetIteratorExt, RevsetModifier,
-    RevsetParseContext, RevsetWorkspaceContext,
+    RevsetParseContext, RevsetWorkspaceContext, SymbolResolverExtension,
 };
 use jj_lib::rewrite::restore_tree;
 use jj_lib::settings::{ConfigResultExt as _, UserSettings};
@@ -91,7 +91,7 @@ use crate::git_util::{
 };
 use crate::merge_tools::{DiffEditor, MergeEditor, MergeToolConfigError};
 use crate::operation_templater::OperationTemplateLanguageExtension;
-use crate::revset_util::RevsetExpressionEvaluator;
+use crate::revset_util::{RevsetExpressionEvaluator, RevsetExtensions};
 use crate::template_builder::TemplateLanguage;
 use crate::template_parser::TemplateAliasesMap;
 use crate::templater::{PropertyPlaceholder, TemplateRenderer};
@@ -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
@@ -402,6 +407,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>>,
@@ -434,6 +440,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,
@@ -877,6 +884,7 @@ impl WorkspaceCommandHelper {
     ) -> Result<RevsetExpressionEvaluator<'_>, CommandError> {
         Ok(RevsetExpressionEvaluator::new(
             self.repo().as_ref(),
+            self.revset_extensions.clone(),
             self.id_prefix_context()?,
             expression,
         ))
@@ -897,7 +905,8 @@ impl WorkspaceCommandHelper {
 
     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.symbol_resolvers().to_vec());
             let revset_string: String = self
                 .settings
                 .config()
@@ -957,6 +966,7 @@ impl WorkspaceCommandHelper {
             self.workspace_id(),
             self.revset_parse_context(),
             self.id_prefix_context()?,
+            self.revset_extensions.clone(),
             &self.commit_template_extensions,
         ))
     }
@@ -1431,12 +1441,14 @@ 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.symbol_resolvers().to_vec());
         let language = CommitTemplateLanguage::new(
             self.tx.repo(),
             self.helper.workspace_id(),
             self.helper.revset_parse_context(),
             &id_prefix_context,
+            self.helper.revset_extensions.clone(),
             &self.helper.commit_template_extensions,
         );
         let template = self
@@ -2506,6 +2518,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,
@@ -2529,6 +2542,7 @@ impl CliRunner {
             extra_configs: vec![],
             store_factories: StoreFactories::default(),
             working_copy_factories: default_working_copy_factories(),
+            revset_extensions: RevsetExtensions::default(),
             commit_template_extensions: vec![],
             operation_template_extensions: vec![],
             dispatch_fn: Box::new(crate::commands::run_command),
@@ -2564,6 +2578,14 @@ impl CliRunner {
         self
     }
 
+    pub fn add_symbol_resolver(
+        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>,
@@ -2698,6 +2720,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

@@ -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, command.revset_extensions().symbol_resolvers());
         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

@@ -17,6 +17,7 @@ use std::cmp::max;
 use std::collections::HashMap;
 use std::io;
 use std::rc::Rc;
+use std::sync::Arc;
 
 use itertools::Itertools as _;
 use jj_lib::backend::{ChangeId, CommitId};
@@ -31,6 +32,7 @@ use jj_lib::revset::{Revset, RevsetParseContext};
 use jj_lib::{git, rewrite};
 use once_cell::unsync::OnceCell;
 
+use crate::revset_util::RevsetExtensions;
 use crate::template_builder::{
     self, merge_fn_map, BuildContext, CoreTemplateBuildFnTable, CoreTemplatePropertyKind,
     IntoTemplateProperty, TemplateBuildMethodFnMap, TemplateLanguage,
@@ -60,6 +62,7 @@ pub struct CommitTemplateLanguage<'repo> {
     id_prefix_context: &'repo IdPrefixContext,
     build_fn_table: CommitTemplateBuildFnTable<'repo>,
     keyword_cache: CommitKeywordCache,
+    revset_extensions: Arc<RevsetExtensions>,
     cache_extensions: ExtensionsMap,
 }
 
@@ -71,12 +74,13 @@ impl<'repo> CommitTemplateLanguage<'repo> {
         workspace_id: &WorkspaceId,
         revset_parse_context: RevsetParseContext<'repo>,
         id_prefix_context: &'repo IdPrefixContext,
-        extensions: &[impl AsRef<dyn CommitTemplateLanguageExtension>],
+        revset_extensions: Arc<RevsetExtensions>,
+        template_extensions: &[impl AsRef<dyn CommitTemplateLanguageExtension>],
     ) -> Self {
         let mut build_fn_table = CommitTemplateBuildFnTable::builtin();
         let mut cache_extensions = ExtensionsMap::empty();
 
-        for extension in extensions {
+        for extension in template_extensions {
             build_fn_table.merge(extension.as_ref().build_fn_table());
             extension
                 .as_ref()
@@ -91,6 +95,7 @@ impl<'repo> CommitTemplateLanguage<'repo> {
             build_fn_table,
             keyword_cache: CommitKeywordCache::default(),
             cache_extensions,
+            revset_extensions,
         }
     }
 }
@@ -645,7 +650,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_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};
@@ -24,7 +25,7 @@ use jj_lib::repo::Repo;
 use jj_lib::revset::{
     self, DefaultSymbolResolver, Revset, RevsetAliasesMap, RevsetCommitRef, RevsetEvaluationError,
     RevsetExpression, RevsetIteratorExt as _, RevsetParseContext, RevsetParseError,
-    RevsetResolutionError,
+    RevsetResolutionError, SymbolResolverExtension,
 };
 use jj_lib::settings::ConfigResultExt as _;
 use thiserror::Error;
@@ -45,21 +46,40 @@ pub enum UserRevsetEvaluationError {
     Evaluation(RevsetEvaluationError),
 }
 
+#[derive(Default)]
+pub struct RevsetExtensions {
+    symbol_resolvers: Vec<Arc<dyn SymbolResolverExtension>>,
+    // TODO: Add more fields for extending the revset language
+}
+
+impl RevsetExtensions {
+    pub fn symbol_resolvers(&self) -> &[Arc<dyn SymbolResolverExtension>] {
+        &self.symbol_resolvers
+    }
+
+    pub fn add_symbol_resolver(&mut self, symbol_resolver: Box<dyn SymbolResolverExtension>) {
+        self.symbol_resolvers.push(symbol_resolver.into())
+    }
+}
+
 /// 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 +97,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,13 +181,14 @@ pub fn evaluate<'a>(
 /// `evaluate()`.
 pub fn default_symbol_resolver<'a>(
     repo: &'a dyn Repo,
+    extensions: &[Arc<dyn SymbolResolverExtension>],
     id_prefix_context: &'a IdPrefixContext,
 ) -> DefaultSymbolResolver<'a> {
     let commit_id_resolver: revset::PrefixResolver<CommitId> =
         Box::new(|repo, prefix| id_prefix_context.resolve_commit_prefix(repo, prefix));
     let change_id_resolver: revset::PrefixResolver<Vec<CommitId>> =
         Box::new(|repo, prefix| id_prefix_context.resolve_change_prefix(repo, prefix));
-    DefaultSymbolResolver::new(repo)
+    DefaultSymbolResolver::new(repo, extensions)
         .with_commit_id_resolver(commit_id_resolver)
         .with_change_id_resolver(change_id_resolver)
 }