fs.Snapshot is declared in Rust (#11328)

### Problem

Our fundamental filesystem types are declared in Python, but declaring them in Rust would cut down on FFI costs, and allow for better error messages by removing some usage of the `externs::getattr*` functions.

Additionally, the previous change to move `fs.Digest` to Rust in #10905 missed the optimization opportunity of directly extracting the inner Rust `Digest` instance.

### Solution

In two commits: port `fs.Snapshot` to Rust, and optimize the conversion of a `PyDigest` to a `Digest`.

### Result

6% faster runs of `./pants dependencies --transitive ::`.

[ci skip-build-wheels]

Author: stuhood

src/python/pants/engine/fs.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, Iterable, Optional, Tuple, Union
 
 from pants.engine.collection import Collection
-from pants.engine.internals.native_engine import PyDigest
+from pants.engine.internals.native_engine import PyDigest, PySnapshot
 from pants.engine.rules import QueryRule, side_effecting
 from pants.option.global_options import GlobMatchErrorBehavior as GlobMatchErrorBehavior
 from pants.util.meta import frozen_after_init
@@ -23,17 +23,12 @@
 Digest = PyDigest
 
 
-@dataclass(frozen=True)
-class Snapshot:
-    """A Snapshot is a collection of sorted file paths and dir paths fingerprinted by their
-    names/content.
-
-    You can lift a `Digest` to a `Snapshot` with `await Get(Snapshot, Digest, my_digest)`.
-    """
+"""A Snapshot is a collection of sorted file paths and dir paths fingerprinted by their
+names/content.
 
-    digest: Digest
-    files: Tuple[str, ...]
-    dirs: Tuple[str, ...]
+You can lift a `Digest` to a `Snapshot` with `await Get(Snapshot, Digest, my_digest)`.
+"""
+Snapshot = PySnapshot
 
 
 @dataclass(frozen=True)
@@ -278,7 +273,7 @@ def write_digest(self, digest: Digest, *, path_prefix: Optional[str] = None) ->
 _EMPTY_FINGERPRINT = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 EMPTY_DIGEST = Digest(fingerprint=_EMPTY_FINGERPRINT, serialized_bytes_length=0)
 EMPTY_FILE_DIGEST = FileDigest(fingerprint=_EMPTY_FINGERPRINT, serialized_bytes_length=0)
-EMPTY_SNAPSHOT = Snapshot(EMPTY_DIGEST, files=(), dirs=())
+EMPTY_SNAPSHOT = Snapshot()
 
 
 @dataclass(frozen=True)

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

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Tuple
 
 # TODO: black and flake8 disagree about the content of this file:
 #   see https://github.com/psf/black/issues/1548
@@ -13,6 +13,15 @@ class PyDigest:
     @property
     def serialized_bytes_length(self) -> int: ...
 
+class PySnapshot:
+    def __init__(self) -> None: ...
+    @property
+    def digest(self) -> PyDigest: ...
+    @property
+    def dirs(self) -> Tuple[str, ...]: ...
+    @property
+    def files(self) -> Tuple[str, ...]: ...
+
 class PyExecutionRequest:
     def __init__(self, **kwargs: Any) -> None: ...
 

src/rust/engine/src/externs/engine_aware.rs

@@ -88,7 +88,7 @@ pub struct Artifacts {}
 impl EngineAwareInformation for Artifacts {
   type MaybeOutput = Vec<(String, Digest)>;
 
-  fn retrieve(types: &Types, value: &Value) -> Option<Self::MaybeOutput> {
+  fn retrieve(_types: &Types, value: &Value) -> Option<Self::MaybeOutput> {
     let artifacts_val = match externs::call_method(&value, "artifacts", &[]) {
       Ok(value) => value,
       Err(py_err) => {
@@ -119,7 +119,7 @@ impl EngineAwareInformation for Artifacts {
           log::error!("Error in EngineAware.artifacts() - no `digest` attr: {}", e);
         })
         .ok()?;
-      let digest = match lift_directory_digest(types, &Value::new(digest_value)) {
+      let digest = match lift_directory_digest(&Value::new(digest_value)) {
         Ok(digest) => digest,
         Err(e) => {
           log::error!("Error in EngineAware.artifacts() implementation: {}", e);

src/rust/engine/src/externs/fs.rs

@@ -38,16 +38,32 @@
 use std::borrow::Cow;
 
 use cpython::{
-  exc, py_class, CompareOp, PyErr, PyObject, PyResult, Python, PythonObject, ToPyObject,
+  exc, py_class, CompareOp, PyErr, PyObject, PyResult, PyString, PyTuple, Python, PythonObject,
+  ToPyObject,
 };
+use fs::PathStat;
 use hashing::{Digest, Fingerprint};
+use store::Snapshot;
 
+///
+/// Data members and `create_instance` methods are module-private by default, so we expose them
+/// with public top-level functions.
+///
 /// TODO: See https://github.com/dgrunwald/rust-cpython/issues/242
-pub fn new_py_digest(digest: Digest) -> PyResult<PyDigest> {
+///
+
+pub fn to_py_digest(digest: Digest) -> PyResult<PyDigest> {
   let gil = Python::acquire_gil();
   PyDigest::create_instance(gil.python(), digest)
 }
 
+pub fn from_py_digest(digest: &PyObject) -> PyResult<Digest> {
+  let gil = Python::acquire_gil();
+  let py = gil.python();
+  let py_digest = digest.extract::<PyDigest>(py)?;
+  Ok(*py_digest.digest(py))
+}
+
 py_class!(pub class PyDigest |py| {
     data digest: Digest;
     def __new__(_cls, fingerprint: Cow<str>, serialized_bytes_length: usize) -> PyResult<Self> {
@@ -84,3 +100,53 @@ py_class!(pub class PyDigest |py| {
       Ok(self.digest(py).0.prefix_hash())
     }
 });
+
+pub fn to_py_snapshot(snapshot: Snapshot) -> PyResult<PySnapshot> {
+  let gil = Python::acquire_gil();
+  PySnapshot::create_instance(gil.python(), snapshot)
+}
+
+py_class!(pub class PySnapshot |py| {
+    data snapshot: Snapshot;
+    def __new__(_cls) -> PyResult<Self> {
+      Self::create_instance(py, Snapshot::empty())
+    }
+
+    @property def digest(&self) -> PyResult<PyDigest> {
+      to_py_digest(self.snapshot(py).digest)
+    }
+
+    @property def files(&self) -> PyResult<PyTuple> {
+      let files = self.snapshot(py).path_stats.iter().filter_map(|ps| match ps {
+        PathStat::File { path, .. } => path.to_str(),
+        _ => None,
+      }).map(|ps| PyString::new(py, ps).into_object()).collect::<Vec<_>>();
+      Ok(PyTuple::new(py, &files))
+    }
+
+    @property def dirs(&self) -> PyResult<PyTuple> {
+      let dirs = self.snapshot(py).path_stats.iter().filter_map(|ps| match ps {
+        PathStat::Dir { path, .. } => path.to_str(),
+        _ => None,
+      }).map(|ps| PyString::new(py, ps).into_object()).collect::<Vec<_>>();
+      Ok(PyTuple::new(py, &dirs))
+    }
+
+    def __richcmp__(&self, other: PySnapshot, op: CompareOp) -> PyResult<PyObject> {
+      match op {
+        CompareOp::Eq => {
+          let res = self.snapshot(py).digest == other.snapshot(py).digest;
+          Ok(res.to_py_object(py).into_object())
+        },
+        CompareOp::Ne => {
+          let res = self.snapshot(py).digest != other.snapshot(py).digest;
+          Ok(res.to_py_object(py).into_object())
+        }
+        _ => Ok(py.NotImplemented()),
+      }
+    }
+
+    def __hash__(&self) -> PyResult<u64> {
+      Ok(self.snapshot(py).digest.0.prefix_hash())
+    }
+});

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

@@ -412,6 +412,7 @@ py_module_initializer!(native_engine, |py, m| {
   m.add_class::<externs::PyGeneratorResponseGetMulti>(py)?;
 
   m.add_class::<externs::fs::PyDigest>(py)?;
+  m.add_class::<externs::fs::PySnapshot>(py)?;
 
   Ok(())
 });
@@ -985,7 +986,7 @@ async fn workunit_to_py_value(
       })?;
     artifact_entries.push((
       externs::store_utf8(artifact_name.as_str()),
-      crate::nodes::Snapshot::store_snapshot(core, &snapshot).map_err(|err_str| {
+      crate::nodes::Snapshot::store_snapshot(snapshot).map_err(|err_str| {
         let gil = Python::acquire_gil();
         let py = gil.python();
         PyErr::new::<exc::Exception, _>(py, (err_str,))
@@ -1561,10 +1562,7 @@ fn capture_snapshots(
             if maybe_digest == externs::none() {
               None
             } else {
-              Some(nodes::lift_directory_digest(
-                &core.types,
-                &Value::new(maybe_digest),
-              )?)
+              Some(nodes::lift_directory_digest(&Value::new(maybe_digest))?)
             }
           };
           path_globs.map(|path_globs| (path_globs, root, digest_hint))
@@ -1585,7 +1583,7 @@ fn capture_snapshots(
               digest_hint,
             )
             .await?;
-            nodes::Snapshot::store_snapshot(&core, &snapshot)
+            nodes::Snapshot::store_snapshot(snapshot)
           }
         })
         .collect::<Vec<_>>();
@@ -1614,7 +1612,7 @@ fn ensure_remote_has_recursive(
       .iter(py)
       .map(|item| {
         let value = item.into();
-        crate::nodes::lift_directory_digest(&core.types, &value)
+        crate::nodes::lift_directory_digest(&value)
           .or_else(|_| crate::nodes::lift_file_digest(&core.types, &value))
       })
       .collect::<Result<Vec<Digest>, _>>()
@@ -1695,7 +1693,7 @@ fn run_local_interactive_process(
 
       let run_in_workspace: bool = externs::getattr(&value, "run_in_workspace").unwrap();
       let input_digest_value: Value = externs::getattr(&value, "input_digest").unwrap();
-      let input_digest: Digest = nodes::lift_directory_digest(types, &input_digest_value)?;
+      let input_digest: Digest = nodes::lift_directory_digest(&input_digest_value)?;
       let hermetic_env: bool = externs::getattr(&value, "hermetic_env").unwrap();
       let env = externs::getattr_from_frozendict(&value, "env");
 
@@ -1737,7 +1735,7 @@ fn write_digest(
       // TODO: A parent_id should be an explicit argument.
       session.workunit_store().init_thread_state(None);
 
-      let lifted_digest = nodes::lift_directory_digest(&scheduler.core.types, &digest.into())
+      let lifted_digest = nodes::lift_directory_digest(&digest.into())
         .map_err(|e| PyErr::new::<exc::ValueError, _>(py, (e,)))?;
 
       // Python will have already validated that path_prefix is a relative path.

src/rust/engine/src/intrinsics.rs

@@ -141,13 +141,12 @@ fn multi_platform_process_request_to_process_result(
     // TODO: The platform will be used in a followup.
     let _platform_val = &args[1];
 
-    let process_request = MultiPlatformExecuteProcess::lift(&context.core.types, process_val)
-      .map_err(|str| {
-        throw(&format!(
-          "Error lifting MultiPlatformExecuteProcess: {}",
-          str
-        ))
-      })?;
+    let process_request = MultiPlatformExecuteProcess::lift(process_val).map_err(|str| {
+      throw(&format!(
+        "Error lifting MultiPlatformExecuteProcess: {}",
+        str
+      ))
+    })?;
     let result = context.get(process_request).await?.0;
 
     let maybe_stdout = context
@@ -205,7 +204,7 @@ fn directory_digest_to_digest_contents(
   args: Vec<Value>,
 ) -> BoxFuture<'static, NodeResult<Value>> {
   async move {
-    let digest = lift_directory_digest(&context.core.types, &args[0]).map_err(|s| throw(&s))?;
+    let digest = lift_directory_digest(&args[0]).map_err(|s| throw(&s))?;
     let snapshot = context
       .core
       .store()
@@ -227,9 +226,8 @@ fn remove_prefix_request_to_digest(
   let store = core.store();
 
   async move {
-    let input_digest =
-      lift_directory_digest(&core.types, &externs::getattr(&args[0], "digest").unwrap())
-        .map_err(|e| throw(&e))?;
+    let input_digest = lift_directory_digest(&externs::getattr(&args[0], "digest").unwrap())
+      .map_err(|e| throw(&e))?;
     let prefix = externs::getattr_as_string(&args[0], "prefix");
     let prefix = RelativePath::new(PathBuf::from(prefix))
       .map_err(|e| throw(&format!("The `prefix` must be relative: {:?}", e)))?;
@@ -249,9 +247,8 @@ fn add_prefix_request_to_digest(
   let core = context.core;
   let store = core.store();
   async move {
-    let input_digest =
-      lift_directory_digest(&core.types, &externs::getattr(&args[0], "digest").unwrap())
-        .map_err(|e| throw(&e))?;
+    let input_digest = lift_directory_digest(&externs::getattr(&args[0], "digest").unwrap())
+      .map_err(|e| throw(&e))?;
     let prefix = externs::getattr_as_string(&args[0], "prefix");
     let prefix = RelativePath::new(PathBuf::from(prefix))
       .map_err(|e| throw(&format!("The `prefix` must be relative: {:?}", e)))?;
@@ -265,12 +262,11 @@ fn add_prefix_request_to_digest(
 }
 
 fn digest_to_snapshot(context: Context, args: Vec<Value>) -> BoxFuture<'static, NodeResult<Value>> {
-  let core = context.core.clone();
   let store = context.core.store();
   async move {
-    let digest = lift_directory_digest(&context.core.types, &args[0])?;
+    let digest = lift_directory_digest(&args[0])?;
     let snapshot = store::Snapshot::from_digest(store, digest).await?;
-    Snapshot::store_snapshot(&core, &snapshot)
+    Snapshot::store_snapshot(snapshot)
   }
   .map_err(|e: String| throw(&e))
   .boxed()
@@ -286,7 +282,7 @@ fn merge_digests_request_to_digest(
     externs::getattr::<Vec<Value>>(&args[0], "digests")
       .unwrap()
       .into_iter()
-      .map(|val: Value| lift_directory_digest(&core.types, &val))
+      .map(|val: Value| lift_directory_digest(&val))
       .collect();
   async move {
     let digest = store
@@ -397,11 +393,8 @@ fn digest_subset_to_digest(
 
   async move {
     let path_globs = Snapshot::lift_prepared_path_globs(&globs).map_err(|e| throw(&e))?;
-    let original_digest = lift_directory_digest(
-      &context.core.types,
-      &externs::getattr(&args[0], "digest").unwrap(),
-    )
-    .map_err(|e| throw(&e))?;
+    let original_digest = lift_directory_digest(&externs::getattr(&args[0], "digest").unwrap())
+      .map_err(|e| throw(&e))?;
     let subset_params = SubsetParams { globs: path_globs };
     let digest = store
       .subset(original_digest, subset_params)