borrow trick

clearer pointer casting

nativelink-store/src/filesystem_store.rs

@@ -35,7 +35,9 @@ use nativelink_util::buf_channel::{
 use nativelink_util::common::{fs, DigestInfo};
 use nativelink_util::evicting_map::{EvictingMap, LenEntry};
 use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};
-use nativelink_util::store_trait::{StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo};
+use nativelink_util::store_trait::{
+    StoreDriver, StoreKey, StoreKeyBorrow, StoreOptimizations, UploadSizeInfo,
+};
 use nativelink_util::{background_spawn, spawn_blocking};
 use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, SeekFrom};
 use tokio::time::{sleep, timeout, Sleep};
@@ -629,7 +631,7 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
 
     pub async fn get_file_entry_for_digest(&self, digest: &DigestInfo) -> Result<Arc<Fe>, Error> {
         self.evicting_map
-            .get(&digest.into())
+            .get::<StoreKey<'static>>(&digest.into())
             .await
             .ok_or_else(|| make_err!(Code::NotFound, "{digest} not found in filesystem store"))
     }
@@ -771,9 +773,11 @@ impl<Fe: FileEntry> StoreDriver for FilesystemStore<Fe> {
         // existence_cache. We need to convert the digests to owned values to be able to
         // insert them into the cache. In theory it should be able to elide this conversion
         // but it seems to be a bit tricky to get right.
-        let keys: Vec<_> = keys.iter().map(|v| v.borrow().into_owned()).collect();
+        let keys: Vec<StoreKey<'static>> = keys.iter().map(|v| v.borrow().into_owned()).collect();
         self.evicting_map
-            .sizes_for_keys(&keys, results, false /* peek */)
+            .sizes_for_keys::<&Vec<StoreKey<'_>>, StoreKey<'static>, &StoreKey<'_>>(
+                &keys, results, false, /* peek */
+            )
             .await;
         // We need to do a special pass to ensure our zero files exist.
         // If our results failed and the result was a zero file, we need to
@@ -880,17 +884,14 @@ impl<Fe: FileEntry> StoreDriver for FilesystemStore<Fe> {
             return Ok(());
         }
 
-        let entry = self
-            .evicting_map
-            .get(&key.borrow().into_owned())
-            .await
-            .ok_or_else(|| {
-                make_err!(
-                    Code::NotFound,
-                    "{} not found in filesystem store here",
-                    key.as_str()
-                )
-            })?;
+        let key_borrow: StoreKeyBorrow<'_> = key.borrow().into();
+        let entry = self.evicting_map.get(&key_borrow).await.ok_or_else(|| {
+            make_err!(
+                Code::NotFound,
+                "{} not found in filesystem store here",
+                key.as_str()
+            )
+        })?;
         let read_limit = length.unwrap_or(u64::MAX);
         let mut resumeable_temp_file = entry.read_file_part(offset, read_limit).await?;
 

nativelink-store/src/memory_store.rs

@@ -86,9 +86,12 @@ impl StoreDriver for MemoryStore {
     ) -> Result<(), Error> {
         // TODO(allada): This is a dirty hack to get around the lifetime issues with the
         // evicting map.
-        let digests: Vec<_> = keys.iter().map(|key| key.borrow().into_owned()).collect();
+        let digests: Vec<StoreKey<'static>> =
+            keys.iter().map(|key| key.borrow().into_owned()).collect();
         self.evicting_map
-            .sizes_for_keys(digests, results, false /* peek */)
+            .sizes_for_keys::<Vec<StoreKey<'_>>, StoreKey<'static>, StoreKey<'_>>(
+                digests, results, false, /* peek */
+            )
             .await;
         // We need to do a special pass to ensure our zero digest exist.
         keys.iter()

nativelink-util/src/store_trait.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::borrow::{BorrowMut, Cow};
+use std::borrow::{Borrow, BorrowMut, Cow};
 use std::collections::hash_map::DefaultHasher as StdHasher;
 use std::convert::Into;
 use std::hash::{Hash, Hasher};
@@ -144,6 +144,49 @@ pub enum StoreOptimizations {
     NoopDownloads,
 }
 
+/// A wrapper struct for [`StoreKey`] to work around
+/// lifetime limitations in `HashMap::get()` as described in
+/// <https://github.com/rust-lang/rust/issues/80389>
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[repr(transparent)]
+pub struct StoreKeyBorrow<'a>(StoreKey<'a>);
+
+impl<'a> From<StoreKey<'a>> for StoreKeyBorrow<'a> {
+    fn from(key: StoreKey<'a>) -> Self {
+        Self(key)
+    }
+}
+
+impl<'a> From<StoreKeyBorrow<'a>> for StoreKey<'a> {
+    fn from(key_borrow: StoreKeyBorrow<'a>) -> Self {
+        key_borrow.0
+    }
+}
+
+impl<'a, 'b> Borrow<StoreKeyBorrow<'a>> for StoreKey<'b>
+where
+    'b: 'a,
+{
+    fn borrow(&self) -> &StoreKeyBorrow<'a> {
+        // we are converting a reference to a StoreKey with lifetime <'b<
+        // to a reference to a StoreKeyBorrow <'a> where b the original
+        // at least outlives 'a. This is similar to the trick in https://blinsay.com/blog/compound-keys/
+        // however the wrapper struct and the struct inside the hashmap
+        // are reversed.
+        //
+        // As such we cast through pointers to convert from one to the other: this is sound
+        // as we use #[repr(transparent)] to ensure the same in memory layout
+
+        // turn a reference to `StoreKey<'b>` into a pointer to `StoreKeyBorrow<'a>`
+        let ptr = (std::ptr::from_ref::<StoreKey<'b>>(self)).cast::<StoreKeyBorrow<'a>>();
+        // turn this pointer back into a reference
+        // we know this pointer is non null as it is just derived from a reference above
+        // and that it is a valid StoreKeyBorrow<'a>, as it has the same layout and as its
+        // lifetime is shorter than the pointer it derives from (it is safe to shorten covariant lifetimes)
+        unsafe { &*ptr }
+    }
+}
+
 /// Holds something that can be converted into a key the
 /// store API can understand. Generally this is a digest
 /// but it can also be a string if the caller wishes to