Directly extract a Digest from a PyDigest: missed optimization from pantsbuild#10905.

# Building wheels and fs_util will be skipped. Delete if not intended.
[ci skip-build-wheels]

Author: stuhood

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

@@ -45,12 +45,25 @@ 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> {
@@ -88,8 +101,7 @@ py_class!(pub class PyDigest |py| {
     }
 });
 
-/// TODO: See https://github.com/dgrunwald/rust-cpython/issues/242
-pub fn new_py_snapshot(snapshot: Snapshot) -> PyResult<PySnapshot> {
+pub fn to_py_snapshot(snapshot: Snapshot) -> PyResult<PySnapshot> {
   let gil = Python::acquire_gil();
   PySnapshot::create_instance(gil.python(), snapshot)
 }
@@ -101,7 +113,7 @@ py_class!(pub class PySnapshot |py| {
     }
 
     @property def digest(&self) -> PyResult<PyDigest> {
-      new_py_digest(self.snapshot(py).digest)
+      to_py_digest(self.snapshot(py).digest)
     }
 
     @property def files(&self) -> PyResult<PyTuple> {

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

@@ -1562,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))
@@ -1615,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>, _>>()
@@ -1696,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");
 
@@ -1738,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)))?;
@@ -267,7 +264,7 @@ fn add_prefix_request_to_digest(
 fn digest_to_snapshot(context: Context, args: Vec<Value>) -> BoxFuture<'static, NodeResult<Value>> {
   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(snapshot)
   }
@@ -285,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
@@ -396,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)

src/rust/engine/src/nodes.rs

@@ -208,19 +208,8 @@ impl From<Select> for NodeKey {
   }
 }
 
-pub fn lift_directory_digest(types: &Types, digest: &Value) -> Result<hashing::Digest, String> {
-  if types.directory_digest != externs::get_type_for(digest) {
-    return Err(format!(
-      "{} is not of type {}.",
-      digest, types.directory_digest
-    ));
-  }
-  let fingerprint = externs::getattr_as_string(&digest, "fingerprint");
-  let digest_length: usize = externs::getattr(&digest, "serialized_bytes_length").unwrap();
-  Ok(hashing::Digest(
-    hashing::Fingerprint::from_hex_string(&fingerprint)?,
-    digest_length,
-  ))
+pub fn lift_directory_digest(digest: &Value) -> Result<hashing::Digest, String> {
+  externs::fs::from_py_digest(digest).map_err(|e| format!("{:?}", e))
 }
 
 pub fn lift_file_digest(types: &Types, digest: &Value) -> Result<hashing::Digest, String> {
@@ -244,11 +233,7 @@ pub struct MultiPlatformExecuteProcess {
 }
 
 impl MultiPlatformExecuteProcess {
-  fn lift_process(
-    types: &Types,
-    value: &Value,
-    platform_constraint: Option<Platform>,
-  ) -> Result<Process, String> {
+  fn lift_process(value: &Value, platform_constraint: Option<Platform>) -> Result<Process, String> {
     let env = externs::getattr_from_frozendict(&value, "env");
 
     let working_directory = {
@@ -261,8 +246,8 @@ impl MultiPlatformExecuteProcess {
     };
 
     let py_digest: Value = externs::getattr(&value, "input_digest").unwrap();
-    let digest = lift_directory_digest(types, &py_digest)
-      .map_err(|err| format!("Error parsing digest {}", err))?;
+    let digest =
+      lift_directory_digest(&py_digest).map_err(|err| format!("Error parsing digest {}", err))?;
 
     let output_files = externs::getattr::<Vec<String>>(&value, "output_files")
       .unwrap()
@@ -336,7 +321,7 @@ impl MultiPlatformExecuteProcess {
     })
   }
 
-  pub fn lift(types: &Types, value: &Value) -> Result<MultiPlatformExecuteProcess, String> {
+  pub fn lift(value: &Value) -> Result<MultiPlatformExecuteProcess, String> {
     let raw_constraints = externs::getattr::<Vec<Option<String>>>(&value, "platform_constraints")?;
     let constraints = raw_constraints
       .into_iter()
@@ -356,8 +341,7 @@ impl MultiPlatformExecuteProcess {
 
     let mut request_by_constraint: BTreeMap<Option<Platform>, Process> = BTreeMap::new();
     for (constraint, execute_process) in constraints.iter().zip(processes.iter()) {
-      let underlying_req =
-        MultiPlatformExecuteProcess::lift_process(types, execute_process, *constraint)?;
+      let underlying_req = MultiPlatformExecuteProcess::lift_process(execute_process, *constraint)?;
       request_by_constraint.insert(*constraint, underlying_req.clone());
     }
 
@@ -640,7 +624,7 @@ impl Snapshot {
   }
 
   pub fn store_directory_digest(item: &hashing::Digest) -> Result<Value, String> {
-    externs::fs::new_py_digest(*item)
+    externs::fs::to_py_digest(*item)
       .map(|d| d.into_object().into())
       .map_err(|e| format!("{:?}", e))
   }
@@ -656,7 +640,7 @@ impl Snapshot {
   }
 
   pub fn store_snapshot(item: store::Snapshot) -> Result<Value, String> {
-    externs::fs::new_py_snapshot(item)
+    externs::fs::to_py_snapshot(item)
       .map(|d| d.into_object().into())
       .map_err(|e| format!("{:?}", e))
   }