[internal] Implement @unions via Querys (#12966)

As described in #12934, we would like plugins to have better defined interfaces, and to further clarify what is available while satisfying a `@union`.

Without making any API changes for `@rule` authors, we can take a first step in this direction by implementing the `Get`s for `@union` types using `Query`s instead. The effect of this is that a `@union` usage may _only_ use the explicitly provided `Param`, and no others. When compared to the usual usage of `Get` (which can consume any `Param`s which are in scope at the "call site"), this makes for a much better defined API boundary.

In order to complete #12934, we will likely want to make interface changes to allow more than the single `@union`-member-`Param` specified to the `Get` to be consumed by a plugin (see the examples in that issue's description). But that is not necessary today, and this change also has the benefit of fixing the blocking issue behind #12889 and #12929 by significantly simplifying the rule graph computation (since plugin boundaries are now "hardcoded", and don't need the `Param` detection executed for `Get`s).

[ci skip-build-wheels]

src/python/pants/core/target_types_test.py

@@ -20,6 +20,7 @@
     RelocateFilesViaCodegenRequest,
 )
 from pants.core.target_types import rules as target_type_rules
+from pants.core.util_rules.archive import rules as archive_rules
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.core.util_rules.source_files import rules as source_files_rules
 from pants.engine.addresses import Address
