feat(tvix/nix-compat/nixhash/NixHashWithMode): from_algo_mode_hash r/6813

Change-Id: Ieae628fab1926a498ae5c3eb27df1b722e6151d7
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9728
Reviewed-by: edef <edef@edef.eu>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

2 files changed, 48 insertions, 2 deletions

diff --git a/tvix/nix-compat/src/nixhash/mod.rs b/tvix/nix-compat/src/nixhash/mod.rs
index 3b2810b8b0bf..9586473d1b9a 100644
--- a/tvix/nix-compat/src/nixhash/mod.rs
+++ b/tvix/nix-compat/src/nixhash/mod.rs
@@ -65,8 +65,8 @@ impl TryFrom<(HashAlgo, &[u8])> for NixHash {
 }
 
 /// Constructs a new [NixHash] by specifying [HashAlgo] and digest.
-// It can fail if the passed digest length doesn't match what's expected for
-// the passed algo.
+/// It can fail if the passed digest length doesn't match what's expected for
+/// the passed algo.
 pub fn from_algo_and_digest(algo: HashAlgo, digest: &[u8]) -> Result<NixHash> {
     if digest.len() != algo.digest_length() {
         return Err(Error::InvalidEncodedDigestLength(digest.len(), algo));
diff --git a/tvix/nix-compat/src/nixhash/with_mode.rs b/tvix/nix-compat/src/nixhash/with_mode.rs
index 86c14dc25d9c..efea3d0e6231 100644
--- a/tvix/nix-compat/src/nixhash/with_mode.rs
+++ b/tvix/nix-compat/src/nixhash/with_mode.rs
@@ -30,6 +30,21 @@ pub enum NixHashWithMode {
 }
 
 impl NixHashWithMode {
+    /// Construct a [NixHashWithMode] from a string containing the algo, and
+    /// optionally a `r:` prefix, and a digest (bytes).
+    pub fn from_algo_mode_hash(algo_and_mode: &str, digest: &[u8]) -> super::Result<Self> {
+        Ok(match algo_and_mode.strip_prefix("r:") {
+            Some(algo) => nixhash::NixHashWithMode::Recursive(nixhash::from_algo_and_digest(
+                algo.try_into()?,
+                &digest,
+            )?),
+            None => nixhash::NixHashWithMode::Flat(nixhash::from_algo_and_digest(
+                algo_and_mode.try_into()?,
+                &digest,
+            )?),
+        })
+    }
+
     pub fn mode(&self) -> NixHashMode {
         match self {
             Self::Flat(_) => NixHashMode::Flat,
@@ -186,3 +201,34 @@ impl<'de> Deserialize<'de> for NixHashWithMode {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::nixhash::{NixHash, NixHashWithMode};
+    use lazy_static::lazy_static;
+    use test_case::test_case;
+
+    const DIGEST_SHA256: [u8; 32] = [
+        0xa5, 0xce, 0x9c, 0x15, 0x5e, 0xd0, 0x93, 0x97, 0x61, 0x46, 0x46, 0xc9, 0x71, 0x7f, 0xc7,
+        0xcd, 0x94, 0xb1, 0x02, 0x3d, 0x7b, 0x76, 0xb6, 0x18, 0xd4, 0x09, 0xe4, 0xfe, 0xfd, 0x6e,
+        0x9d, 0x39,
+    ];
+    lazy_static! {
+        pub static ref NIXHASH_SHA256: NixHash = NixHash::Sha256(DIGEST_SHA256);
+    }
+
+    #[test_case("sha256", &DIGEST_SHA256, NixHashWithMode::Flat(NIXHASH_SHA256.clone()); "sha256 flat")]
+    #[test_case("r:sha256", &DIGEST_SHA256, NixHashWithMode::Recursive(NIXHASH_SHA256.clone()); "sha256 recursive")]
+    fn from_from_algo_mode_hash(algo_and_mode: &str, digest: &[u8], expected: NixHashWithMode) {
+        assert_eq!(
+            expected,
+            NixHashWithMode::from_algo_mode_hash(algo_and_mode, digest).unwrap()
+        );
+    }
+
+    #[test]
+    fn from_algo_mode_failure() {
+        assert!(NixHashWithMode::from_algo_mode_hash("r:sha256", &[]).is_err());
+        assert!(NixHashWithMode::from_algo_mode_hash("ha256", &DIGEST_SHA256).is_err());
+    }
+}