Query-ify global limit attribute handling

2 files changed, 23 insertions, 20 deletions

diff --git a/compiler/rustc_middle/src/middle/limits.rs b/compiler/rustc_middle/src/middle/limits.rs
index 601198fd0de..eb706b57004 100644
--- a/compiler/rustc_middle/src/middle/limits.rs
+++ b/compiler/rustc_middle/src/middle/limits.rs
@@ -10,38 +10,37 @@
 //! just peeks and looks for that attribute.
 
 use crate::bug;
-use rustc_ast as ast;
-use rustc_data_structures::sync::OnceCell;
+use crate::ty;
+use rustc_ast::Attribute;
+use rustc_session::Limit;
 use rustc_session::Session;
 use rustc_span::symbol::{sym, Symbol};
 
 use std::num::IntErrorKind;
 
-pub fn update_limits(sess: &Session, krate: &ast::Crate) {
-    update_limit(sess, krate, &sess.recursion_limit, sym::recursion_limit, 128);
-    update_limit(sess, krate, &sess.move_size_limit, sym::move_size_limit, 0);
-    update_limit(sess, krate, &sess.type_length_limit, sym::type_length_limit, 1048576);
-    update_limit(sess, krate, &sess.const_eval_limit, sym::const_eval_limit, 1_000_000);
+pub fn provide(providers: &mut ty::query::Providers) {
+    providers.recursion_limit = |tcx, ()| get_recursion_limit(tcx.hir().krate_attrs(), tcx.sess);
+    providers.move_size_limit =
+        |tcx, ()| get_limit(tcx.hir().krate_attrs(), tcx.sess, sym::move_size_limit, 0).0;
+    providers.type_length_limit =
+        |tcx, ()| get_limit(tcx.hir().krate_attrs(), tcx.sess, sym::type_length_limit, 1048576);
+    providers.const_eval_limit =
+        |tcx, ()| get_limit(tcx.hir().krate_attrs(), tcx.sess, sym::const_eval_limit, 1_000_000);
 }
 
-fn update_limit(
-    sess: &Session,
-    krate: &ast::Crate,
-    limit: &OnceCell<impl From<usize> + std::fmt::Debug>,
-    name: Symbol,
-    default: usize,
-) {
-    for attr in &krate.attrs {
+pub fn get_recursion_limit(krate_attrs: &[Attribute], sess: &Session) -> Limit {
+    get_limit(krate_attrs, sess, sym::recursion_limit, 128)
+}
+
+fn get_limit(krate_attrs: &[Attribute], sess: &Session, name: Symbol, default: usize) -> Limit {
+    for attr in krate_attrs {
         if !sess.check_name(attr, name) {
             continue;
         }
 
         if let Some(s) = attr.value_str() {
             match s.as_str().parse() {
-                Ok(n) => {
-                    limit.set(From::from(n)).unwrap();
-                    return;
-                }
+                Ok(n) => return Limit::new(n),
                 Err(e) => {
                     let mut err =
                         sess.struct_span_err(attr.span, "`limit` must be a non-negative integer");
@@ -68,5 +67,5 @@ fn update_limit(
             }
         }
     }
-    limit.set(From::from(default)).unwrap();
+    return Limit::new(default);
 }
diff --git a/compiler/rustc_middle/src/middle/mod.rs b/compiler/rustc_middle/src/middle/mod.rs
index a369e85306b..b370ec152e8 100644
--- a/compiler/rustc_middle/src/middle/mod.rs
+++ b/compiler/rustc_middle/src/middle/mod.rs
@@ -32,3 +32,7 @@ pub mod privacy;
 pub mod region;
 pub mod resolve_lifetime;
 pub mod stability;
+
+pub fn provide(providers: &mut crate::ty::query::Providers) {
+    limits::provide(providers);
+}