From 674735b109cdfb815c02b9b60fb65e5495f7cfe1 Mon Sep 17 00:00:00 2001
From: Jack Huey <jack.huey@umassmed.edu>
Date: Sun, 17 Jan 2021 02:49:30 -0500
Subject: [PATCH] Impl EncodableWithShorthand for PredicateKind

---
 compiler/rustc_metadata/src/rmeta/encoder.rs  |  6 ++++
 compiler/rustc_middle/src/ty/codec.rs         | 31 +++++++++++++++++++
 compiler/rustc_middle/src/ty/mod.rs           |  2 +-
 .../src/ty/query/on_disk_cache.rs             |  5 +++
 4 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index 0b398c35db0c..5e2674254b29 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -46,6 +46,7 @@ pub(super) struct EncodeContext<'a, 'tcx> {
 
     lazy_state: LazyState,
     type_shorthands: FxHashMap<Ty<'tcx>, usize>,
+    predicate_shorthands: FxHashMap<ty::PredicateKind<'tcx>, usize>,
 
     interpret_allocs: FxIndexSet<interpret::AllocId>,
 
@@ -327,6 +328,10 @@ impl<'a, 'tcx> TyEncoder<'tcx> for EncodeContext<'a, 'tcx> {
         &mut self.type_shorthands
     }
 
+    fn predicate_shorthands(&mut self) -> &mut FxHashMap<ty::PredicateKind<'tcx>, usize> {
+        &mut self.predicate_shorthands
+    }
+
     fn encode_alloc_id(
         &mut self,
         alloc_id: &rustc_middle::mir::interpret::AllocId,
@@ -2146,6 +2151,7 @@ fn encode_metadata_impl(tcx: TyCtxt<'_>) -> EncodedMetadata {
         tables: Default::default(),
         lazy_state: LazyState::NoNode,
         type_shorthands: Default::default(),
+        predicate_shorthands: Default::default(),
         source_file_cache,
         interpret_allocs: Default::default(),
         required_source_files,
diff --git a/compiler/rustc_middle/src/ty/codec.rs b/compiler/rustc_middle/src/ty/codec.rs
index 380fb8172a75..0aaba81819b0 100644
--- a/compiler/rustc_middle/src/ty/codec.rs
+++ b/compiler/rustc_middle/src/ty/codec.rs
@@ -43,11 +43,21 @@ impl<'tcx, E: TyEncoder<'tcx>> EncodableWithShorthand<'tcx, E> for Ty<'tcx> {
     }
 }
 
+impl<'tcx, E: TyEncoder<'tcx>> EncodableWithShorthand<'tcx, E> for ty::PredicateKind<'tcx> {
+    type Variant = ty::PredicateKind<'tcx>;
+
+    #[inline]
+    fn variant(&self) -> &Self::Variant {
+        self
+    }
+}
+
 pub trait TyEncoder<'tcx>: Encoder {
     const CLEAR_CROSS_CRATE: bool;
 
     fn position(&self) -> usize;
     fn type_shorthands(&mut self) -> &mut FxHashMap<Ty<'tcx>, usize>;
+    fn predicate_shorthands(&mut self) -> &mut FxHashMap<ty::PredicateKind<'tcx>, usize>;
     fn encode_alloc_id(&mut self, alloc_id: &AllocId) -> Result<(), Self::Error>;
 }
 
@@ -110,6 +120,12 @@ impl<'tcx, E: TyEncoder<'tcx>> Encodable<E> for Ty<'tcx> {
     }
 }
 
+impl<'tcx, E: TyEncoder<'tcx>> Encodable<E> for ty::PredicateKind<'tcx> {
+    fn encode(&self, e: &mut E) -> Result<(), E::Error> {
+        encode_with_shorthand(e, self, TyEncoder::predicate_shorthands)
+    }
+}
+
 impl<'tcx, E: TyEncoder<'tcx>> Encodable<E> for ty::Predicate<'tcx> {
     fn encode(&self, e: &mut E) -> Result<(), E::Error> {
         self.kind().encode(e)
@@ -210,6 +226,21 @@ impl<'tcx, D: TyDecoder<'tcx>> Decodable<D> for Ty<'tcx> {
     }
 }
 
+impl<'tcx, D: TyDecoder<'tcx>> Decodable<D> for ty::PredicateKind<'tcx> {
+    fn decode(decoder: &mut D) -> Result<ty::PredicateKind<'tcx>, D::Error> {
+        // Handle shorthands first, if we have an usize > 0x80.
+        if decoder.positioned_at_shorthand() {
+            let pos = decoder.read_usize()?;
+            assert!(pos >= SHORTHAND_OFFSET);
+            let shorthand = pos - SHORTHAND_OFFSET;
+
+            decoder.with_position(shorthand, ty::PredicateKind::decode)
+        } else {
+            Ok(ty::PredicateKind::decode(decoder)?)
+        }
+    }
+}
+
 impl<'tcx, D: TyDecoder<'tcx>> Decodable<D> for ty::Predicate<'tcx> {
     fn decode(decoder: &mut D) -> Result<ty::Predicate<'tcx>, D::Error> {
         let predicate_kind = Decodable::decode(decoder)?;
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index 099c5aa84e51..88a2aac010c9 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -1081,7 +1081,7 @@ impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for Predicate<'tcx> {
     }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq, Hash, TyEncodable, TyDecodable)]
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
 #[derive(HashStable, TypeFoldable)]
 pub enum PredicateKind<'tcx> {
     /// Corresponds to `where Foo: Bar<A, B, C>`. `Foo` here would be
diff --git a/compiler/rustc_middle/src/ty/query/on_disk_cache.rs b/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
index 914937134b47..cfe47004e01b 100644
--- a/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
+++ b/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
@@ -293,6 +293,7 @@ impl<'sess> OnDiskCache<'sess> {
                 tcx,
                 encoder,
                 type_shorthands: Default::default(),
+                predicate_shorthands: Default::default(),
                 interpret_allocs: Default::default(),
                 source_map: CachingSourceMapView::new(tcx.sess.source_map()),
                 file_to_file_index,
@@ -988,6 +989,7 @@ struct CacheEncoder<'a, 'tcx, E: OpaqueEncoder> {
     tcx: TyCtxt<'tcx>,
     encoder: &'a mut E,
     type_shorthands: FxHashMap<Ty<'tcx>, usize>,
+    predicate_shorthands: FxHashMap<ty::PredicateKind<'tcx>, usize>,
     interpret_allocs: FxIndexSet<interpret::AllocId>,
     source_map: CachingSourceMapView<'tcx>,
     file_to_file_index: FxHashMap<*const SourceFile, SourceFileIndex>,
@@ -1101,6 +1103,9 @@ where
     fn type_shorthands(&mut self) -> &mut FxHashMap<Ty<'tcx>, usize> {
         &mut self.type_shorthands
     }
+    fn predicate_shorthands(&mut self) -> &mut FxHashMap<ty::PredicateKind<'tcx>, usize> {
+        &mut self.predicate_shorthands
+    }
     fn encode_alloc_id(&mut self, alloc_id: &interpret::AllocId) -> Result<(), Self::Error> {
         let (index, _) = self.interpret_allocs.insert_full(*alloc_id);