feat(tvix/store/digests): use bytes::Bytes instead of Vec<u8> r/6437

This will save us some copies, because a clone will simply create an
additional pointer to the same data.

Change-Id: I017a5d6b4c85a861b5541ebad2858ad4fbf8e8fa
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8978
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

4 files changed, 23 insertions, 17 deletions

diff --git a/tvix/store/src/proto/grpc_blobservice_wrapper.rs b/tvix/store/src/proto/grpc_blobservice_wrapper.rs
index fee97c7d2d..e60ff2ef1d 100644
--- a/tvix/store/src/proto/grpc_blobservice_wrapper.rs
+++ b/tvix/store/src/proto/grpc_blobservice_wrapper.rs
@@ -1,6 +1,4 @@
-use crate::{
-    blobservice::BlobService, proto::sync_read_into_async_read::SyncReadIntoAsyncRead, B3Digest,
-};
+use crate::{blobservice::BlobService, proto::sync_read_into_async_read::SyncReadIntoAsyncRead};
 use std::{
     collections::VecDeque,
     io,
@@ -96,7 +94,9 @@ impl super::blob_service_server::BlobService for GRPCBlobServiceWrapper {
         request: Request<super::StatBlobRequest>,
     ) -> Result<Response<super::BlobMeta>, Status> {
         let rq = request.into_inner();
-        let req_digest = B3Digest::from_vec(rq.digest)
+        let req_digest = rq
+            .digest
+            .try_into()
             .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
 
         if rq.include_chunks || rq.include_bao {
@@ -117,7 +117,9 @@ impl super::blob_service_server::BlobService for GRPCBlobServiceWrapper {
     ) -> Result<Response<Self::ReadStream>, Status> {
         let rq = request.into_inner();
 
-        let req_digest = B3Digest::from_vec(rq.digest)
+        let req_digest = rq
+            .digest
+            .try_into()
             .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
 
         match self.blob_service.open_read(&req_digest) {
diff --git a/tvix/store/src/proto/grpc_directoryservice_wrapper.rs b/tvix/store/src/proto/grpc_directoryservice_wrapper.rs
index ec9e3cb123..22fcd2fa6a 100644
--- a/tvix/store/src/proto/grpc_directoryservice_wrapper.rs
+++ b/tvix/store/src/proto/grpc_directoryservice_wrapper.rs
@@ -38,8 +38,10 @@ impl proto::directory_service_server::DirectoryService for GRPCDirectoryServiceW
             // look at the digest in the request and put it in the top of the queue.
             match &req_inner.by_what {
                 None => return Err(Status::invalid_argument("by_what needs to be specified")),
-                Some(proto::get_directory_request::ByWhat::Digest(digest)) => {
-                    let digest = B3Digest::from_vec(digest.to_vec())
+                Some(proto::get_directory_request::ByWhat::Digest(ref digest)) => {
+                    let digest: B3Digest = digest
+                        .clone()
+                        .try_into()
                         .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
 
                     task::spawn(async move {
@@ -91,6 +93,8 @@ impl proto::directory_service_server::DirectoryService for GRPCDirectoryServiceW
         // This keeps track of the seen directory keys, and their size.
         // This is used to validate the size field of a reference to a previously sent directory.
         // We don't need to keep the contents around, they're stored in the DB.
+        // https://github.com/rust-lang/rust-clippy/issues/5812
+        #[allow(clippy::mutable_key_type)]
         let mut seen_directories_sizes: HashMap<B3Digest, u32> = HashMap::new();
         let mut last_directory_dgst: Option<B3Digest> = None;
 
@@ -110,7 +114,10 @@ impl proto::directory_service_server::DirectoryService for GRPCDirectoryServiceW
             // to ensure it has been seen already in this stream, and that the size
             // matches what we recorded.
             for child_directory in &directory.directories {
-                let child_directory_digest = B3Digest::from_vec(child_directory.digest.to_vec())
+                let child_directory_digest: B3Digest = child_directory
+                    .digest
+                    .clone()
+                    .try_into()
                     .map_err(|_e| Status::internal("invalid child directory digest len"))?;
 
                 match seen_directories_sizes.get(&child_directory_digest) {
diff --git a/tvix/store/src/proto/mod.rs b/tvix/store/src/proto/mod.rs
index 7e81efc517..126a8b0edc 100644
--- a/tvix/store/src/proto/mod.rs
+++ b/tvix/store/src/proto/mod.rs
@@ -247,12 +247,11 @@ impl Directory {
     pub fn digest(&self) -> B3Digest {
         let mut hasher = blake3::Hasher::new();
 
-        let vec = hasher
+        hasher
             .update(&self.encode_to_vec())
             .finalize()
             .as_bytes()
-            .to_vec();
-        B3Digest::from_vec(vec).unwrap()
+            .into()
     }
 
     /// validate checks the directory for invalid data, such as:
diff --git a/tvix/store/src/proto/tests/directory.rs b/tvix/store/src/proto/tests/directory.rs
index 48eeaa7b5f..22b10ca746 100644
--- a/tvix/store/src/proto/tests/directory.rs
+++ b/tvix/store/src/proto/tests/directory.rs
@@ -1,7 +1,4 @@
-use crate::{
-    proto::{Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError},
-    B3Digest,
-};
+use crate::proto::{Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError};
 use lazy_static::lazy_static;
 
 lazy_static! {
@@ -69,11 +66,12 @@ fn digest() {
 
     assert_eq!(
         d.digest(),
-        B3Digest::from_vec(vec![
+        vec![
             0xaf, 0x13, 0x49, 0xb9, 0xf5, 0xf9, 0xa1, 0xa6, 0xa0, 0x40, 0x4d, 0xea, 0x36, 0xdc,
             0xc9, 0x49, 0x9b, 0xcb, 0x25, 0xc9, 0xad, 0xc1, 0x12, 0xb7, 0xcc, 0x9a, 0x93, 0xca,
             0xe4, 0x1f, 0x32, 0x62
-        ])
+        ]
+        .try_into()
         .unwrap()
     )
 }