5 files changed, 46 insertions, 72 deletions
diff --git a/tvix/store/src/pathinfoservice/grpc.rs b/tvix/store/src/pathinfoservice/grpc.rs
index 821331fbdef5..0e4cbb758c05 100644
--- a/tvix/store/src/pathinfoservice/grpc.rs
+++ b/tvix/store/src/pathinfoservice/grpc.rs
@@ -28,10 +28,7 @@ impl GRPCPathInfoService {
 }
 
 impl PathInfoService for GRPCPathInfoService {
-    fn get(
-        &self,
-        by_what: proto::get_path_info_request::ByWhat,
-    ) -> Result<Option<proto::PathInfo>, crate::Error> {
+    fn get(&self, digest: [u8; 20]) -> Result<Option<proto::PathInfo>, crate::Error> {
         // Get a new handle to the gRPC client.
         let mut grpc_client = self.grpc_client.clone();
 
@@ -39,7 +36,9 @@ impl PathInfoService for GRPCPathInfoService {
             self.tokio_handle.spawn(async move {
                 let path_info = grpc_client
                     .get(proto::GetPathInfoRequest {
-                        by_what: Some(by_what),
+                        by_what: Some(proto::get_path_info_request::ByWhat::ByOutputHash(
+                            digest.to_vec(),
+                        )),
                     })
                     .await?
                     .into_inner();
diff --git a/tvix/store/src/pathinfoservice/memory.rs b/tvix/store/src/pathinfoservice/memory.rs
index 4ac21883823b..d0ff1976efab 100644
--- a/tvix/store/src/pathinfoservice/memory.rs
+++ b/tvix/store/src/pathinfoservice/memory.rs
@@ -1,36 +1,22 @@
+use super::PathInfoService;
+use crate::{proto, Error};
 use std::{
     collections::HashMap,
     sync::{Arc, RwLock},
 };
 
-use crate::{proto, Error};
-use nix_compat::store_path::DIGEST_SIZE;
-
-use super::PathInfoService;
-
 #[derive(Default)]
 pub struct MemoryPathInfoService {
-    db: Arc<RwLock<HashMap<Vec<u8>, proto::PathInfo>>>,
+    db: Arc<RwLock<HashMap<[u8; 20], proto::PathInfo>>>,
 }
 
 impl PathInfoService for MemoryPathInfoService {
-    fn get(
-        &self,
-        by_what: proto::get_path_info_request::ByWhat,
-    ) -> Result<Option<proto::PathInfo>, Error> {
-        match by_what {
-            proto::get_path_info_request::ByWhat::ByOutputHash(digest) => {
-                if digest.len() != DIGEST_SIZE {
-                    return Err(Error::InvalidRequest("invalid digest length".to_string()));
-                }
+    fn get(&self, digest: [u8; 20]) -> Result<Option<proto::PathInfo>, Error> {
+        let db = self.db.read().unwrap();
 
-                let db = self.db.read().unwrap();
-
-                match db.get(&digest) {
-                    None => Ok(None),
-                    Some(path_info) => Ok(Some(path_info.clone())),
-                }
-            }
+        match db.get(&digest) {
+            None => Ok(None),
+            Some(path_info) => Ok(Some(path_info.clone())),
         }
     }
 
@@ -46,7 +32,7 @@ impl PathInfoService for MemoryPathInfoService {
             // This overwrites existing PathInfo objects.
             Ok(nix_path) => {
                 let mut db = self.db.write().unwrap();
-                db.insert(nix_path.digest.to_vec(), path_info.clone());
+                db.insert(nix_path.digest, path_info.clone());
 
                 Ok(path_info)
             }
diff --git a/tvix/store/src/pathinfoservice/mod.rs b/tvix/store/src/pathinfoservice/mod.rs
index 410bd205be1d..6a34e09af478 100644
--- a/tvix/store/src/pathinfoservice/mod.rs
+++ b/tvix/store/src/pathinfoservice/mod.rs
@@ -11,11 +11,8 @@ pub use self::sled::SledPathInfoService;
 /// The base trait all PathInfo services need to implement.
 /// This is a simple get and put of [proto::Directory], returning their digest.
 pub trait PathInfoService {
-    /// Retrieve a PathInfo message.
-    fn get(
-        &self,
-        by_what: proto::get_path_info_request::ByWhat,
-    ) -> Result<Option<proto::PathInfo>, Error>;
+    /// Retrieve a PathInfo message by the output digest.
+    fn get(&self, digest: [u8; 20]) -> Result<Option<proto::PathInfo>, Error>;
 
     /// Store a PathInfo message. Implementations MUST call validate and reject
     /// invalid messages.
diff --git a/tvix/store/src/pathinfoservice/sled.rs b/tvix/store/src/pathinfoservice/sled.rs
index b629d869f0fe..bccd6b141315 100644
--- a/tvix/store/src/pathinfoservice/sled.rs
+++ b/tvix/store/src/pathinfoservice/sled.rs
@@ -1,11 +1,9 @@
+use super::PathInfoService;
 use crate::{proto, Error};
-use nix_compat::store_path::DIGEST_SIZE;
 use prost::Message;
 use std::path::PathBuf;
 use tracing::warn;
 
-use super::PathInfoService;
-
 /// SledPathInfoService stores PathInfo in a [sled](https://github.com/spacejam/sled).
 ///
 /// The PathInfo messages are stored as encoded protos, and keyed by their output hash,
@@ -31,36 +29,25 @@ impl SledPathInfoService {
 }
 
 impl PathInfoService for SledPathInfoService {
-    fn get(
-        &self,
-        by_what: proto::get_path_info_request::ByWhat,
-    ) -> Result<Option<proto::PathInfo>, Error> {
-        match by_what {
-            proto::get_path_info_request::ByWhat::ByOutputHash(digest) => {
-                if digest.len() != DIGEST_SIZE {
-                    return Err(Error::InvalidRequest("invalid digest length".to_string()));
-                }
-
-                match self.db.get(digest) {
-                    Ok(None) => Ok(None),
-                    Ok(Some(data)) => match proto::PathInfo::decode(&*data) {
-                        Ok(path_info) => Ok(Some(path_info)),
-                        Err(e) => {
-                            warn!("failed to decode stored PathInfo: {}", e);
-                            Err(Error::StorageError(format!(
-                                "failed to decode stored PathInfo: {}",
-                                e
-                            )))
-                        }
-                    },
-                    Err(e) => {
-                        warn!("failed to retrieve PathInfo: {}", e);
-                        Err(Error::StorageError(format!(
-                            "failed to retrieve PathInfo: {}",
-                            e
-                        )))
-                    }
+    fn get(&self, digest: [u8; 20]) -> Result<Option<proto::PathInfo>, Error> {
+        match self.db.get(digest) {
+            Ok(None) => Ok(None),
+            Ok(Some(data)) => match proto::PathInfo::decode(&*data) {
+                Ok(path_info) => Ok(Some(path_info)),
+                Err(e) => {
+                    warn!("failed to decode stored PathInfo: {}", e);
+                    Err(Error::StorageError(format!(
+                        "failed to decode stored PathInfo: {}",
+                        e
+                    )))
                 }
+            },
+            Err(e) => {
+                warn!("failed to retrieve PathInfo: {}", e);
+                Err(Error::StorageError(format!(
+                    "failed to retrieve PathInfo: {}",
+                    e
+                )))
             }
         }
     }
diff --git a/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs b/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs
index 21a65185de74..8050ce10cc54 100644
--- a/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs
+++ b/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs
@@ -31,14 +31,19 @@ impl<
     ) -> Result<Response<proto::PathInfo>> {
         match request.into_inner().by_what {
             None => Err(Status::unimplemented("by_what needs to be specified")),
-            Some(by_what) => match self.path_info_service.get(by_what) {
-                Ok(None) => Err(Status::not_found("PathInfo not found")),
-                Ok(Some(path_info)) => Ok(Response::new(path_info)),
-                Err(e) => {
-                    warn!("failed to retrieve PathInfo: {}", e);
-                    Err(e.into())
+            Some(proto::get_path_info_request::ByWhat::ByOutputHash(digest)) => {
+                let digest: [u8; 20] = digest
+                    .try_into()
+                    .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
+                match self.path_info_service.get(digest) {
+                    Ok(None) => Err(Status::not_found("PathInfo not found")),
+                    Ok(Some(path_info)) => Ok(Response::new(path_info)),
+                    Err(e) => {
+                        warn!("failed to retrieve PathInfo: {}", e);
+                        Err(e.into())
+                    }
                 }
-            },
+            }
         }
     }