1 files changed, 24 insertions, 32 deletions

diff --git a/tvix/store/src/directoryservice/memory.rs b/tvix/store/src/directoryservice/memory.rs
index 87cd038888..7440be112c 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)
     }