Auto merge of #55617 - oli-obk:stacker, r=nagisa,oli-obk

Prevent compiler stack overflow for deeply recursive code

I was unable to write a test that

1. runs in under 1s
2. overflows on my machine without this patch

The following reproduces the issue, but I don't think it's sensible to include a test that takes 30s to compile. We can now easily squash newly appearing overflows by the strategic insertion of calls to `ensure_sufficient_stack`.

```rust
// compile-pass

#![recursion_limit="1000000"]

macro_rules! chain {
    (EE $e:expr) => {$e.sin()};
    (RECURSE $i:ident $e:expr) => {chain!($i chain!($i chain!($i chain!($i $e))))};
    (Z $e:expr) => {chain!(RECURSE EE $e)};
    (Y $e:expr) => {chain!(RECURSE Z $e)};
    (X $e:expr) => {chain!(RECURSE Y $e)};
    (A $e:expr) => {chain!(RECURSE X $e)};
    (B $e:expr) => {chain!(RECURSE A $e)};
    (C $e:expr) => {chain!(RECURSE B $e)};
    // causes overflow on x86_64 linux
    // less than 1 second until overflow on test machine
    // after overflow has been fixed, takes 30s to compile :/
    (D $e:expr) => {chain!(RECURSE C $e)};
    (E $e:expr) => {chain!(RECURSE D $e)};
    (F $e:expr) => {chain!(RECURSE E $e)};
    // more than 10 seconds
    (G $e:expr) => {chain!(RECURSE F $e)};
    (H $e:expr) => {chain!(RECURSE G $e)};
    (I $e:expr) => {chain!(RECURSE H $e)};
    (J $e:expr) => {chain!(RECURSE I $e)};
    (K $e:expr) => {chain!(RECURSE J $e)};
    (L $e:expr) => {chain!(RECURSE L $e)};
}

fn main() {
    let x = chain!(D 42.0_f32);
}
```

fixes #55471
fixes #41884
fixes #40161
fixes #34844
fixes #32594

cc @alexcrichton @rust-lang/compiler

I looked at all code that checks the recursion limit and inserted stack growth calls where appropriate.

Cargo.lock

@@ -2639,6 +2639,15 @@ dependencies = [
  "core",
 ]
 
+[[package]]
+name = "psm"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "659ecfea2142a458893bb7673134bad50b752fea932349c213d6a23874ce3aa7"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "publicsuffix"
 version = "1.5.3"
@@ -3708,6 +3717,7 @@ dependencies = [
  "serialize",
  "smallvec 1.4.0",
  "stable_deref_trait",
+ "stacker",
  "winapi 0.3.8",
 ]
 
@@ -4669,6 +4679,19 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ffbc596e092fe5f598b12ef46cc03754085ac2f4d8c739ad61c4ae266cc3b3fa"
 
+[[package]]
+name = "stacker"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32c2467b8abbb417e4e62fd62229719b9c9d77714a7fa989f1afad16ba9c9743"
+dependencies = [
+ "cc",
+ "cfg-if",
+ "libc",
+ "psm",
+ "winapi 0.3.8",
+]
+
 [[package]]
 name = "std"
 version = "0.0.0"

src/librustc_ast_lowering/pat.rs

@@ -2,83 +2,89 @@ use super::{ImplTraitContext, LoweringContext, ParamMode};
 
 use rustc_ast::ast::*;
 use rustc_ast::ptr::P;
