From 51dfe498f14417f1b833361d4a298c6b4c0ff1cd Mon Sep 17 00:00:00 2001
From: Mark Rousskov <mark.simulacrum@gmail.com>
Date: Sat, 30 Dec 2023 18:22:58 -0500
Subject: [PATCH] Embed length of offset/position into Span tag byte

This cuts the average bytes/relative span from 3.5 to 3.2 on libcore,
ultimately saving ~400kb of data.
---
 compiler/rustc_metadata/src/rmeta/decoder.rs |  6 +++++-
 compiler/rustc_metadata/src/rmeta/encoder.rs | 20 ++++++++++++++++----
 compiler/rustc_metadata/src/rmeta/mod.rs     |  4 +++-
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/compiler/rustc_metadata/src/rmeta/decoder.rs b/compiler/rustc_metadata/src/rmeta/decoder.rs
index 49e849964be..eb0ac7fbf28 100644
--- a/compiler/rustc_metadata/src/rmeta/decoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/decoder.rs
@@ -502,7 +502,11 @@ impl<'a, 'tcx> SpanDecoder for DecodeContext<'a, 'tcx> {
         let data = if tag.kind() == SpanKind::Indirect {
             // Skip past the tag we just peek'd.
             self.read_u8();
-            let offset_or_position = self.read_usize();
+            // indirect tag lengths are safe to access, since they're (0, 8)
+            let bytes_needed = tag.length().unwrap().0 as usize;
+            let mut total = [0u8; usize::BITS as usize / 8];
+            total[..bytes_needed].copy_from_slice(self.read_raw_bytes(bytes_needed));
+            let offset_or_position = usize::from_le_bytes(total);
             let position = if tag.is_relative_offset() {
                 start - offset_or_position
             } else {
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index a458b528a97..e8237f2fc9f 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -172,11 +172,19 @@ impl<'a, 'tcx> SpanEncoder for EncodeContext<'a, 'tcx> {
                 // previously saved offset must be smaller than the current position.
                 let offset = self.opaque.position() - last_location;
                 if offset < last_location {
-                    SpanTag::indirect(true).encode(self);
-                    offset.encode(self);
+                    let needed = bytes_needed(offset);
+                    SpanTag::indirect(true, needed as u8).encode(self);
+                    self.opaque.write_with(|dest| {
+                        *dest = offset.to_le_bytes();
+                        needed
+                    });
                 } else {
-                    SpanTag::indirect(false).encode(self);
-                    last_location.encode(self);
+                    let needed = bytes_needed(last_location);
+                    SpanTag::indirect(false, needed as u8).encode(self);
+                    self.opaque.write_with(|dest| {
+                        *dest = last_location.to_le_bytes();
+                        needed
+                    });
                 }
             }
             Entry::Vacant(v) => {
@@ -212,6 +220,10 @@ impl<'a, 'tcx> SpanEncoder for EncodeContext<'a, 'tcx> {
     }
 }
 
+fn bytes_needed(n: usize) -> usize {
+    (usize::BITS - n.leading_zeros()).div_ceil(u8::BITS) as usize
+}
+
 impl<'a, 'tcx> Encodable<EncodeContext<'a, 'tcx>> for SpanData {
     fn encode(&self, s: &mut EncodeContext<'a, 'tcx>) {
         // Don't serialize any `SyntaxContext`s from a proc-macro crate,
diff --git a/compiler/rustc_metadata/src/rmeta/mod.rs b/compiler/rustc_metadata/src/rmeta/mod.rs
index 2f775882693..c02a1a6d1ca 100644
--- a/compiler/rustc_metadata/src/rmeta/mod.rs
+++ b/compiler/rustc_metadata/src/rmeta/mod.rs
@@ -529,11 +529,13 @@ impl SpanTag {
         SpanTag(data)
     }
 
-    fn indirect(relative: bool) -> SpanTag {
+    fn indirect(relative: bool, length_bytes: u8) -> SpanTag {
         let mut tag = SpanTag(SpanKind::Indirect as u8);
         if relative {
             tag.0 |= 0b100;
         }
+        assert!(length_bytes <= 8);
+        tag.0 |= length_bytes << 3;
         tag
     }