perf: Use rkyv to deserialize paths, on a need-to-use basis

LukeMathWalker · LukeMathWalker · commit d0d6ec617b15 · 2026-01-17T09:57:45.000+01:00
diff --git a/compiler/pavex_rustdoc_types/src/lib.rs b/compiler/pavex_rustdoc_types/src/lib.rs
@@ -798,6 +798,8 @@ pub struct Id(pub u32);
     rkyv::Deserialize,
 )]
 #[rkyv(derive(Debug))]
+#[rkyv(compare(PartialEq))]
+#[repr(u8)]
 #[serde(rename_all = "snake_case")]
 pub enum ItemKind {
     /// A module declaration, e.g. `mod foo;` or `mod foo {}`
diff --git a/compiler/pavexc/src/rustdoc/compute/cache.rs b/compiler/pavexc/src/rustdoc/compute/cache.rs
@@ -18,7 +18,8 @@ use crate::{
     rustdoc::{
         annotations::AnnotatedItems,
         queries::{
-            CrateData, CrateItemIndex, ImportPath2Id, LazyCrateItemIndex, LazyImportPath2Id,
+            CrateData, CrateItemIndex, CrateItemPaths, ImportPath2Id, LazyCrateItemIndex,
+            LazyCrateItemPaths, LazyImportPath2Id,
         },
     },
 };
@@ -380,7 +381,7 @@ impl ToolchainCache {
         let krate = CacheEntry {
             root_item_id,
             external_crates: Cow::Borrowed(external_crates),
-            paths: Cow::Borrowed(paths),
+            paths: RkyvCowBytes::Borrowed(paths),
             format_version,
             items: RkyvCowBytes::Borrowed(items),
             secondary_indexes: Some(SecondaryIndexes {
@@ -601,7 +602,7 @@ impl ThirdPartyCrateCache {
             let krate = CacheEntry {
                 root_item_id,
                 external_crates: Cow::Borrowed(external_crates),
-                paths: Cow::Borrowed(paths),
+                paths: RkyvCowBytes::Borrowed(paths),
                 format_version,
                 items: RkyvCowBytes::Borrowed(items),
                 secondary_indexes,
@@ -786,7 +787,7 @@ impl ThirdPartyCrateCache {
 pub(in crate::rustdoc) struct CacheEntry<'a> {
     root_item_id: u32,
     external_crates: Cow<'a, [u8]>,
-    paths: Cow<'a, [u8]>,
+    paths: RkyvCowBytes<'a>,
     format_version: i64,
     items: RkyvCowBytes<'a>,
     secondary_indexes: Option<SecondaryIndexes<'a>>,
@@ -879,19 +880,26 @@ impl<'a> CacheEntry<'a> {
                 the same crate twice? This is a bug."
             );
         };
+        let CrateItemPaths::Eager(paths) = &crate_data.paths else {
+            anyhow::bail!(
+                "The crate item paths is not deserialized. Are we trying to cache \
+                the same crate twice? This is a bug."
+            );
+        };
 
-        // Serialize the items HashMap using rkyv for zero-copy deserialization later.
         let items = rkyv::to_bytes::<rkyv::rancor::Error>(&index.index)
             .map_err(|e| anyhow::anyhow!(e).context("Failed to serialize crate items with rkyv"))?;
 
         let external_crates =
             bincode::serde::encode_to_vec(&crate_data.external_crates, BINCODE_CONFIG)?;
-        let paths = bincode::serde::encode_to_vec(&crate_data.paths, BINCODE_CONFIG)?;
+        let paths = rkyv::to_bytes::<rkyv::rancor::Error>(&paths.paths).map_err(|e| {
+            anyhow::anyhow!(e).context("Failed to serialize item summaries with rkyv")
+        })?;
 
         Ok(CacheEntry {
             root_item_id: crate_data.root_item_id.0,
             external_crates: Cow::Owned(external_crates),
-            paths: Cow::Owned(paths),
+            paths: RkyvCowBytes::Owned(paths),
             format_version: crate_data.format_version as i64,
             items: RkyvCowBytes::Owned(items),
             secondary_indexes: None,
@@ -904,17 +912,14 @@ impl<'a> CacheEntry<'a> {
     /// since it can be quite large and deserialization can be slow for large crates.
     /// The item index is stored as rkyv-serialized bytes for zero-copy access.
     pub(super) fn hydrate(self, package_id: PackageId) -> Result<RustdocCacheEntry, anyhow::Error> {
-        let paths = tracing::trace_span!("Deserialize paths")
-            .in_scope(|| bincode::decode_from_slice(&self.paths, BINCODE_CONFIG))
-            .context("Failed to deserialize paths")?
-            .0;
-
         let crate_data = CrateData {
             root_item_id: rustdoc_types::Id(self.root_item_id.to_owned()),
             external_crates: bincode::decode_from_slice(&self.external_crates, BINCODE_CONFIG)
                 .context("Failed to deserialize external_crates")?
                 .0,
-            paths,
+            paths: CrateItemPaths::Lazy(LazyCrateItemPaths {
+                bytes: self.paths.into_owned(),
+            }),
             format_version: self.format_version.try_into()?,
             index: CrateItemIndex::Lazy(LazyCrateItemIndex {
                 bytes: self.items.into_owned(),
diff --git a/compiler/pavexc/src/rustdoc/queries.rs b/compiler/pavexc/src/rustdoc/queries.rs
@@ -11,13 +11,15 @@ use guppy::{PackageId, Version};
 use indexmap::IndexSet;
 use rayon::iter::IntoParallelRefIterator;
 use rkyv::collections::swiss_table::ArchivedHashMap;
+use rkyv::hash::FxHasher64;
 use rkyv::rancor::Panic;
 use rkyv::string::ArchivedString;
 use rkyv::util::AlignedVec;
 use rkyv::vec::ArchivedVec;
 use rustc_hash::FxHashMap;
 use rustdoc_types::{
-    ArchivedId, ArchivedItem, ExternalCrate, Item, ItemEnum, ItemKind, ItemSummary, Visibility,
+    ArchivedId, ArchivedItem, ArchivedItemSummary, ExternalCrate, Item, ItemEnum, ItemKind,
+    ItemSummary, Visibility,
 };
 use tracing::Span;
 use tracing_log_error::log_error;
@@ -989,13 +991,131 @@ pub(crate) struct CrateData {
     pub external_crates: FxHashMap<u32, ExternalCrate>,
     /// A mapping from the id of a type to its fully qualified path.
     /// Primarily useful for foreign items that are being re-exported by this crate.
-    #[allow(clippy::disallowed_types)]
-    pub paths: FxHashMap<rustdoc_types::Id, ItemSummary>,
+    pub paths: CrateItemPaths,
     /// The version of the JSON format used by rustdoc.
     pub format_version: u32,
     /// The index of all the items in the crate.
     pub index: CrateItemIndex,
 }
+
+#[derive(Debug, Clone)]
+/// A mapping from the id of a type to its fully qualified path.
+///
+/// Primarily useful for foreign items that are being re-exported by this crate.
+pub(crate) enum CrateItemPaths {
+    Eager(EagerCrateItemPaths),
+    Lazy(LazyCrateItemPaths),
+}
+
+impl CrateItemPaths {
+    /// Retrieve an item summary from the index given its id.
+    pub fn get(&self, id: &rustdoc_types::Id) -> Option<Cow<'_, ItemSummary>> {
+        match self {
+            Self::Eager(m) => m.paths.get(id).map(Cow::Borrowed),
+            Self::Lazy(m) => {
+                let item = m.get_deserialized(id)?;
+                Some(Cow::Owned(item))
+            }
+        }
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (rustdoc_types::Id, ItemSummaryRef<'_>)> {
+        match self {
+            CrateItemPaths::Eager(paths) => CrateItemPathsIter::Eager(paths.paths.iter()),
+            CrateItemPaths::Lazy(paths) => CrateItemPathsIter::Lazy(paths.archived().iter()),
+        }
+    }
+}
+
+pub enum CrateItemPathsIter<'a> {
+    Eager(std::collections::hash_map::Iter<'a, rustdoc_types::Id, ItemSummary>),
+    Lazy(
+        rkyv::collections::swiss_table::map::Iter<'a, ArchivedId, ArchivedItemSummary, FxHasher64>,
+    ),
+}
+
+pub enum ItemSummaryRef<'a> {
+    Eager(&'a ItemSummary),
+    Lazy(&'a ArchivedItemSummary),
+}
+
+impl<'a> ItemSummaryRef<'a> {
+    pub fn crate_id(&self) -> u32 {
+        match self {
+            ItemSummaryRef::Eager(s) => s.crate_id,
+            ItemSummaryRef::Lazy(s) => s.crate_id.to_native(),
+        }
+    }
+
+    pub fn kind(&self) -> ItemKind {
+        match self {
+            ItemSummaryRef::Eager(s) => s.kind,
+            ItemSummaryRef::Lazy(s) => {
+                // Safe to do since the enum is repr(u8)
+                rkyv::deserialize::<_, rkyv::rancor::Infallible>(&s.kind).unwrap()
+            }
+        }
+    }
+
+    pub fn path(&self) -> Cow<'_, [String]> {
+        match self {
+            ItemSummaryRef::Eager(s) => Cow::Borrowed(&s.path),
+            ItemSummaryRef::Lazy(s) => {
+                Cow::Owned(s.path.iter().map(|s| s.as_str().to_owned()).collect())
+            }
+        }
+    }
+}
+
+impl<'a> Iterator for CrateItemPathsIter<'a> {
+    type Item = (rustdoc_types::Id, ItemSummaryRef<'a>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            Self::Eager(iter) => iter.next().map(|(k, v)| (*k, ItemSummaryRef::Eager(v))),
+            Self::Lazy(iter) => iter
+                .next()
+                .map(|(k, v)| (rustdoc_types::Id(k.0.to_native()), ItemSummaryRef::Lazy(v))),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+/// See [`CrateItemPaths`] for more information.
+pub(crate) struct EagerCrateItemPaths {
+    #[allow(clippy::disallowed_types)]
+    pub paths: FxHashMap<rustdoc_types::Id, ItemSummary>,
+}
+
+/// See [`CrateItemPaths`] for more information.
+#[derive(Debug, Clone)]
+pub(crate) struct LazyCrateItemPaths {
+    pub(super) bytes: AlignedVec,
+}
+
+impl LazyCrateItemPaths {
+    /// Get zero-copy access to the archived HashMap.
+    #[inline]
+    fn archived(&self) -> &ArchivedHashMap<ArchivedId, ArchivedItemSummary> {
+        // SAFETY: The bytes were serialized by rkyv from a valid HashMap<Id, ItemSummary>.
+        // We trust the cache to contain valid data.
+        unsafe {
+            rkyv::access_unchecked::<ArchivedHashMap<ArchivedId, ArchivedItemSummary>>(&self.bytes)
+        }
+    }
+
+    /// Get an item by its ID, returning a reference to the archived summary.
+    pub fn get(&self, id: &rustdoc_types::Id) -> Option<&ArchivedItemSummary> {
+        self.archived().get(&ArchivedId(id.0.into()))
+    }
+
+    /// Deserialize a summary by its ID.
+    pub fn get_deserialized(&self, id: &rustdoc_types::Id) -> Option<ItemSummary> {
+        let archived = self.get(id)?;
+        Some(rkyv::deserialize::<ItemSummary, Panic>(archived).unwrap())
+    }
+}
+
 #[derive(Debug, Clone)]
 /// The index of all the items in the crate.
 ///
@@ -1107,7 +1227,7 @@ impl Crate {
             index: CrateItemIndex::Eager(EagerCrateItemIndex { index: krate.index }),
             external_crates: krate.external_crates,
             format_version: krate.format_version,
-            paths: krate.paths,
+            paths: CrateItemPaths::Eager(EagerCrateItemPaths { paths: krate.paths }),
         };
         Self::index(crate_data, package_id, diagnostics)
     }
@@ -1123,18 +1243,18 @@ impl Crate {
             .iter()
             .filter_map(|(id, summary)| {
                 // We only want types, no macros
-                if matches!(summary.kind, ItemKind::Macro | ItemKind::ProcDerive) {
+                if matches!(summary.kind(), ItemKind::Macro | ItemKind::ProcDerive) {
                     return None;
                 }
                 // We will index local items on our own.
                 // We don't get them from `paths` because it may include private items
                 // as well, and we don't have a way to figure out if an item is private
                 // or not from the summary info.
-                if summary.crate_id == 0 {
+                if summary.crate_id() == 0 {
                     return None;
                 }
 
-                Some((summary.path.clone(), id.to_owned()))
+                Some((summary.path().into_owned(), id.to_owned()))
             })
             .collect();
 
@@ -1312,7 +1432,7 @@ impl Crate {
     fn get_summary_by_local_type_id(
         &self,
         id: &rustdoc_types::Id,
-    ) -> Result<&rustdoc_types::ItemSummary, anyhow::Error> {
+    ) -> Result<Cow<'_, rustdoc_types::ItemSummary>, anyhow::Error> {
         self.core.krate.paths.get(id).ok_or_else(|| {
             anyhow!(
                 "Failed to look up the type id `{}` in the rustdoc's path index for `{}`. \
