Optionally resolve rules by a provided id (pantsbuild#19755)

This is work towards supporting call-by-name for rules, as spec'd
out in pantsbuild#19730.

This change:

1. Introduces a RuleId struct.
2. Adds an optional RuleId-typed field to DependencyKey.
3. Sets a RuleId on every Task and Intrinsic, of the form
    `path.to.module.function`.
4. Validate that RuleIds are unique.
5. During rule graph construction, when considering candidates
    to satisfy a DependencyKey, if the DependencyKey has a RuleId,
    only consider candidates with that RuleId (there should be only one).
6. A simple test to demonstrate 5.

The idea is that the call-by-name syntax will set the
path.to.module.function as the id on the DependencyKey, which
will cause this id to be used to select the rule and be done with it,
instead of the current, complex solving logic.

src/python/pants/backend/build_files/utils.py

@@ -19,10 +19,11 @@ def _get_build_file_partitioner_rules(cls) -> Iterable:
     """Returns the BUILD file partitioner rule."""
 
     @rule(
+        canonical_name_suffix=cls.__name__,
         _param_type_overrides={
             "request": cls.PartitionRequest,
             "subsystem": cls.tool_subsystem,
-        }
+        },
     )
     async def partition_build_files(
         request: FixFilesRequest.PartitionRequest,

src/python/pants/core/goals/run.py

@@ -275,7 +275,7 @@ async def run(
 def _unsupported_debug_adapter_rules(cls: type[RunFieldSet]) -> Iterable:
     """Returns a rule that implements DebugAdapterRequest by raising an error."""
 
-    @rule(_param_type_overrides={"request": cls})
+    @rule(canonical_name_suffix=cls.__name__, _param_type_overrides={"request": cls})
     async def get_run_debug_adapter_request(request: RunFieldSet) -> RunDebugAdapterRequest:
         raise NotImplementedError(
             "Running this target type with a debug adapter is not yet supported."
@@ -297,17 +297,17 @@ def _run_in_sandbox_behavior_rule(cls: type[RunFieldSet]) -> Iterable:
     a `RunInSandboxRequest`.
     """
 
-    @rule(_param_type_overrides={"request": cls})
+    @rule(canonical_name_suffix=cls.__name__, _param_type_overrides={"request": cls})
     async def not_supported(request: RunFieldSet) -> RunInSandboxRequest:
         raise NotImplementedError(
             "Running this target type within the sandbox is not yet supported."
         )
 
-    @rule(_param_type_overrides={"request": cls})
+    @rule(canonical_name_suffix=cls.__name__, _param_type_overrides={"request": cls})
     async def run_request_hermetic(request: RunFieldSet) -> RunInSandboxRequest:
         return await _run_request(request)
 
-    @_uncacheable_rule(_param_type_overrides={"request": cls})
+    @_uncacheable_rule(canonical_name_suffix=cls.__name__, _param_type_overrides={"request": cls})
     async def run_request_not_hermetic(request: RunFieldSet) -> RunInSandboxRequest:
         return await _run_request(request)
 

src/python/pants/core/goals/test.py

@@ -1032,7 +1032,7 @@ async def get_filtered_environment(test_env_aware: TestSubsystem.EnvironmentAwar
 def _unsupported_debug_rules(cls: type[TestRequest]) -> Iterable:
     """Returns a rule that implements TestDebugRequest by raising an error."""
 
-    @rule(_param_type_overrides={"request": cls.Batch})
+    @rule(canonical_name_suffix=cls.__name__, _param_type_overrides={"request": cls.Batch})
     async def get_test_debug_request(request: TestRequest.Batch) -> TestDebugRequest:
         raise NotImplementedError("Testing this target with --debug is not yet supported.")
 
@@ -1043,7 +1043,7 @@ async def get_test_debug_request(request: TestRequest.Batch) -> TestDebugRequest
 def _unsupported_debug_adapter_rules(cls: type[TestRequest]) -> Iterable:
     """Returns a rule that implements TestDebugAdapterRequest by raising an error."""
 
-    @rule(_param_type_overrides={"request": cls.Batch})
+    @rule(canonical_name_suffix=cls.__name__, _param_type_overrides={"request": cls.Batch})
     async def get_test_debug_adapter_request(request: TestRequest.Batch) -> TestDebugAdapterRequest:
         raise NotImplementedError(
             "Testing this target type with a debug adapter is not yet supported."

src/python/pants/core/util_rules/partitions.py

@@ -190,10 +190,11 @@ def _single_partition_field_set_rules(cls) -> Iterable:
     one partition."""
 
     @rule(
+        canonical_name_suffix=cls.__name__,
         _param_type_overrides={
             "request": cls.PartitionRequest,
             "subsystem": cls.tool_subsystem,
-        }
+        },
     )
     async def partitioner(
         request: _PartitionFieldSetsRequestBase, subsystem: SkippableSubsystem
@@ -213,10 +214,11 @@ def _single_partition_file_rules(cls) -> Iterable:
     sources_field_name = _get_sources_field_name(cls.field_set_type)
 
     @rule(
+        canonical_name_suffix=cls.__name__,
         _param_type_overrides={
             "request": cls.PartitionRequest,
             "subsystem": cls.tool_subsystem,
-        }
+        },
     )
     async def partitioner(
         request: _PartitionFieldSetsRequestBase, subsystem: SkippableSubsystem
@@ -246,10 +248,11 @@ def _partition_per_input_field_set_rules(cls) -> Iterable:
     a single-element partition per input."""
 
     @rule(
+        canonical_name_suffix=cls.__name__,
         _param_type_overrides={
             "request": cls.PartitionRequest,
             "subsystem": cls.tool_subsystem,
-        }
+        },
     )
     async def partitioner(
         request: _PartitionFieldSetsRequestBase, subsystem: SkippableSubsystem
@@ -275,10 +278,11 @@ def _partition_per_input_file_rules(cls) -> Iterable:
     sources_field_name = _get_sources_field_name(cls.field_set_type)
 
     @rule(
+        canonical_name_suffix=cls.__name__,
         _param_type_overrides={
             "request": cls.PartitionRequest,
             "subsystem": cls.tool_subsystem,
-        }
+        },
     )
     async def partitioner(
         request: _PartitionFieldSetsRequestBase, subsystem: SkippableSubsystem

src/python/pants/core/util_rules/wrap_source.py

@@ -136,7 +136,10 @@ class WrapSourceTarget(Target):
         )
 
     # need to use `_param_type_overrides` to stop `@rule` from inspecting the function's source
-    @rule(_param_type_overrides={"request": GenerateWrapSourceSourcesRequest})
+    @rule(
+        canonical_name_suffix=source_field_type.__name__,
+        _param_type_overrides={"request": GenerateWrapSourceSourcesRequest},
+    )
     async def wrap_source(request: GenerateSourcesRequest) -> GeneratedSources:
         return await _wrap_source(request)
 

src/python/pants/engine/rules.py

@@ -144,7 +144,7 @@ def _ensure_type_annotation(
     return type_annotation
 
 
-PUBLIC_RULE_DECORATOR_ARGUMENTS = {"canonical_name", "desc", "level"}
+PUBLIC_RULE_DECORATOR_ARGUMENTS = {"canonical_name", "canonical_name_suffix", "desc", "level"}
 # We aren't sure if these'll stick around or be removed at some point, so they are "private"
 # and should only be used in Pants' codebase.
 PRIVATE_RULE_DECORATOR_ARGUMENTS = {
@@ -217,9 +217,18 @@ def rule_decorator(func, **kwargs) -> Callable:
     is_goal_cls = issubclass(return_type, Goal)
 
     # Set a default canonical name if one is not explicitly provided to the module and name of the
-    # function that implements it. This is used as the workunit name.
+    # function that implements it, plus an optional suffix. This is used as the workunit name.
+    # The suffix is a convenient way to disambiguate multiple rules registered dynamically from the
+    # same static code (by overriding the inferred param types in the @rule decorator).
+    # TODO: It is not yet clear how dynamically registered rules whose names are generated
+    #  with a suffix will work in practice with the new call-by-name semantics.
+    #  For now the suffix serves to ensure unique names. Whether they are useful is another matter.
+    suffix = kwargs.get("canonical_name_suffix", "")
     effective_name = kwargs.get(
-        "canonical_name", f"{func.__module__}.{func.__qualname__}".replace(".<locals>", "")
+        "canonical_name",
+        f"{func.__module__}.{func.__qualname__}{('_' + suffix) if suffix else ''}".replace(
+            ".<locals>", ""
+        ),
     )
 
     # Set a default description, which is used in the dynamic UI and stacktraces.

src/python/pants/jvm/target_types.py

@@ -826,7 +826,11 @@ def jvm_resolve_field_default_factory(
 def _jvm_source_run_request_rule(cls: type[JvmRunnableSourceFieldSet]) -> Iterable[Rule]:
     from pants.jvm.run import rules as run_rules
 
-    @rule(_param_type_overrides={"request": cls}, level=LogLevel.DEBUG)
+    @rule(
+        canonical_name_suffix=cls.__name__,
+        _param_type_overrides={"request": cls},
+        level=LogLevel.DEBUG,
+    )
     async def jvm_source_run_request(request: JvmRunnableSourceFieldSet) -> RunRequest:
         return await Get(RunRequest, GenericJvmRunRequest(request))
 

src/rust/engine/rule_graph/Cargo.toml

@@ -10,6 +10,7 @@ deepsize = { workspace = true, features = ["internment", "smallvec"] }
 fnv = { workspace = true }
 indexmap = { workspace = true }
 internment = { workspace = true }
+itertools = { workspace = true }
 log = { workspace = true }
 petgraph = { workspace = true }
 smallvec = { version = "1", features = ["union"] }

src/rust/engine/rule_graph/src/builder.rs

@@ -1,7 +1,7 @@
 // Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
 // Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-use crate::rules::{DependencyKey, ParamTypes, Query, Rule};
+use crate::rules::{DependencyKey, ParamTypes, Query, Rule, RuleId};
 use crate::{
   params_str, Entry, EntryWithDeps, Reentry, RootEntry, RuleEdges, RuleEntry, RuleGraph,
 };
@@ -11,6 +11,7 @@ use std::collections::{BTreeMap, VecDeque};
 use fnv::{FnvHashMap as HashMap, FnvHashSet as HashSet};
 use indexmap::IndexSet;
 use internment::Intern;
+use itertools::Itertools;
 use petgraph::graph::{DiGraph, EdgeReference, NodeIndex};
 use petgraph::visit::{DfsPostOrder, EdgeRef, IntoNodeReferences, NodeRef, VisitMap, Visitable};
 use petgraph::Direction;
@@ -256,6 +257,8 @@ impl<R: Rule> Builder<R> {
   }
 
   pub fn graph(self) -> Result<RuleGraph<R>, String> {
+    // 0. validate that the rules all have unique rule ids.
+    self.validate_rule_ids()?;
     // 1. build a polymorphic graph, where nodes might have multiple legal sources of dependencies
     let initial_polymorphic_graph = self.initial_polymorphic();
     // 2. run live variable analysis on the polymorphic graph to gather a conservative (ie, overly
@@ -272,6 +275,29 @@ impl<R: Rule> Builder<R> {
     self.finalize(pruned_edges_graph)
   }
 
+  ///
+  /// Validate that all rules have unique RuleIds.
+  ///
+  fn validate_rule_ids(&self) -> Result<(), String> {
+    let mut invalid_rule_ids: Vec<&RuleId> = self
+      .rules
+      .values()
+      .flatten()
+      .map(|rule| rule.id())
+      .duplicates()
+      .collect();
+    match invalid_rule_ids.len() {
+      0 => Ok(()),
+      _ => {
+        invalid_rule_ids.sort();
+        Err(format!(
+          "The following rule ids were each used by more than one rule: {}",
+          invalid_rule_ids.iter().join(", ")
+        ))
+      }
+    }
+  }
+
   ///
   /// Builds a polymorphic graph while computing an out_set for each node in the graph by accounting
   /// for which `Param`s are available at each use site. During this phase, nodes may have multiple
@@ -308,6 +334,9 @@ impl<R: Rule> Builder<R> {
       })
       .collect::<HashMap<_, _>>();
 
+    let rules_by_id: HashMap<&RuleId, &R> =
+      self.rules.values().flatten().map(|r| (r.id(), r)).collect();
+
     // Rules and Reentries are created on the fly based on the out_set of dependents.
     let mut rules: HashMap<(R, ParamTypes<R::TypeId>), NodeIndex<u32>> = HashMap::default();
     #[allow(clippy::type_complexity)]
@@ -362,16 +391,27 @@ impl<R: Rule> Builder<R> {
           }
 
           let mut candidates = Vec::new();
-          if dependency_key.provided_params.is_empty()
-            && graph[node_id].1.contains(&dependency_key.product())
-            && params.contains_key(&dependency_key.product())
-          {
-            candidates.push(Node::Param(dependency_key.product()));
-          }
+          if let Some(rule_id) = &dependency_key.rule_id {
+            // New call-by-name semantics.
+            candidates.extend(rules_by_id.get(rule_id).map(|&r| Node::Rule(r.clone())));
+            // TODO: Once we are entirely call-by-name, we can get rid of the entire edifice
+            // of multiple candidates and the unsatisfiable_nodes mechanism, and modify this
+            // function to return a Result, which will be Err if there is no rule with a
+            // matching RuleId for some node.
+            assert!(candidates.len() < 2);
+          } else {
+            // Old call-by-type semantics.
+            if dependency_key.provided_params.is_empty()
+              && graph[node_id].1.contains(&dependency_key.product())
+              && params.contains_key(&dependency_key.product())
+            {
+              candidates.push(Node::Param(dependency_key.product()));
+            }
 
-          if let Some(rules) = self.rules.get(&dependency_key.product()) {
-            candidates.extend(rules.iter().map(|r| Node::Rule(r.clone())));
-          };
+            if let Some(rules) = self.rules.get(&dependency_key.product()) {
+              candidates.extend(rules.iter().map(|r| Node::Rule(r.clone())));
+            };
+          }
 
           (dependency_key, candidates)
         })

src/rust/engine/rule_graph/src/lib.rs

@@ -37,7 +37,7 @@ use internment::Intern;
 
 pub use crate::builder::Builder;
 pub use crate::rules::{
-  DependencyKey, DisplayForGraph, DisplayForGraphArgs, ParamTypes, Query, Rule, TypeId,
+  DependencyKey, DisplayForGraph, DisplayForGraphArgs, ParamTypes, Query, Rule, RuleId, TypeId,
 };
 
 #[derive(Eq, Hash, PartialEq, Clone, Debug)]

src/rust/engine/rule_graph/src/rules.rs

@@ -23,10 +23,35 @@ pub trait TypeId:
     I: Iterator<Item = Self>;
 }
 
+// Identifies a specific Rule when called by name.
+// TODO: Turn into a generic type param? For now this is just a newtype, for clarity.
+#[derive(DeepSizeOf, Eq, Hash, PartialEq, Clone, Debug, PartialOrd, Ord)]
+pub struct RuleId(String);
+
+impl RuleId {
+  pub fn new(id: &str) -> Self {
+    Self(id.into())
+  }
+
+  pub fn from_string(s: String) -> Self {
+    Self(s)
+  }
+}
+
+impl Display for RuleId {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    write!(f, "{}", self.0)
+  }
+}
+
 // NB: Most of our expected usecases for multiple-provided-parameters involve two parameters, hence
 // the SmallVec sizing here. See also `Self::provides`.
 #[derive(DeepSizeOf, Eq, Hash, PartialEq, Clone, Debug, PartialOrd, Ord)]
 pub struct DependencyKey<T: TypeId> {
+  // The id of the rule represented by this DependencyKey, if provided at the callsite
+  // (e.g., via call-by-name semantics).
+  pub rule_id: Option<RuleId>,
+
   pub product: T,
   // The param types which are introduced into scope at the callsite ("provided").
   pub provided_params: SmallVec<[T; 2]>,
@@ -41,6 +66,16 @@ pub struct DependencyKey<T: TypeId> {
 impl<T: TypeId> DependencyKey<T> {
   pub fn new(product: T) -> Self {
     DependencyKey {
+      rule_id: None,
+      product,
+      provided_params: SmallVec::default(),
+      in_scope_params: None,
+    }
+  }
+
+  pub fn for_known_rule(rule_id: RuleId, product: T) -> Self {
+    DependencyKey {
+      rule_id: Some(rule_id),
       product,
       provided_params: SmallVec::default(),
       in_scope_params: None,
@@ -121,7 +156,13 @@ impl<T: TypeId> DependencyKey<T> {
 
 impl<T: TypeId> Display for DependencyKey<T> {
   fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-    if self.provided_params.is_empty() {
+    if let Some(rule_id) = &self.rule_id {
+      write!(
+        f,
+        "{}({:?}, ...) -> {}",
+        rule_id, self.provided_params, self.product
+      )
+    } else if self.provided_params.is_empty() {
       write!(f, "{}", self.product)
     } else {
       write!(f, "Get({}, {:?})", self.product, self.provided_params)
@@ -167,6 +208,11 @@ pub trait Rule:
 {
   type TypeId: TypeId;
 
+  ///
+  /// Returns the id of this Rule.
+  ///
+  fn id(&self) -> &RuleId;
+
   ///
   /// Returns the product (output) type for this Rule.
   ///