Auto merge of #118490 - Nadrieril:arena-alloc-matrix, r=<try>

bors · bors · commit c0020fdde2ee · 2023-12-01T08:50:10.000Z
[Experiment] Exhaustiveness: allocate memory better

Exhaustiveness is a recursive algorithm that allocates a bunch of slices at every step. Let's see if I can improve performance by improving allocations.

Already just using `Vec::with_capacity` is showing impressive improvements on my local measurements.

r? `@ghost`
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,11 +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![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())),
         })
     }
 
@@ -614,9 +617,13 @@ impl<'p, 'tcx> PatStack<'p, 'tcx> {
     ) -> PatStack<'p, 'tcx> {
         // We pop the head pattern and push the new fields extracted from the arguments of
         // `self.head()`.
-        let mut new_pats = self.head().specialize(pcx, ctor);
-        new_pats.extend_from_slice(&self.pats[1..]);
-        PatStack { pats: new_pats }
+        let ctor_fields = self.head().specialize(pcx, ctor);
+        PatStack {
+            pats: pcx
+                .cx
+                .dropless_arena
+                .alloc_from_iter(ctor_fields.into_iter().chain(self.pats[1..].iter().copied())),
+        }
     }
 }
 
@@ -667,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,
@@ -711,7 +721,7 @@ impl<'p, 'tcx> fmt::Debug for MatrixRow<'p, 'tcx> {
 /// the matrix will correspond to `scrutinee.0.Some.0` and the second column to `scrutinee.1`.
 #[derive(Clone)]
 struct Matrix<'p, 'tcx> {
-    rows: Vec<MatrixRow<'p, 'tcx>>,
+    rows: SmallVec<[MatrixRow<'p, 'tcx>; 8]>,
     /// Stores an extra fictitious row full of wildcards. Mostly used to keep track of the type of
     /// each column. This must obey the same invariants as the real rows.
     wildcard_row: PatStack<'p, 'tcx>,
@@ -720,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 {
@@ -732,25 +742,18 @@ impl<'p, 'tcx> Matrix<'p, 'tcx> {
     }
 
     /// Build a new matrix from an iterator of `MatchArm`s.
-    fn new<'a>(
-        cx: &MatchCheckCtxt<'p, 'tcx>,
-        iter: impl Iterator<Item = &'a MatchArm<'p, 'tcx>>,
-        scrut_ty: Ty<'tcx>,
-    ) -> Self
-    where
-        'p: 'a,
-    {
+    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 mut matrix = Matrix { rows: vec![], wildcard_row };
-        for (row_id, arm) in iter.enumerate() {
+        let wildcard_row = PatStack::from_pattern(cx, wild_pattern);
+        let mut matrix = Matrix { rows: SmallVec::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
     }
@@ -806,11 +809,12 @@ impl<'p, 'tcx> Matrix<'p, 'tcx> {
         ctor: &Constructor<'tcx>,
     ) -> Matrix<'p, 'tcx> {
         let wildcard_row = self.wildcard_row.pop_head_constructor(pcx, ctor);
-        let mut matrix = Matrix { rows: vec![], wildcard_row };
+        let rows = SmallVec::with_capacity(self.rows.len()); // Better waste capacity than reallocate a lot.
+        let mut matrix = Matrix { rows, wildcard_row };
         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
@@ -1386,7 +1390,7 @@ pub(crate) fn compute_match_usefulness<'p, 'tcx>(
     arms: &[MatchArm<'p, 'tcx>],
     scrut_ty: Ty<'tcx>,
 ) -> UsefulnessReport<'p, 'tcx> {
-    let mut matrix = Matrix::new(cx, arms.iter(), scrut_ty);
+    let mut matrix = Matrix::new(cx, arms, scrut_ty);
     let non_exhaustiveness_witnesses =
         compute_exhaustiveness_and_reachability(cx, &mut matrix, true);
 