@@ -37,6 +38,7 @@ def test_relocated_files() -> None:
     rule_runner = RuleRunner(
         rules=[
             *target_type_rules(),
+            *archive_rules(),
             *source_files_rules(),
             QueryRule(GeneratedSources, [RelocateFilesViaCodegenRequest]),
             QueryRule(TransitiveTargets, [TransitiveTargetsRequest]),

src/python/pants/engine/internals/native_engine.pyi

@@ -79,6 +79,7 @@ def tasks_task_begin(
 ) -> None: ...
 def tasks_task_end(tasks: PyTasks) -> None: ...
 def tasks_add_get(tasks: PyTasks, output: type, input: type) -> None: ...
+def tasks_add_union(tasks: PyTasks, output_type: type, input_type: tuple[type, ...]) -> None: ...
 def tasks_add_select(tasks: PyTasks, selector: type) -> None: ...
 def tasks_add_query(tasks: PyTasks, output_type: type, input_type: tuple[type, ...]) -> None: ...
 def execution_add_root_select(

src/python/pants/engine/internals/scheduler.py

@@ -659,18 +659,15 @@ def register_task(rule: TaskRule) -> None:
         for selector in rule.input_selectors:
             native_engine.tasks_add_select(tasks, selector)
 
-        def add_get_edge(product: type, subject: type) -> None:
-            native_engine.tasks_add_get(tasks, product, subject)
-
         for the_get in rule.input_gets:
             if is_union(the_get.input_type):
-                # If the registered subject type is a union, add Get edges to all registered
-                # union members.
+                # Register a union. TODO: See #12934: this should involve an explicit interface
+                # soon, rather than one being implicitly created with only the provided Param.
                 for union_member in union_membership.get(the_get.input_type):
-                    add_get_edge(the_get.output_type, union_member)
+                    native_engine.tasks_add_union(tasks, the_get.output_type, (union_member,))
             else:
                 # Otherwise, the Get subject is a "concrete" type, so add a single Get edge.
-                add_get_edge(the_get.output_type, the_get.input_type)
+                native_engine.tasks_add_get(tasks, the_get.output_type, the_get.input_type)
 
         native_engine.tasks_task_end(tasks)
 

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

@@ -189,12 +189,12 @@ type InLabeledGraph<R> =
 ///
 pub struct Builder<R: Rule> {
   rules: BTreeMap<R::TypeId, Vec<R>>,
-  queries: Vec<Query<R>>,
+  queries: IndexSet<Query<R>>,
   params: ParamTypes<R::TypeId>,
 }
 
 impl<R: Rule> Builder<R> {
-  pub fn new(rules: Vec<R>, queries: Vec<Query<R>>) -> Builder<R> {
+  pub fn new(rules: IndexSet<R>, queries: IndexSet<Query<R>>) -> Builder<R> {
     // Group rules by product/return type.
     let mut rules_by_type = BTreeMap::new();
     for rule in rules {
@@ -1228,7 +1228,7 @@ impl<R: Rule> Builder<R> {
     }
 
     Ok(RuleGraph {
-      queries: self.queries,
+      queries: self.queries.into_iter().collect(),
       rule_dependency_edges,
       // TODO
       unreachable_rules: Vec::default(),

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

@@ -33,6 +33,8 @@ mod rules;
 use std::collections::{HashMap, HashSet};
 use std::io;
 
+use indexmap::IndexSet;
+
 pub use crate::builder::Builder;
 pub use crate::rules::{
   DependencyKey, DisplayForGraph, DisplayForGraphArgs, ParamTypes, Query, Rule, TypeId,
@@ -243,7 +245,7 @@ fn entry_with_deps_str<R: Rule>(entry: &EntryWithDeps<R>) -> String {
 }
 
 impl<R: Rule> RuleGraph<R> {
-  pub fn new(rules: Vec<R>, queries: Vec<Query<R>>) -> Result<RuleGraph<R>, String> {
+  pub fn new(rules: IndexSet<R>, queries: IndexSet<Query<R>>) -> Result<RuleGraph<R>, String> {
     Builder::new(rules, queries).graph()
   }
 

src/rust/engine/rule_graph/src/tests.rs

@@ -1,6 +1,9 @@
+use std::fmt;
+
+use indexmap::{indexset, IndexSet};
+
 use crate::builder::combinations_of_one;
 use crate::{Palette, Query, RuleGraph};
-use std::fmt;
 
 #[test]
 fn combinations_of_one_test() {
@@ -29,8 +32,8 @@ fn combinations_of_one_test() {
 
 #[test]
 fn basic() {
-  let rules = vec![Rule("a", "a_from_b", vec![DependencyKey("b", None)])];
-  let queries = vec![Query::new("a", vec!["b"])];
+  let rules = indexset![Rule("a", "a_from_b", vec![DependencyKey("b", None)])];
+  let queries = indexset![Query::new("a", vec!["b"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -39,8 +42,8 @@ fn basic() {
 
 #[test]
 fn singleton() {
-  let rules = vec![Rule("a", "a_singleton", vec![])];
-  let queries = vec![Query::new("a", vec![])];
+  let rules = indexset![Rule("a", "a_singleton", vec![])];
+  let queries = indexset![Query::new("a", vec![])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -49,8 +52,8 @@ fn singleton() {
 
 #[test]
 fn insufficient_query() {
-  let rules = vec![Rule("a", "a_from_b", vec![DependencyKey("b", None)])];
-  let queries = vec![Query::new("a", vec![])];
+  let rules = indexset![Rule("a", "a_from_b", vec![DependencyKey("b", None)])];
+  let queries = indexset![Query::new("a", vec![])];
 
   assert!(RuleGraph::new(rules, queries)
     .err()
@@ -60,8 +63,8 @@ fn insufficient_query() {
 
 #[test]
 fn no_rules() {
-  let rules: Vec<Rule> = vec![];
-  let queries = vec![Query::new("a", vec![])];
+  let rules: IndexSet<Rule> = indexset![];
+  let queries = indexset![Query::new("a", vec![])];
 
   assert!(RuleGraph::new(rules, queries)
     .err()
@@ -71,11 +74,11 @@ fn no_rules() {
 
 #[test]
 fn ambiguity() {
-  let rules = vec![
+  let rules = indexset![
     Rule("a", "a_from_b", vec![DependencyKey("b", None)]),
     Rule("a", "a_from_c", vec![DependencyKey("c", None)]),
   ];
-  let queries = vec![Query::new("a", vec!["b", "c"])];
+  let queries = indexset![Query::new("a", vec!["b", "c"])];
 
   assert!(RuleGraph::new(rules, queries)
     .err()
@@ -85,15 +88,15 @@ fn ambiguity() {
 
 #[test]
 fn nested_single() {
-  let rules = vec![
+  let rules = indexset![
     Rule("a", "a_from_b", vec![DependencyKey("b", Some("c"))]),
     Rule(
       "b",
       "b_from_c",
       vec![DependencyKey("c", None), DependencyKey("d", None)],
     ),
   ];
-  let queries = vec![Query::new("a", vec!["d"])];
+  let queries = indexset![Query::new("a", vec!["d"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -102,7 +105,7 @@ fn nested_single() {
 
 #[test]
 fn nested_multiple() {
-  let rules = vec![
+  let rules = indexset![
     Rule("a", "a_from_b", vec![DependencyKey("b", Some("c"))]),
     Rule(
       "b",
@@ -115,7 +118,7 @@ fn nested_multiple() {
       vec![DependencyKey("d", None)],
     ),
   ];
-  let queries = vec![Query::new("a", vec!["d"])];
+  let queries = indexset![Query::new("a", vec!["d"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -124,15 +127,15 @@ fn nested_multiple() {
 
 #[test]
 fn self_cycle_simple() {
-  let rules = vec![Rule(
+  let rules = indexset![Rule(
     "Fib",
     "fib",
     vec![
       DependencyKey("int", None),
       DependencyKey("Fib", Some("int")),
     ],
   )];
-  let queries = vec![
+  let queries = indexset![
     Query::new("Fib", vec!["int"]),
     Query::new("Fib", vec!["Fib"]),
   ];
@@ -145,7 +148,7 @@ fn self_cycle_simple() {
 
 #[test]
 fn self_cycle_with_external_dep() {
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "Thing",
       "transitive_thing",
@@ -163,7 +166,7 @@ fn self_cycle_with_external_dep() {
       vec![DependencyKey("ExternalDep", None)],
     ),
   ];
-  let queries = vec![Query::new("Thing", vec!["int"])];
+  let queries = indexset![Query::new("Thing", vec!["int"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -173,7 +176,7 @@ fn self_cycle_with_external_dep() {
 #[test]
 fn ambiguous_cycle() {
   let _logger = env_logger::try_init();
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "Root",
       "me",
@@ -193,7 +196,7 @@ fn ambiguous_cycle() {
     Rule("FPR", "fpr_for_p", vec![DependencyKey("P", None)]),
     Rule("FPR", "fpr_for_mpp", vec![DependencyKey("MPP", None)]),
   ];
-  let queries = vec![Query::new("Root", vec![])];
+  let queries = indexset![Query::new("Root", vec![])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -202,7 +205,7 @@ fn ambiguous_cycle() {
 
 #[test]
 fn natural_loop() {
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "A",
       "a",
@@ -219,7 +222,7 @@ fn natural_loop() {
       vec![DependencyKey("F", None), DependencyKey("A", Some("D"))],
     ),
   ];
-  let queries = vec![Query::new("A", vec!["D"])];
+  let queries = indexset![Query::new("A", vec!["D"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -229,7 +232,7 @@ fn natural_loop() {
 #[test]
 fn multi_path_cycle() {
   let _logger = env_logger::try_init();
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "A",
       "sao",
@@ -245,7 +248,7 @@ fn multi_path_cycle() {
       vec![DependencyKey("AS", None), DependencyKey("A", None)],
     ),
   ];
-  let queries = vec![Query::new("A", vec![])];
+  let queries = indexset![Query::new("A", vec![])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -254,7 +257,7 @@ fn multi_path_cycle() {
 
 #[test]
 fn mutual_recursion() {
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "IsEven",
       "is_even",
@@ -272,7 +275,7 @@ fn mutual_recursion() {
       ],
     ),
   ];
-  let queries = vec![
+  let queries = indexset![
     Query::new("IsEven", vec!["int"]),
     Query::new("IsOdd", vec!["int"]),
   ];
@@ -286,7 +289,7 @@ fn mutual_recursion() {
 #[test]
 fn wide() {
   let _logger = env_logger::try_init();
-  let rules = vec![
+  let rules = indexset![
     Rule("Output", "one", vec![DependencyKey("Output", Some("A"))]),
     Rule(
       "Output",
@@ -304,7 +307,7 @@ fn wide() {
       vec![DependencyKey("C", None), DependencyKey("D", None)],
     ),
   ];
-  let queries = vec![Query::new("Output", vec!["D"])];
+  let queries = indexset![Query::new("Output", vec!["D"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -314,7 +317,7 @@ fn wide() {
 #[test]
 fn reduced_source_roots() {
   let _logger = env_logger::try_init();
-  let rules = vec![
+  let rules = indexset![
     Rule("SourceRootConfig", "construct_scope_source", vec![]),
     Rule(
       "OptionalSourceRootsResult",
@@ -381,7 +384,7 @@ fn reduced_source_roots() {
       ],
     ),
   ];
-  let queries = vec![Query::new("StrippedSourceFiles", vec!["SourceFiles"])];
+  let queries = indexset![Query::new("StrippedSourceFiles", vec!["SourceFiles"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -390,7 +393,7 @@ fn reduced_source_roots() {
 #[test]
 fn reduced_codegen_cycle() {
   let _logger = env_logger::try_init();
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "Process",
       "setup_pex_cli_process",
@@ -423,7 +426,7 @@ fn reduced_codegen_cycle() {
       vec![DependencyKey("MultiPlatformProcess", None)],
     ),
   ];
-  let queries = vec![Query::new("Process", vec!["PexCliProcess"])];
+  let queries = indexset![Query::new("Process", vec!["PexCliProcess"])];
   let graph = RuleGraph::new(rules, queries).unwrap();
 
   graph.validate_reachability().unwrap();
@@ -432,7 +435,7 @@ fn reduced_codegen_cycle() {
 #[test]
 fn full_scale_target() {
   let _logger = env_logger::try_init();
-  let rules = vec![
+  let rules = indexset![
     Rule(
       "InferredDependencies",
       "infer_python_conftest_dependencies",
@@ -797,7 +800,7 @@ fn full_scale_target() {
       ],
     ),
   ];
-  let queries = vec![
+  let queries = indexset![
     Query::new("AddressesWithOrigins", vec!["OptionsBootstrapper", "Specs"]),
     Query::new("UnexpandedTargets", vec!["OptionsBootstrapper", "Specs"]),
     Query::new("Addresses", vec!["OptionsBootstrapper", "Specs"]),