+use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir as hir;
 use rustc_hir::def::Res;
 use rustc_span::{source_map::Spanned, Span};
 
 impl<'a, 'hir> LoweringContext<'a, 'hir> {
     crate fn lower_pat(&mut self, p: &Pat) -> &'hir hir::Pat<'hir> {
-        let node = match p.kind {
-            PatKind::Wild => hir::PatKind::Wild,
-            PatKind::Ident(ref binding_mode, ident, ref sub) => {
-                let lower_sub = |this: &mut Self| sub.as_ref().map(|s| this.lower_pat(&*s));
-                let node = self.lower_pat_ident(p, binding_mode, ident, lower_sub);
-                node
-            }
-            PatKind::Lit(ref e) => hir::PatKind::Lit(self.lower_expr(e)),
-            PatKind::TupleStruct(ref path, ref pats) => {
-                let qpath = self.lower_qpath(
-                    p.id,
-                    &None,
-                    path,
-                    ParamMode::Optional,
-                    ImplTraitContext::disallowed(),
-                );
-                let (pats, ddpos) = self.lower_pat_tuple(pats, "tuple struct");
-                hir::PatKind::TupleStruct(qpath, pats, ddpos)
-            }
-            PatKind::Or(ref pats) => {
-                hir::PatKind::Or(self.arena.alloc_from_iter(pats.iter().map(|x| self.lower_pat(x))))
-            }
-            PatKind::Path(ref qself, ref path) => {
-                let qpath = self.lower_qpath(
-                    p.id,
-                    qself,
-                    path,
-                    ParamMode::Optional,
-                    ImplTraitContext::disallowed(),
-                );
-                hir::PatKind::Path(qpath)
-            }
-            PatKind::Struct(ref path, ref fields, etc) => {
-                let qpath = self.lower_qpath(
-                    p.id,
-                    &None,
-                    path,
-                    ParamMode::Optional,
-                    ImplTraitContext::disallowed(),
-                );
+        ensure_sufficient_stack(|| {
+            let node = match p.kind {
+                PatKind::Wild => hir::PatKind::Wild,
+                PatKind::Ident(ref binding_mode, ident, ref sub) => {
+                    let lower_sub = |this: &mut Self| sub.as_ref().map(|s| this.lower_pat(&*s));
+                    let node = self.lower_pat_ident(p, binding_mode, ident, lower_sub);
+                    node
+                }
+                PatKind::Lit(ref e) => hir::PatKind::Lit(self.lower_expr(e)),
+                PatKind::TupleStruct(ref path, ref pats) => {
+                    let qpath = self.lower_qpath(
+                        p.id,
+                        &None,
+                        path,
+                        ParamMode::Optional,
+                        ImplTraitContext::disallowed(),
+                    );
+                    let (pats, ddpos) = self.lower_pat_tuple(pats, "tuple struct");
+                    hir::PatKind::TupleStruct(qpath, pats, ddpos)
+                }
+                PatKind::Or(ref pats) => hir::PatKind::Or(
+                    self.arena.alloc_from_iter(pats.iter().map(|x| self.lower_pat(x))),
+                ),
+                PatKind::Path(ref qself, ref path) => {
+                    let qpath = self.lower_qpath(
+                        p.id,
+                        qself,
+                        path,
+                        ParamMode::Optional,
+                        ImplTraitContext::disallowed(),
+                    );
+                    hir::PatKind::Path(qpath)
+                }
+                PatKind::Struct(ref path, ref fields, etc) => {
+                    let qpath = self.lower_qpath(
+                        p.id,
+                        &None,
+                        path,
+                        ParamMode::Optional,
+                        ImplTraitContext::disallowed(),
+                    );
 
-                let fs = self.arena.alloc_from_iter(fields.iter().map(|f| hir::FieldPat {
-                    hir_id: self.next_id(),
-                    ident: f.ident,
-                    pat: self.lower_pat(&f.pat),
-                    is_shorthand: f.is_shorthand,
-                    span: f.span,
-                }));
-                hir::PatKind::Struct(qpath, fs, etc)
-            }
-            PatKind::Tuple(ref pats) => {
-                let (pats, ddpos) = self.lower_pat_tuple(pats, "tuple");
-                hir::PatKind::Tuple(pats, ddpos)
-            }
-            PatKind::Box(ref inner) => hir::PatKind::Box(self.lower_pat(inner)),
-            PatKind::Ref(ref inner, mutbl) => hir::PatKind::Ref(self.lower_pat(inner), mutbl),
-            PatKind::Range(ref e1, ref e2, Spanned { node: ref end, .. }) => hir::PatKind::Range(
-                e1.as_deref().map(|e| self.lower_expr(e)),
-                e2.as_deref().map(|e| self.lower_expr(e)),
-                self.lower_range_end(end, e2.is_some()),
-            ),
-            PatKind::Slice(ref pats) => self.lower_pat_slice(pats),
-            PatKind::Rest => {
-                // If we reach here the `..` pattern is not semantically allowed.
-                self.ban_illegal_rest_pat(p.span)
-            }
-            PatKind::Paren(ref inner) => return self.lower_pat(inner),
-            PatKind::MacCall(_) => panic!("Shouldn't exist here"),
-        };
+                    let fs = self.arena.alloc_from_iter(fields.iter().map(|f| hir::FieldPat {
+                        hir_id: self.next_id(),
+                        ident: f.ident,
+                        pat: self.lower_pat(&f.pat),
+                        is_shorthand: f.is_shorthand,
+                        span: f.span,
+                    }));
+                    hir::PatKind::Struct(qpath, fs, etc)
+                }
+                PatKind::Tuple(ref pats) => {
+                    let (pats, ddpos) = self.lower_pat_tuple(pats, "tuple");
+                    hir::PatKind::Tuple(pats, ddpos)
+                }
+                PatKind::Box(ref inner) => hir::PatKind::Box(self.lower_pat(inner)),
+                PatKind::Ref(ref inner, mutbl) => hir::PatKind::Ref(self.lower_pat(inner), mutbl),
+                PatKind::Range(ref e1, ref e2, Spanned { node: ref end, .. }) => {
+                    hir::PatKind::Range(
+                        e1.as_deref().map(|e| self.lower_expr(e)),
+                        e2.as_deref().map(|e| self.lower_expr(e)),
+                        self.lower_range_end(end, e2.is_some()),
+                    )
+                }
+                PatKind::Slice(ref pats) => self.lower_pat_slice(pats),
+                PatKind::Rest => {
+                    // If we reach here the `..` pattern is not semantically allowed.
+                    self.ban_illegal_rest_pat(p.span)
+                }
+                // FIXME: consider not using recursion to lower this.
+                PatKind::Paren(ref inner) => return self.lower_pat(inner),
+                PatKind::MacCall(_) => panic!("{:?} shouldn't exist here", p.span),
+            };
 
-        self.pat_with_node_id_of(p, node)
+            self.pat_with_node_id_of(p, node)
+        })
     }
 
     fn lower_pat_tuple(

src/librustc_data_structures/Cargo.toml

@@ -28,6 +28,7 @@ rustc_index = { path = "../librustc_index", package = "rustc_index" }
 bitflags = "1.2.1"
 measureme = "0.7.1"
 libc = "0.2"
+stacker = "0.1.6"
 
 [dependencies.parking_lot]
 version = "0.10"

src/librustc_data_structures/lib.rs

@@ -80,6 +80,7 @@ pub mod stable_set;
 #[macro_use]
 pub mod stable_hasher;
 pub mod sharded;
+pub mod stack;
 pub mod sync;
 pub mod thin_vec;
 pub mod tiny_list;

src/librustc_data_structures/stack.rs

@@ -0,0 +1,17 @@
+// This is the amount of bytes that need to be left on the stack before increasing the size.
+// It must be at least as large as the stack required by any code that does not call
+// `ensure_sufficient_stack`.
+const RED_ZONE: usize = 100 * 1024; // 100k
+
+// Only the first stack that is pushed, grows exponentially (2^n * STACK_PER_RECURSION) from then
+// on. This flag has performance relevant characteristics. Don't set it too high.
+const STACK_PER_RECURSION: usize = 1 * 1024 * 1024; // 1MB
+
+/// Grows the stack on demand to prevent stack overflow. Call this in strategic locations
+/// to "break up" recursive calls. E.g. almost any call to `visit_expr` or equivalent can benefit
+/// from this.
+///
+/// Should not be sprinkled around carelessly, as it causes a little bit of overhead.
+pub fn ensure_sufficient_stack<R>(f: impl FnOnce() -> R) -> R {
+    stacker::maybe_grow(RED_ZONE, STACK_PER_RECURSION, f)
+}

src/librustc_interface/util.rs

@@ -81,14 +81,7 @@ pub fn create_session(
     (Lrc::new(sess), Lrc::new(codegen_backend), source_map)
 }
 
-// Temporarily have stack size set to 32MB to deal with various crates with long method
-// chains or deep syntax trees, except when on Haiku.
-// FIXME(oli-obk): get https://github.com/rust-lang/rust/pull/55617 the finish line
-#[cfg(not(target_os = "haiku"))]
-const STACK_SIZE: usize = 32 * 1024 * 1024;
-
-#[cfg(target_os = "haiku")]
-const STACK_SIZE: usize = 16 * 1024 * 1024;
+const STACK_SIZE: usize = 8 * 1024 * 1024;
 
 fn get_stack_size() -> Option<usize> {
     // FIXME: Hacks on hacks. If the env is trying to override the stack size

src/librustc_middle/ty/inhabitedness/mod.rs

@@ -6,6 +6,7 @@ use crate::ty::TyKind::*;
 use crate::ty::{AdtDef, FieldDef, Ty, TyS, VariantDef};
 use crate::ty::{AdtKind, Visibility};
 use crate::ty::{DefId, SubstsRef};
+use rustc_data_structures::stack::ensure_sufficient_stack;
 
 mod def_id_forest;
 
@@ -196,7 +197,9 @@ impl<'tcx> TyS<'tcx> {
     /// Calculates the forest of `DefId`s from which this type is visibly uninhabited.
     fn uninhabited_from(&self, tcx: TyCtxt<'tcx>, param_env: ty::ParamEnv<'tcx>) -> DefIdForest {
         match self.kind {
-            Adt(def, substs) => def.uninhabited_from(tcx, substs, param_env),
+            Adt(def, substs) => {
+                ensure_sufficient_stack(|| def.uninhabited_from(tcx, substs, param_env))
+            }
 
             Never => DefIdForest::full(tcx),
 

src/librustc_middle/ty/query/plumbing.rs

@@ -68,7 +68,9 @@ impl QueryContext for TyCtxt<'tcx> {
             };
 
             // Use the `ImplicitCtxt` while we execute the query.
-            tls::enter_context(&new_icx, |_| compute(*self))
+            tls::enter_context(&new_icx, |_| {
+                rustc_data_structures::stack::ensure_sufficient_stack(|| compute(*self))
+            })
         })
     }
 }

src/librustc_mir/monomorphize/collector.rs

@@ -369,7 +369,9 @@ fn collect_items_rec<'tcx>(
             recursion_depth_reset = Some(check_recursion_limit(tcx, instance, recursion_depths));
             check_type_length_limit(tcx, instance);
 
-            collect_neighbours(tcx, instance, &mut neighbors);
+            rustc_data_structures::stack::ensure_sufficient_stack(|| {
+                collect_neighbours(tcx, instance, &mut neighbors);
+            });
         }
         MonoItem::GlobalAsm(..) => {
             recursion_depth_reset = None;
@@ -1146,7 +1148,9 @@ fn collect_miri<'tcx>(tcx: TyCtxt<'tcx>, alloc_id: AllocId, output: &mut Vec<Mon
         Some(GlobalAlloc::Memory(alloc)) => {
             trace!("collecting {:?} with {:#?}", alloc_id, alloc);
             for &((), inner) in alloc.relocations().values() {
-                collect_miri(tcx, inner, output);
+                rustc_data_structures::stack::ensure_sufficient_stack(|| {
+                    collect_miri(tcx, inner, output);
+                });
             }
         }
         Some(GlobalAlloc::Function(fn_instance)) => {

src/librustc_mir_build/build/expr/as_temp.rs

@@ -3,6 +3,7 @@
 use crate::build::scope::DropKind;
 use crate::build::{BlockAnd, BlockAndExtension, Builder};
 use crate::hair::*;
+use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir as hir;
 use rustc_middle::middle::region;
 use rustc_middle::mir::*;
@@ -21,7 +22,11 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         M: Mirror<'tcx, Output = Expr<'tcx>>,
     {
         let expr = self.hir.mirror(expr);
-        self.expr_as_temp(block, temp_lifetime, expr, mutability)
+        //
+        // this is the only place in mir building that we need to truly need to worry about
+        // infinite recursion. Everything else does recurse, too, but it always gets broken up
+        // at some point by inserting an intermediate temporary
+        ensure_sufficient_stack(|| self.expr_as_temp(block, temp_lifetime, expr, mutability))
     }
 
     fn expr_as_temp(

src/librustc_trait_selection/traits/project.rs

@@ -18,6 +18,7 @@ use crate::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
 use crate::infer::{InferCtxt, InferOk, LateBoundRegionConversionTime};
 use crate::traits::error_reporting::InferCtxtExt;
 use rustc_ast::ast::Ident;
+use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_errors::ErrorReported;
 use rustc_hir::def_id::DefId;
 use rustc_middle::ty::fold::{TypeFoldable, TypeFolder};
@@ -261,7 +262,7 @@ where
 {
     debug!("normalize_with_depth(depth={}, value={:?})", depth, value);
     let mut normalizer = AssocTypeNormalizer::new(selcx, param_env, cause, depth, obligations);
-    let result = normalizer.fold(value);
+    let result = ensure_sufficient_stack(|| normalizer.fold(value));
     debug!(
         "normalize_with_depth: depth={} result={:?} with {} obligations",
         depth,

src/librustc_trait_selection/traits/query/normalize.rs

@@ -7,6 +7,7 @@ use crate::infer::canonical::OriginalQueryValues;
 use crate::infer::{InferCtxt, InferOk};
 use crate::traits::error_reporting::InferCtxtExt;
 use crate::traits::{Obligation, ObligationCause, PredicateObligation, Reveal};
+use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_infer::traits::Normalized;
 use rustc_middle::ty::fold::{TypeFoldable, TypeFolder};
 use rustc_middle::ty::subst::Subst;
@@ -131,7 +132,7 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for QueryNormalizer<'cx, 'tcx> {
                                 ty
                             );
                         }
-                        let folded_ty = self.fold_ty(concrete_ty);
+                        let folded_ty = ensure_sufficient_stack(|| self.fold_ty(concrete_ty));
                         self.anon_depth -= 1;
                         folded_ty
                     }