1 files changed, 16 insertions, 20 deletions

diff --git a/tvix/store/src/directoryservice/utils.rs b/tvix/store/src/directoryservice/utils.rs
index 7d41e8d3ba01..3661808734f3 100644
--- a/tvix/store/src/directoryservice/utils.rs
+++ b/tvix/store/src/directoryservice/utils.rs
@@ -1,6 +1,7 @@
 use super::DirectoryPutter;
 use super::DirectoryService;
 use crate::proto;
+use crate::B3Digest;
 use crate::Error;
 use std::collections::{HashSet, VecDeque};
 use tracing::{debug_span, instrument, warn};
@@ -13,17 +14,17 @@ pub struct DirectoryTraverser<DS: DirectoryService> {
     /// The list of all directories that still need to be traversed. The next
     /// element is picked from the front, new elements are enqueued at the
     /// back.
-    worklist_directory_digests: VecDeque<[u8; 32]>,
+    worklist_directory_digests: VecDeque<B3Digest>,
     /// The list of directory digests already sent to the consumer.
     /// We omit sending the same directories multiple times.
-    sent_directory_digests: HashSet<[u8; 32]>,
+    sent_directory_digests: HashSet<B3Digest>,
 }
 
 impl<DS: DirectoryService> DirectoryTraverser<DS> {
-    pub fn with(directory_service: DS, root_directory_digest: &[u8; 32]) -> Self {
+    pub fn with(directory_service: DS, root_directory_digest: &B3Digest) -> Self {
         Self {
             directory_service,
-            worklist_directory_digests: VecDeque::from([*root_directory_digest]),
+            worklist_directory_digests: VecDeque::from([root_directory_digest.clone()]),
             sent_directory_digests: HashSet::new(),
         }
     }
@@ -33,12 +34,8 @@ impl<DS: DirectoryService> DirectoryTraverser<DS> {
     // This panics if the digest looks invalid, it's supposed to be checked first.
     fn enqueue_child_directories(&mut self, directory: &proto::Directory) {
         for child_directory_node in &directory.directories {
-            let child_digest: [u8; 32] = child_directory_node
-                .digest
-                .as_slice()
-                .try_into()
-                .map_err(|_e| Error::StorageError("invalid digest length".to_string()))
-                .unwrap();
+            // TODO: propagate error
+            let child_digest = B3Digest::from_vec(child_directory_node.digest.clone()).unwrap();
 
             if self.worklist_directory_digests.contains(&child_digest)
                 || self.sent_directory_digests.contains(&child_digest)
@@ -59,8 +56,7 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
         match self.worklist_directory_digests.pop_front() {
             None => None,
             Some(current_directory_digest) => {
-                let current_directory_b64 = data_encoding::BASE64.encode(&current_directory_digest);
-                let span = debug_span!("directory.digest", current_directory_b64);
+                let span = debug_span!("directory.digest", "{}", current_directory_digest);
                 let _ = span.enter();
 
                 // look up the directory itself.
@@ -73,24 +69,24 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
                             warn!("directory failed validation: {}", e.to_string());
                             return Some(Err(Error::StorageError(format!(
                                 "invalid directory: {}",
-                                current_directory_b64
+                                current_directory_digest
                             ))));
                         }
                         current_directory
                     }
                     // if it's not there, we have an inconsistent store!
                     Ok(None) => {
-                        warn!("directory {} does not exist", current_directory_b64);
+                        warn!("directory {} does not exist", current_directory_digest);
                         return Some(Err(Error::StorageError(format!(
                             "directory {} does not exist",
-                            current_directory_b64
+                            current_directory_digest
                         ))));
                     }
                     Err(e) => {
                         warn!("failed to look up directory");
                         return Some(Err(Error::StorageError(format!(
                             "unable to look up directory {}: {}",
-                            current_directory_b64, e
+                            current_directory_digest, e
                         ))));
                     }
                 };
@@ -110,7 +106,7 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
 /// TODO: verify connectivity? Factor out these checks into generic helpers?
 pub struct SimplePutter<DS: DirectoryService> {
     directory_service: DS,
-    last_directory_digest: Option<[u8; 32]>,
+    last_directory_digest: Option<B3Digest>,
 }
 
 impl<DS: DirectoryService> SimplePutter<DS> {
@@ -133,9 +129,9 @@ impl<DS: DirectoryService> DirectoryPutter for SimplePutter<DS> {
     }
 
     /// We need to be mutable here, as that's the signature of the trait.
-    fn close(&mut self) -> Result<[u8; 32], Error> {
-        match self.last_directory_digest {
-            Some(last_digest) => Ok(last_digest),
+    fn close(&mut self) -> Result<B3Digest, Error> {
+        match &self.last_directory_digest {
+            Some(last_digest) => Ok(last_digest.clone()),
             None => Err(Error::InvalidRequest(
                 "no directories sent, can't show root digest".to_string(),
             )),