Arena-allocate PatStack

Nadrieril · Nadrieril · commit 98a9d5c30636 · 2023-12-01T08:31:23.000+01:00
diff --git a/compiler/rustc_mir_build/src/thir/pattern/check_match.rs b/compiler/rustc_mir_build/src/thir/pattern/check_match.rs
@@ -5,7 +5,7 @@ use super::usefulness::{
 
 use crate::errors::*;
 
-use rustc_arena::TypedArena;
+use rustc_arena::{DroplessArena, TypedArena};
 use rustc_ast::Mutability;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::stack::ensure_sufficient_stack;
@@ -29,13 +29,15 @@ pub(crate) fn check_match(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Result<(), Err
     let (thir, expr) = tcx.thir_body(def_id)?;
     let thir = thir.borrow();
     let pattern_arena = TypedArena::default();
+    let dropless_arena = DroplessArena::default();
     let mut visitor = MatchVisitor {
         tcx,
         thir: &*thir,
         param_env: tcx.param_env(def_id),
         lint_level: tcx.local_def_id_to_hir_id(def_id),
         let_source: LetSource::None,
         pattern_arena: &pattern_arena,
+        dropless_arena: &dropless_arena,
         error: Ok(()),
     };
     visitor.visit_expr(&thir[expr]);
@@ -80,6 +82,7 @@ struct MatchVisitor<'thir, 'p, 'tcx> {
     lint_level: HirId,
     let_source: LetSource,
     pattern_arena: &'p TypedArena<DeconstructedPat<'p, 'tcx>>,
+    dropless_arena: &'p DroplessArena,
     /// Tracks if we encountered an error while checking this body. That the first function to
     /// report it stores it here. Some functions return `Result` to allow callers to short-circuit
     /// on error, but callers don't need to store it here again.
@@ -302,6 +305,7 @@ impl<'thir, 'p, 'tcx> MatchVisitor<'thir, 'p, 'tcx> {
             param_env: self.param_env,
             module: self.tcx.parent_module(self.lint_level).to_def_id(),
             pattern_arena: self.pattern_arena,
+            dropless_arena: self.dropless_arena,
             match_lint_level: self.lint_level,
             match_span,
             scrut_span,
diff --git a/compiler/rustc_mir_build/src/thir/pattern/usefulness.rs b/compiler/rustc_mir_build/src/thir/pattern/usefulness.rs
@@ -499,7 +499,7 @@ use crate::errors::{
 
 use rustc_data_structures::captures::Captures;
 
-use rustc_arena::TypedArena;
+use rustc_arena::{DroplessArena, TypedArena};
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir::def_id::DefId;
 use rustc_hir::HirId;
@@ -508,8 +508,8 @@ use rustc_session::lint;
 use rustc_session::lint::builtin::NON_EXHAUSTIVE_OMITTED_PATTERNS;
 use rustc_span::{Span, DUMMY_SP};
 
-use smallvec::{smallvec, SmallVec};
-use std::fmt;
+use smallvec::SmallVec;
+use std::{fmt, iter::once};
 
 pub(crate) struct MatchCheckCtxt<'p, 'tcx> {
     pub(crate) tcx: TyCtxt<'tcx>,
@@ -521,6 +521,7 @@ pub(crate) struct MatchCheckCtxt<'p, 'tcx> {
     pub(crate) module: DefId,
     pub(crate) param_env: ty::ParamEnv<'tcx>,
     pub(crate) pattern_arena: &'p TypedArena<DeconstructedPat<'p, 'tcx>>,
+    pub(crate) dropless_arena: &'p DroplessArena,
     /// Lint level at the match.
     pub(crate) match_lint_level: HirId,
     /// The span of the whole match, if applicable.
@@ -570,13 +571,12 @@ impl<'a, 'p, 'tcx> fmt::Debug for PatCtxt<'a, 'p, 'tcx> {
 /// Represents a pattern-tuple under investigation.
 #[derive(Clone)]
 struct PatStack<'p, 'tcx> {
-    // Rows of len 1 are very common, which is why `SmallVec[_; 2]` works well.
-    pats: SmallVec<[&'p DeconstructedPat<'p, 'tcx>; 2]>,
+    pats: &'p [&'p DeconstructedPat<'p, 'tcx>],
 }
 
 impl<'p, 'tcx> PatStack<'p, 'tcx> {
-    fn from_pattern(pat: &'p DeconstructedPat<'p, 'tcx>) -> Self {
-        PatStack { pats: smallvec![pat] }
+    fn from_pattern(cx: &MatchCheckCtxt<'p, '_>, pat: &'p DeconstructedPat<'p, 'tcx>) -> Self {
+        PatStack { pats: cx.dropless_arena.alloc_from_iter(once(pat)) }
     }
 
     fn is_empty(&self) -> bool {
@@ -597,12 +597,14 @@ impl<'p, 'tcx> PatStack<'p, 'tcx> {
 
     // Recursively expand the first or-pattern into its subpatterns. Only useful if the pattern is
     // an or-pattern. Panics if `self` is empty.
-    fn expand_or_pat<'a>(&'a self) -> impl Iterator<Item = PatStack<'p, 'tcx>> + Captures<'a> {
-        self.head().flatten_or_pat().into_iter().map(move |pat| {
-            let mut new_pats = SmallVec::<[_; 2]>::with_capacity(self.len());
-            new_pats.push(pat);
-            new_pats.extend_from_slice(&self.pats[1..]);
-            PatStack { pats: new_pats }
+    fn expand_or_pat<'a>(
+        &'a self,
+        cx: &'a MatchCheckCtxt<'p, 'tcx>,
+    ) -> impl Iterator<Item = PatStack<'p, 'tcx>> + Captures<'a> {
+        self.head().flatten_or_pat().into_iter().map(move |pat| PatStack {
+            pats: cx
+                .dropless_arena
+                .alloc_from_iter(once(pat).chain(self.pats[1..].iter().copied())),
         })
     }
 
@@ -616,10 +618,12 @@ impl<'p, 'tcx> PatStack<'p, 'tcx> {
         // We pop the head pattern and push the new fields extracted from the arguments of
         // `self.head()`.
         let ctor_fields = self.head().specialize(pcx, ctor);
-        let mut new_pats = SmallVec::<[_; 2]>::with_capacity(self.len() + ctor_fields.len() - 1);
-        new_pats.extend_from_slice(&ctor_fields);
-        new_pats.extend_from_slice(&self.pats[1..]);
-        PatStack { pats: new_pats }
+        PatStack {
+            pats: pcx
+                .cx
+                .dropless_arena
+                .alloc_from_iter(ctor_fields.into_iter().chain(self.pats[1..].iter().copied())),
+        }
     }
 }
 
@@ -670,8 +674,11 @@ impl<'p, 'tcx> MatrixRow<'p, 'tcx> {
 
     // Recursively expand the first or-pattern into its subpatterns. Only useful if the pattern is
     // an or-pattern. Panics if `self` is empty.
-    fn expand_or_pat<'a>(&'a self) -> impl Iterator<Item = MatrixRow<'p, 'tcx>> + Captures<'a> {
-        self.pats.expand_or_pat().map(|patstack| MatrixRow {
+    fn expand_or_pat<'a>(
+        &'a self,
+        cx: &'a MatchCheckCtxt<'p, 'tcx>,
+    ) -> impl Iterator<Item = MatrixRow<'p, 'tcx>> + Captures<'a> {
+        self.pats.expand_or_pat(cx).map(|patstack| MatrixRow {
             pats: patstack,
             parent_row: self.parent_row,
             is_under_guard: self.is_under_guard,
@@ -723,10 +730,10 @@ struct Matrix<'p, 'tcx> {
 impl<'p, 'tcx> Matrix<'p, 'tcx> {
     /// Pushes a new row to the matrix. If the row starts with an or-pattern, this recursively
     /// expands it. Internal method, prefer [`Matrix::new`].
-    fn expand_and_push(&mut self, row: MatrixRow<'p, 'tcx>) {
+    fn expand_and_push(&mut self, cx: &MatchCheckCtxt<'p, 'tcx>, row: MatrixRow<'p, 'tcx>) {
         if !row.is_empty() && row.head().is_or_pat() {
             // Expand nested or-patterns.
-            for new_row in row.expand_or_pat() {
+            for new_row in row.expand_or_pat(cx) {
                 self.rows.push(new_row);
             }
         } else {
@@ -737,16 +744,16 @@ impl<'p, 'tcx> Matrix<'p, 'tcx> {
     /// Build a new matrix from an iterator of `MatchArm`s.
     fn new(cx: &MatchCheckCtxt<'p, 'tcx>, arms: &[MatchArm<'p, 'tcx>], scrut_ty: Ty<'tcx>) -> Self {
         let wild_pattern = cx.pattern_arena.alloc(DeconstructedPat::wildcard(scrut_ty, DUMMY_SP));
-        let wildcard_row = PatStack::from_pattern(wild_pattern);
+        let wildcard_row = PatStack::from_pattern(cx, wild_pattern);
         let mut matrix = Matrix { rows: Vec::with_capacity(arms.len()), wildcard_row };
         for (row_id, arm) in arms.iter().enumerate() {
             let v = MatrixRow {
-                pats: PatStack::from_pattern(arm.pat),
+                pats: PatStack::from_pattern(cx, arm.pat),
                 parent_row: row_id, // dummy, we won't read it
                 is_under_guard: arm.has_guard,
                 reachable: false,
             };
-            matrix.expand_and_push(v);
+            matrix.expand_and_push(cx, v);
         }
         matrix
     }
@@ -807,7 +814,7 @@ impl<'p, 'tcx> Matrix<'p, 'tcx> {
         for (i, row) in self.rows().enumerate() {
             if ctor.is_covered_by(pcx, row.head().ctor()) {
                 let new_row = row.pop_head_constructor(pcx, ctor, i);
-                matrix.expand_and_push(new_row);
+                matrix.expand_and_push(pcx.cx, new_row);
             }
         }
         matrix
