Move scc implementation to rust-cpython.

[ci skip-build-wheels]

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

@@ -36,6 +36,7 @@
     Snapshot,
     SpecsSnapshot,
 )
+from pants.engine.internals import native_engine
 from pants.engine.internals.target_adaptor import TargetAdaptor
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import (
@@ -323,14 +324,16 @@ async def transitive_targets(dependency_mapping: _DependencyMapping) -> Transiti
 @rule
 async def coarsened_targets(address: Address) -> CoarsenedTargets:
     dependency_mapping = await Get(_DependencyMapping, TransitiveTargetsRequest([address]))
-
-    # Run SCC.
-
-    # Build a mapping from Address to CoarsenedTargets.
-
-    # Return the relevant one.
-
-    return ...
+    components = native_engine.strongly_connected_components(list(dependency_mapping))
+
+    # Return the relevant component.
+    # TODO: This is very inefficient. Before landing, this should apply a strategy similar to
+    # https://github.com/pantsbuild/pants/issues/11270 to share subgraphs which have already been
+    # computed.
+    for component in components:
+        if address in component:
+            return CoarsenedTargets(component)
+    raise AssertionError("No component contained the relevant Address.")
 
 
 # -----------------------------------------------------------------------------------------------

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

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 from io import RawIOBase
-from typing import Any, Sequence, TextIO
+from typing import Any, Sequence, TextIO, Tuple
 
 from typing_extensions import Protocol
 
@@ -139,6 +139,9 @@ def rule_subgraph_visualize(
 ) -> None: ...
 def garbage_collect_store(scheduler: PyScheduler, target_size_bytes: int) -> None: ...
 def lease_files_in_graph(scheduler: PyScheduler, session: PySession) -> None: ...
+def strongly_connected_components(
+    adjacency_lists: Sequence[Tuple[Any, Sequence[Any]]]
+) -> Sequence[Sequence[Any]]: ...
 
 class PyDigest:
     def __init__(self, fingerprint: str, serialized_bytes_length: int) -> None: ...

src/rust/engine/Cargo.toml

@@ -122,6 +122,7 @@ logging = { path = "logging" }
 nailgun = { path = "nailgun" }
 num_enum = "0.4"
 parking_lot = "0.11"
+petgraph = "0.4.5"
 process_execution = { path = "process_execution" }
 rand = "0.8"
 regex = "1"

src/rust/engine/engine_pyo3/Cargo.toml

@@ -21,7 +21,6 @@ default = []
 hashing = { path = "../hashing" }
 nailgun = { path = "../nailgun" }
 parking_lot = "0.11"
-petgraph = "0.4.5"
 # We must disable the `auto-initialize` feature because we do not enable `extension-module` normally
 # (see above comment in `features`), so `auto-initialize` would try to link to a static Python interpreter during
 # tests, which may fail. However, we need to then re-activate the `macros` feature. See

src/rust/engine/engine_pyo3/src/externs/interface.rs

@@ -1,12 +1,8 @@
 // Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 // Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-use std::collections::HashMap;
-
-use petgraph::graph::{DiGraph, Graph};
 use pyo3::exceptions::PyException;
 use pyo3::prelude::*;
-use pyo3::wrap_pyfunction;
 
 mod nailgun;
 mod testutil;
@@ -18,46 +14,9 @@ fn native_engine_pyo3(py: Python, m: &PyModule) -> PyResult<()> {
 
   m.add_class::<PyExecutor>()?;
 
-  m.add_function(wrap_pyfunction!(strongly_connected_components, m)?)
-    .unwrap();
-
   Ok(())
 }
 
-// TODO: Would be nice to be able to accept any Python object, rather than only u32.
-#[pyfunction]
-fn strongly_connected_components(
-  adjacency_lists: HashMap<u32, Vec<u32>>,
-) -> PyResult<Vec<Vec<u32>>> {
-  let mut graph: DiGraph<u32, (), u32> = Graph::new();
-  let mut node_ids = HashMap::new();
-
-  for (node, adjacency_list) in adjacency_lists {
-    let node_id = node_ids
-      .entry(node)
-      .or_insert_with(|| graph.add_node(node))
-      .clone();
-    for dependency in adjacency_list {
-      let dependency_id = node_ids
-        .entry(dependency)
-        .or_insert_with(|| graph.add_node(dependency));
-      graph.add_edge(node_id, *dependency_id, ());
-    }
-  }
-
-  Ok(
-    petgraph::algo::tarjan_scc(&graph)
-      .into_iter()
-      .map(|component| {
-        component
-          .into_iter()
-          .map(|node_id| graph[node_id])
-          .collect::<Vec<_>>()
-      })
-      .collect(),
-  )
-}
-
 #[pyclass]
 #[derive(Debug, Clone)]
 struct PyExecutor {

src/rust/engine/src/externs/interface.rs

@@ -66,6 +66,7 @@ use hashing::Digest;
 use log::{self, debug, error, warn, Log};
 use logging::logger::PANTS_LOGGER;
 use logging::{Logger, PythonLogLevel};
+use petgraph::graph::{DiGraph, Graph};
 use process_execution::RemoteCacheWarningsBehavior;
 use regex::Regex;
 use rule_graph::{self, RuleGraph};
@@ -77,7 +78,7 @@ use workunit_store::{
 
 use crate::{
   externs, nodes, Core, ExecutionRequest, ExecutionStrategyOptions, ExecutionTermination, Failure,
-  Function, Intrinsics, LocalStoreOptions, Params, RemotingOptions, Rule, Scheduler, Session,
+  Function, Intrinsics, Key, LocalStoreOptions, Params, RemotingOptions, Rule, Scheduler, Session,
   Tasks, Types, Value,
 };
 
@@ -396,6 +397,15 @@ py_module_initializer!(native_engine, |py, m| {
     py_fn!(py, ensure_remote_has_recursive(a: PyScheduler, b: PyList)),
   )?;
 
+  m.add(
+    py,
+    "strongly_connected_components",
+    py_fn!(
+      py,
+      strongly_connected_components(a: Vec<(PyObject, Vec<PyObject>)>)
+    ),
+  )?;
+
   m.add_class::<PyExecutionRequest>(py)?;
   m.add_class::<PyExecutionStrategyOptions>(py)?;
   m.add_class::<PyExecutor>(py)?;
@@ -822,6 +832,41 @@ fn nailgun_server_await_shutdown(py: Python, nailgun_server_ptr: PyNailgunServer
   })
 }
 
+fn strongly_connected_components(
+  py: Python,
+  adjacency_lists: Vec<(PyObject, Vec<PyObject>)>,
+) -> CPyResult<Vec<Vec<PyObject>>> {
+  let mut graph: DiGraph<Key, (), u32> = Graph::new();
+  let mut node_ids: HashMap<Key, _> = HashMap::new();
+
+  for (node, adjacency_list) in adjacency_lists {
+    let node_key = externs::key_for(node.clone_ref(py).into())?;
+    let node_id = node_ids
+      .entry(node_key)
+      .or_insert_with(|| graph.add_node(node_key))
+      .clone();
+    for dependency in adjacency_list {
+      let dependency_key = externs::key_for(dependency.clone_ref(py).into())?;
+      let dependency_id = node_ids
+        .entry(dependency_key)
+        .or_insert_with(|| graph.add_node(dependency_key));
+      graph.add_edge(node_id, *dependency_id, ());
+    }
+  }
+
+  Ok(
+    petgraph::algo::tarjan_scc(&graph)
+      .into_iter()
+      .map(|component| {
+        component
+          .into_iter()
+          .map(|node_id| externs::val_for(&graph[node_id]).consume_into_py_object(py))
+          .collect::<Vec<_>>()
+      })
+      .collect(),
+  )
+}
+
 ///
 /// Given a set of Tasks and type information, creates a Scheduler.
 ///