Working, with one punt in #16751.

[ci skip-build-wheels]

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

@@ -13,6 +13,7 @@
 from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.rules import Get, rule
 from pants.engine.unions import UnionRule, union
+from pants.engine.internals.selectors import Params
 from pants.testutil.rule_runner import QueryRule, RuleRunner
 
 # -----------------------------------------------------------------------------------------------
@@ -125,8 +126,12 @@ def test_strict_equals() -> None:
 # Test unions
 # -----------------------------------------------------------------------------------------------
 
+@dataclass(frozen=True)
+class Fuel:
+    pass
 
-@union
+
+@union(in_scope_types=[Fuel])
 class Vehicle(ABC):
     @abstractmethod
     def num_wheels(self) -> int:
@@ -144,7 +149,7 @@ def num_wheels(self) -> int:
 
 
 @rule
-def car_num_wheels(car: Car) -> int:
+def car_num_wheels(car: Car, _: Fuel) -> int:
     return car.num_wheels()
 
 
@@ -163,30 +168,56 @@ async def generic_num_wheels(wrapped_vehicle: WrappedVehicle) -> int:
     return await Get(int, Vehicle, wrapped_vehicle.vehicle)
 
 
-def test_union_rules() -> None:
+def test_union_rules_in_scope_via_query() -> None:
     rule_runner = RuleRunner(
         rules=[
             car_num_wheels,
             motorcycle_num_wheels,
             UnionRule(Vehicle, Car),
             UnionRule(Vehicle, Motorcycle),
             generic_num_wheels,
-            QueryRule(int, [WrappedVehicle]),
+            QueryRule(int, [WrappedVehicle, Fuel]),
         ],
     )
-    assert rule_runner.request(int, [WrappedVehicle(Car())]) == 4
-    assert rule_runner.request(int, [WrappedVehicle(Motorcycle())]) == 2
+    assert rule_runner.request(int, [WrappedVehicle(Car()), Fuel()]) == 4
+    assert rule_runner.request(int, [WrappedVehicle(Motorcycle()), Fuel()]) == 2
 
     # Fails due to no union relationship between Vehicle -> str.
     with pytest.raises(ExecutionError) as exc:
-        rule_runner.request(int, [WrappedVehicle("not a vehicle")])  # type: ignore[arg-type]
+        rule_runner.request(int, [WrappedVehicle("not a vehicle"), Fuel()])  # type: ignore[arg-type]
     assert (
         "Invalid Get. Because an input type for `Get(int, Vehicle, not a vehicle)` was "
         "annotated with `@union`, the value for that type should be a member of that union. Did you "
         "intend to register a `UnionRule`?"
     ) in str(exc.value.args[0])
 
 
+def test_union_rules_in_scope_computed() -> None:
+    @rule
+    def fuel_singleton() -> Fuel:
+        return Fuel()
+
+    rule_runner = RuleRunner(
+        rules=[
+            car_num_wheels,
+            motorcycle_num_wheels,
+            UnionRule(Vehicle, Car),
+            UnionRule(Vehicle, Motorcycle),
+            generic_num_wheels,
+            fuel_singleton,
+            QueryRule(int, [WrappedVehicle]),
+        ],
+    )
+    assert rule_runner.request(int, [WrappedVehicle(Car())]) == 4
+    assert rule_runner.request(int, [WrappedVehicle(Motorcycle())]) == 2
+    assert False
+
+
+# -----------------------------------------------------------------------------------------------
+# Test invalid Gets
+# -----------------------------------------------------------------------------------------------
+
+
 def create_outlined_get() -> Get[int]:
     return Get(int, str, "hello")
 

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

@@ -37,7 +37,7 @@ impl<R: Rule> std::fmt::Display for Node<R> {
       Node::Query(q) => write!(f, "{}", q),
       Node::Rule(r) => write!(f, "{}", r),
       Node::Param(p) => write!(f, "Param({})", p),
-      Node::Reentry(q, in_scope) => write!(f, "Reentry({}, {})", q.product, params_str(&in_scope)),
+      Node::Reentry(q, in_scope) => write!(f, "Reentry({}, {})", q.product, params_str(in_scope)),
     }
   }
 }
