From 529a3f4ce609f4bc98215185838351d9815fbdc1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9my=20Rakic?= <remy.rakic+github@gmail.com>
Date: Wed, 3 Jul 2024 19:36:14 +0000
Subject: [PATCH] cache type sizes in type-size limit visitor

---
 compiler/rustc_middle/src/ty/instance.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/compiler/rustc_middle/src/ty/instance.rs b/compiler/rustc_middle/src/ty/instance.rs
index ae54411d788ab..c50a98e88fd63 100644
--- a/compiler/rustc_middle/src/ty/instance.rs
+++ b/compiler/rustc_middle/src/ty/instance.rs
@@ -5,6 +5,7 @@ use crate::ty::{
     self, EarlyBinder, GenericArgs, GenericArgsRef, Ty, TyCtxt, TypeFoldable, TypeSuperFoldable,
     TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
 };
+use rustc_data_structures::fx::FxHashMap;
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir as hir;
 use rustc_hir::def::Namespace;
@@ -388,13 +389,25 @@ impl<'tcx> InstanceKind<'tcx> {
 }
 
 fn type_length<'tcx>(item: impl TypeVisitable<TyCtxt<'tcx>>) -> usize {
-    struct Visitor {
+    struct Visitor<'tcx> {
         type_length: usize,
+        cache: FxHashMap<Ty<'tcx>, usize>,
     }
-    impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for Visitor {
+    impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for Visitor<'tcx> {
         fn visit_ty(&mut self, t: Ty<'tcx>) {
+            if let Some(&value) = self.cache.get(&t) {
+                self.type_length += value;
+                return;
+            }
+
+            let prev = self.type_length;
             self.type_length += 1;
             t.super_visit_with(self);
+
+            // We don't try to use the cache if the type is fairly small.
+            if self.type_length > 16 {
+                self.cache.insert(t, self.type_length - prev);
+            }
         }
 
         fn visit_const(&mut self, ct: ty::Const<'tcx>) {
@@ -402,7 +415,7 @@ fn type_length<'tcx>(item: impl TypeVisitable<TyCtxt<'tcx>>) -> usize {
             ct.super_visit_with(self);
         }
     }
-    let mut visitor = Visitor { type_length: 0 };
+    let mut visitor = Visitor { type_length: 0, cache: Default::default() };
     item.visit_with(&mut visitor);
 
     visitor.type_length