Auto merge of #164484 - TaKO8Ki:sparse-metadata-tables, r=<try>

Use adaptive sparse encoding for metadata tables

Author: bors

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -195,8 +195,14 @@ pub(super) trait LazyDecoder: BlobDecoder {
         self.read_lazy_offset_then(|pos| LazyArray::from_position_and_num_elems(pos, len))
     }
 
-    fn read_lazy_table<Ie, Id, T>(&mut self, width: usize, len: usize) -> LazyTable<Ie, Id, T> {
-        self.read_lazy_offset_then(|pos| LazyTable::from_position_and_encoded_size(pos, width, len))
+    fn read_lazy_table<Ie, Id, T>(
+        &mut self,
+        width_and_sparse: usize,
+        len: usize,
+    ) -> LazyTable<Ie, Id, T> {
+        self.read_lazy_offset_then(|pos| {
+            LazyTable::from_position_and_encoded_size(pos, width_and_sparse, len)
+        })
     }
 
     #[inline]
@@ -688,9 +694,9 @@ impl<D: LazyDecoder, T> Decodable<D> for LazyArray<T> {
 
 impl<Ie: Idx, Id: Idx, D: LazyDecoder, T> Decodable<D> for LazyTable<Ie, Id, T> {
     fn decode(decoder: &mut D) -> Self {
-        let width = decoder.read_usize();
+        let width_and_sparse = decoder.read_usize();
         let len = decoder.read_usize();
-        decoder.read_lazy_table(width, len)
+        decoder.read_lazy_table(width_and_sparse, len)
     }
 }
 

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -135,7 +135,7 @@ impl<'a, 'tcx, T> Encodable<EncodeContext<'a, 'tcx>> for LazyArray<T> {
 
 impl<'a, 'tcx, Ie, Id, T> Encodable<EncodeContext<'a, 'tcx>> for LazyTable<Ie, Id, T> {
     fn encode(&self, e: &mut EncodeContext<'a, 'tcx>) {
-        e.emit_usize(self.width);
+        e.emit_usize(self.width_and_sparse);
         e.emit_usize(self.len);
         e.emit_lazy_distance(self.position);
     }

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -66,7 +66,7 @@ pub(crate) fn rustc_version(cfg_version: &'static str) -> String {
 /// Metadata encoding version.
 /// N.B., increment this if you change the format of metadata such that
 /// the rustc version can't be found to compare with `rustc_version()`.
-const METADATA_VERSION: u8 = 11;
+const METADATA_VERSION: u8 = 12;
 
 /// Metadata header which includes `METADATA_VERSION`.
 ///
@@ -153,8 +153,9 @@ impl<T> LazyArray<T> {
 struct LazyTable<IdxEncode, IdxDecode, T> {
     position: NonZero<usize>,
     /// The encoded size of the elements of a table is selected at runtime to drop
-    /// trailing zeroes. This is the number of bytes used for each table element.
-    width: usize,
+    /// trailing zeroes. The low bit selects sparse encoding; the remaining bits
+    /// store the number of bytes used for each table element.
+    width_and_sparse: usize,
     /// How many elements are in the table.
     len: usize,
     _marker: PhantomData<fn(IdxEncode, IdxDecode) -> T>,
@@ -163,10 +164,10 @@ struct LazyTable<IdxEncode, IdxDecode, T> {
 impl<Ie, Id, T> LazyTable<Ie, Id, T> {
     fn from_position_and_encoded_size(
         position: NonZero<usize>,
-        width: usize,
+        width_and_sparse: usize,
         len: usize,
     ) -> LazyTable<Ie, Id, T> {
-        LazyTable { position, width, len, _marker: PhantomData }
+        LazyTable { position, width_and_sparse, len, _marker: PhantomData }
     }
 }
 

compiler/rustc_metadata/src/rmeta/table.rs

@@ -4,6 +4,9 @@ use rustc_index::Idx;
 use crate::rmeta::decoder::MetaBlob;
 use crate::rmeta::*;
 
+const SPARSE_BLOCK_LEN: usize = 64;
+const SPARSE_INDEX_BYTES: usize = 12;
+
 pub(super) trait IsDefault: Default {
     fn is_default(&self) -> bool;
 }
@@ -465,11 +468,6 @@ impl<Ie: Idx, Id: Idx, const N: usize, T: FixedSizeEncoding<ByteArray = [u8; N]>
             );
         }
         if !value.is_default() {
-            // FIXME(eddyb) investigate more compact encodings for sparse tables.
-            // On the PR @michaelwoerister mentioned:
-            // > Space requirements could perhaps be optimized by using the HAMT `popcnt`
-            // > trick (i.e. divide things into buckets of 32 or 64 items and then
-            // > store bit-masks of which item in each bucket is actually serialized).
             let block = self.blocks.ensure_contains_elem(i, || [0; N]);
             value.write_to_bytes(block);
             if self.width != N {
@@ -483,17 +481,45 @@ impl<Ie: Idx, Id: Idx, const N: usize, T: FixedSizeEncoding<ByteArray = [u8; N]>
         let pos = buf.position();
 
         let width = self.width;
+        let len = self.blocks.len();
+        let index_bytes = len.div_ceil(SPARSE_BLOCK_LEN) * SPARSE_INDEX_BYTES;
+        let non_defaults = if index_bytes < len * width {
+            self.blocks.iter().filter(|block| **block != [0; N]).count()
+        } else {
+            len
+        };
+        let sparse =
+            u32::try_from(non_defaults).is_ok() && index_bytes + non_defaults * width < len * width;
+
+        if sparse {
+            let mut base = 0u32;
+            for chunk in self.blocks.raw.chunks(SPARSE_BLOCK_LEN) {
+                let mut bitmap = 0u64;
+                for (i, block) in chunk.iter().enumerate() {
+                    if *block != [0; N] {
+                        bitmap |= 1 << i;
+                    }
+                }
+                buf.write_array(bitmap.to_le_bytes());
+                buf.write_array(base.to_le_bytes());
+                base += bitmap.count_ones();
+            }
+        }
+
         for block in &self.blocks {
-            buf.write_with(|dest| {
-                *dest = *block;
-                width
-            });
+            if !sparse || *block != [0; N] {
+                buf.write_with(|dest| {
+                    *dest = *block;
+                    width
+                });
+            }
         }
 
+        debug_assert!(width <= usize::MAX >> 1);
         LazyTable::from_position_and_encoded_size(
             NonZero::new(pos).unwrap(),
-            width,
-            self.blocks.len(),
+            (width << 1) | usize::from(sparse),
+            len,
         )
     }
 }
@@ -518,8 +544,23 @@ where
             return Default::default();
         }
 
-        let width = self.width;
-        let start = self.position.get() + (width * i.index());
+        let width = self.width_and_sparse >> 1;
+        let mut position = self.position.get();
+        let mut index = i.index();
+        if self.width_and_sparse & 1 != 0 {
+            let block = index / SPARSE_BLOCK_LEN;
+            let bit = 1u64 << (index % SPARSE_BLOCK_LEN);
+            let start = position + block * SPARSE_INDEX_BYTES;
+            let bitmap = u64::from_le_bytes(metadata.blob()[start..start + 8].try_into().unwrap());
+            if bitmap & bit == 0 {
+                return Default::default();
+            }
+            let base =
+                u32::from_le_bytes(metadata.blob()[start + 8..start + 12].try_into().unwrap());
+            index = base as usize + (bitmap & (bit - 1)).count_ones() as usize;
+            position += self.len.div_ceil(SPARSE_BLOCK_LEN) * SPARSE_INDEX_BYTES;
+        }
+        let start = position + width * index;
         let end = start + width;
         let bytes = &metadata.blob()[start..end];