@@ -59,9 +59,15 @@ impl<R: Rule> Node<R> {
 
   fn add_inherent_in_set(&self, in_set: &mut ParamTypes<R::TypeId>) {
     match self {
-      Node::Reentry(query, _) => {
+      Node::Reentry(query, in_scope_params) => {
         // Reentry nodes include in_sets computed from their Query and their dependencies.
-        in_set.extend(query.params.iter().cloned());
+        in_set.extend(
+          query
+            .params
+            .iter()
+            .filter(|p| !in_scope_params.contains(p))
+            .cloned(),
+        );
       }
       Node::Param(p) => {
         // Params are always leaves with an in-set of their own value, and no out-set.
@@ -300,6 +306,7 @@ impl<R: Rule> Builder<R> {
 
     // Rules and Reentries are created on the fly based on the out_set of dependees.
     let mut rules: HashMap<(R, ParamTypes<R::TypeId>), NodeIndex<u32>> = HashMap::default();
+    #[allow(clippy::type_complexity)]
     let mut reentries: HashMap<
       (
         Query<R::TypeId>,
@@ -894,11 +901,10 @@ impl<R: Rule> Builder<R> {
       |_edge_id, edge_weight| edge_weight.clone(),
     );
 
-    // Because the leaves of the graph (generally Param nodes) are the most significant source of
-    // information, we start there. But we will eventually visit all reachable nodes, possibly
-    // multiple times. Information flows up (the in_sets) this graph.
+    // Information flows up (the in_sets) this graph.
     let mut to_visit = graph
-      .externals(Direction::Outgoing)
+      .node_references()
+      .map(|nr| nr.id())
       .collect::<VecDeque<_>>();
     while let Some(node_id) = to_visit.pop_front() {
       if graph[node_id].is_deleted() {
@@ -1227,8 +1233,8 @@ impl<R: Rule> Builder<R> {
         }))),
         Node::Query(q) => Entry::WithDeps(Intern::new(EntryWithDeps::Root(RootEntry(q.clone())))),
         Node::Param(p) => Entry::Param(*p),
-        Node::Reentry(q, params) => Entry::WithDeps(Intern::new(EntryWithDeps::Reentry(Reentry {
-          params: params.clone(),
+        Node::Reentry(q, _) => Entry::WithDeps(Intern::new(EntryWithDeps::Reentry(Reentry {
+          params: in_set.clone(),
           query: q.clone(),
         }))),
       }
@@ -1263,10 +1269,10 @@ impl<R: Rule> Builder<R> {
         Entry::WithDeps(wd) => {
           rule_dependency_edges.insert(wd, RuleEdges { dependencies });
         }
-        e @ (Entry::Param(_) | Entry::Reentry(_)) => {
+        Entry::Param(p) => {
           if !dependencies.is_empty() {
             return Err(format!(
-              "Entry {e:?} should not have had dependencies, but had: {dependencies:#?}",
+              "Param {p} should not have had dependencies, but had: {dependencies:#?}",
             ));
           }
         }

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

@@ -87,7 +87,6 @@ impl<R: Rule> EntryWithDeps<R> {
 #[derive(DeepSizeOf, Eq, Hash, PartialEq, Clone, Debug)]
 pub enum Entry<R: Rule> {
   Param(R::TypeId),
-  Reentry(Query<R::TypeId>),
   WithDeps(Intern<EntryWithDeps<R>>),
 }
 
@@ -97,7 +96,7 @@ pub struct RootEntry<R: Rule>(Query<R::TypeId>);
 #[derive(DeepSizeOf, Eq, Hash, PartialEq, Clone, Debug)]
 pub struct Reentry<T: TypeId> {
   params: ParamTypes<T>,
-  query: Query<T>,
+  pub query: Query<T>,
 }
 
 #[derive(DeepSizeOf, Eq, Hash, PartialEq, Clone, Debug)]
@@ -192,7 +191,6 @@ impl<R: Rule> DisplayForGraph for Entry<R> {
     match self {
       Entry::WithDeps(e) => e.fmt_for_graph(display_args),
       Entry::Param(type_id) => format!("Param({})", type_id),
-      Entry::Reentry(query) => format!("Reentry({})", query),
     }
   }
 }
@@ -245,8 +243,8 @@ pub fn visualize_entry<R: Rule>(
           .map(|color| color.fmt_for_graph(display_args))
       }
     }
-    &Entry::Param(_) | &Entry::Reentry(_) => {
-      // Color "Param"s and re-entry "Query"s.
+    &Entry::Param(_) => {
+      // Color "Param"s.
       Some(Palette::Orange.fmt_for_graph(display_args))
     }
   };

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

@@ -204,7 +204,6 @@ fn multiple_provided() {
 
 #[test]
 fn in_scope_basic() {
-  let _logger = env_logger::try_init();
   let rules = indexset![
     Rule(
       "a",
@@ -252,6 +251,7 @@ fn in_scope_filtered() {
 
 #[test]
 fn in_scope_computed() {
+  let _logger = env_logger::try_init();
   let rules = indexset![
     Rule(
       "a",
@@ -260,11 +260,11 @@ fn in_scope_computed() {
         .provided_params(vec!["c"])
         .in_scope_params(vec!["d"])]
     ),
-    Rule("b", "b_from_c_and_d", vec![DependencyKey::new("c")],),
+    Rule("b", "b_from_c", vec![DependencyKey::new("c")],),
     Rule("d", "d", vec![],),
   ];
 
-  // Valid when the `in_scope` param can be computed, or is a singleton
+  // Valid when the `in_scope` param can be computed or is a singleton.
   let queries = indexset![Query::new("a", vec![])];
   let graph = RuleGraph::new(rules.clone(), queries).unwrap();
   graph.validate_reachability().unwrap();

src/rust/engine/src/nodes.rs

@@ -144,7 +144,6 @@ impl Select {
   ) -> Select {
     params.retain(|k| match entry.as_ref() {
       rule_graph::Entry::Param(type_id) => type_id == k.type_id(),
-      rule_graph::Entry::Reentry(query) => query.params.contains(k.type_id()),
       rule_graph::Entry::WithDeps(with_deps) => with_deps.params().contains(k.type_id()),
     });
     Select {
@@ -218,7 +217,7 @@ impl Select {
   async fn run_node(self, context: Context) -> NodeResult<Value> {
     match self.entry.as_ref() {
       &rule_graph::Entry::WithDeps(wd) => match wd.as_ref() {
-        rule_graph::EntryWithDeps::Inner(ref inner) => match inner.rule() {
+        rule_graph::EntryWithDeps::Rule(ref rule) => match rule.rule() {
           &tasks::Rule::Task(ref task) => {
             context
               .get(Task {
@@ -247,18 +246,28 @@ impl Select {
               .await
           }
         },
+        &rule_graph::EntryWithDeps::Reentry(ref reentry) => {
+          // TODO: Actually using the `RuleEdges` of this entry to compute inputs is not
+          // implemented: doing so would involve doing something similar to what we do for
+          // intrinsics above, and waiting to compute inputs before executing the query here.
+          //
+          // That doesn't block using a singleton to provide an API type, but it would block a more
+          // complex use case.
+          //
+          // see https://github.com/pantsbuild/pants/issues/16751
+          self.reenter(context, &reentry.query).await
+        }
         &rule_graph::EntryWithDeps::Root(_) => {
           panic!("Not a runtime-executable entry! {:?}", self.entry)
         }
       },
-      &rule_graph::Entry::Reentry(ref query) => self.reenter(context, query).await,
       &rule_graph::Entry::Param(type_id) => {
         if let Some(key) = self.params.find(type_id) {
           Ok(key.to_value())
         } else {
           Err(throw(format!(
-            "Expected a Param of type {} to be present.",
-            type_id
+            "Expected a Param of type {} to be present, but had only: {}",
+            type_id, self.params,
           )))
         }
       }