3 files changed, 54 insertions, 73 deletions
diff --git a/tvix/store/src/directoryservice/memory.rs b/tvix/store/src/directoryservice/memory.rs
index 87cd038888e6..7440be112cfd 100644
--- a/tvix/store/src/directoryservice/memory.rs
+++ b/tvix/store/src/directoryservice/memory.rs
@@ -8,46 +8,38 @@ use super::DirectoryService;
 
 #[derive(Clone, Default)]
 pub struct MemoryDirectoryService {
-    db: Arc<RwLock<HashMap<Vec<u8>, proto::Directory>>>,
+    db: Arc<RwLock<HashMap<[u8; 32], proto::Directory>>>,
 }
 
 impl DirectoryService for MemoryDirectoryService {
-    // TODO: change api to only be by digest
-    #[instrument(skip(self, by_what))]
-    fn get(
-        &self,
-        by_what: &proto::get_directory_request::ByWhat,
-    ) -> Result<Option<proto::Directory>, Error> {
-        match by_what {
-            proto::get_directory_request::ByWhat::Digest(digest) => {
-                let db = self.db.read()?;
-
-                match db.get(digest) {
-                    // The directory was not found, return
-                    None => Ok(None),
-
-                    // The directory was found, try to parse the data as Directory message
-                    Some(directory) => {
-                        // Validate the retrieved Directory indeed has the
-                        // digest we expect it to have, to detect corruptions.
-                        let actual_digest = directory.digest();
-                        if actual_digest.as_slice() != digest {
-                            return Err(Error::StorageError(format!(
-                                "requested directory with digest {}, but got {}",
-                                BASE64.encode(digest),
-                                BASE64.encode(&actual_digest)
-                            )));
-                        }
-
-                        Ok(Some(directory.clone()))
-                    }
+    #[instrument(skip(self, digest), fields(directory.digest = BASE64.encode(digest)))]
+    fn get(&self, digest: &[u8; 32]) -> Result<Option<proto::Directory>, Error> {
+        let db = self.db.read()?;
+
+        match db.get(digest) {
+            // The directory was not found, return
+            None => Ok(None),
+
+            // The directory was found, try to parse the data as Directory message
+            Some(directory) => {
+                // Validate the retrieved Directory indeed has the
+                // digest we expect it to have, to detect corruptions.
+                let actual_digest = directory.digest();
+                if actual_digest.as_slice() != digest {
+                    return Err(Error::StorageError(format!(
+                        "requested directory with digest {}, but got {}",
+                        BASE64.encode(digest),
+                        BASE64.encode(&actual_digest)
+                    )));
                 }
+
+                Ok(Some(directory.clone()))
             }
         }
     }
 
     #[instrument(skip(self, directory), fields(directory.digest = BASE64.encode(&directory.digest())))]
-    fn put(&self, directory: proto::Directory) -> Result<Vec<u8>, Error> {
+    fn put(&self, directory: proto::Directory) -> Result<[u8; 32], Error> {
         let digest = directory.digest();
 
         // validate the directory itself.
@@ -61,7 +53,7 @@ impl DirectoryService for MemoryDirectoryService {
 
         // store it
         let mut db = self.db.write()?;
-        db.insert(digest.clone(), directory);
+        db.insert(digest, directory);
 
         Ok(digest)
     }
diff --git a/tvix/store/src/directoryservice/mod.rs b/tvix/store/src/directoryservice/mod.rs
index 4abf823b23b4..a27f6745332a 100644
--- a/tvix/store/src/directoryservice/mod.rs
+++ b/tvix/store/src/directoryservice/mod.rs
@@ -11,11 +11,8 @@ pub use self::sled::SledDirectoryService;
 pub trait DirectoryService {
     /// Get looks up a single Directory message by its digest.
     /// In case the directory is not found, Ok(None) is returned.
-    fn get(
-        &self,
-        by_what: &proto::get_directory_request::ByWhat,
-    ) -> Result<Option<proto::Directory>, Error>;
+    fn get(&self, digest: &[u8; 32]) -> Result<Option<proto::Directory>, Error>;
     /// Get uploads a single Directory message, and returns the calculated
     /// digest, or an error.
-    fn put(&self, directory: proto::Directory) -> Result<Vec<u8>, Error>;
+    fn put(&self, directory: proto::Directory) -> Result<[u8; 32], Error>;
 }
diff --git a/tvix/store/src/directoryservice/sled.rs b/tvix/store/src/directoryservice/sled.rs
index 064f86ceaa6e..1f729a594c19 100644
--- a/tvix/store/src/directoryservice/sled.rs
+++ b/tvix/store/src/directoryservice/sled.rs
@@ -29,48 +29,40 @@ impl SledDirectoryService {
 }
 
 impl DirectoryService for SledDirectoryService {
-    // TODO: change api to only be by digest
-    #[instrument(name = "SledDirectoryService::get", skip(self, by_what))]
-    fn get(
-        &self,
-        by_what: &proto::get_directory_request::ByWhat,
-    ) -> Result<Option<proto::Directory>, Error> {
-        match by_what {
-            proto::get_directory_request::ByWhat::Digest(digest) => {
-                match self.db.get(digest) {
-                    // The directory was not found, return
-                    Ok(None) => Ok(None),
+    #[instrument(name = "SledDirectoryService::get", skip(self, digest), fields(directory.digest = BASE64.encode(digest)))]
+    fn get(&self, digest: &[u8; 32]) -> Result<Option<proto::Directory>, Error> {
+        match self.db.get(digest) {
+            // The directory was not found, return
+            Ok(None) => Ok(None),
 
-                    // The directory was found, try to parse the data as Directory message
-                    Ok(Some(data)) => match Directory::decode(&*data) {
-                        Ok(directory) => {
-                            // Validate the retrieved Directory indeed has the
-                            // digest we expect it to have, to detect corruptions.
-                            let actual_digest = directory.digest();
-                            if actual_digest.as_slice() != digest {
-                                return Err(Error::StorageError(format!(
-                                    "requested directory with digest {}, but got {}",
-                                    BASE64.encode(digest),
-                                    BASE64.encode(&actual_digest)
-                                )));
-                            }
+            // The directory was found, try to parse the data as Directory message
+            Ok(Some(data)) => match Directory::decode(&*data) {
+                Ok(directory) => {
+                    // Validate the retrieved Directory indeed has the
+                    // digest we expect it to have, to detect corruptions.
+                    let actual_digest = directory.digest();
+                    if actual_digest.as_slice() != digest {
+                        return Err(Error::StorageError(format!(
+                            "requested directory with digest {}, but got {}",
+                            BASE64.encode(digest),
+                            BASE64.encode(&actual_digest)
+                        )));
+                    }
 
-                            Ok(Some(directory))
-                        }
-                        Err(e) => {
-                            warn!("unable to parse directory {}: {}", BASE64.encode(digest), e);
-                            Err(Error::StorageError(e.to_string()))
-                        }
-                    },
-                    // some storage error?
-                    Err(e) => Err(Error::StorageError(e.to_string())),
+                    Ok(Some(directory))
                 }
-            }
+                Err(e) => {
+                    warn!("unable to parse directory {}: {}", BASE64.encode(digest), e);
+                    Err(Error::StorageError(e.to_string()))
+                }
+            },
+            // some storage error?
+            Err(e) => Err(Error::StorageError(e.to_string())),
         }
     }
 
     #[instrument(name = "SledDirectoryService::put", skip(self, directory), fields(directory.digest = BASE64.encode(&directory.digest())))]
-    fn put(&self, directory: proto::Directory) -> Result<Vec<u8>, Error> {
+    fn put(&self, directory: proto::Directory) -> Result<[u8; 32], Error> {
         let digest = directory.digest();
 
         // validate the directory itself.
@@ -82,7 +74,7 @@ impl DirectoryService for SledDirectoryService {
             )));
         }
         // store it
-        let result = self.db.insert(&digest, directory.encode_to_vec());
+        let result = self.db.insert(digest, directory.encode_to_vec());
         if let Err(e) = result {
             return Err(Error::StorageError(e.to_string()));
         }