valtree simplification

Author: xacrimon

compiler/rustc_middle/src/arena.rs

@@ -114,6 +114,7 @@ macro_rules! arena_types {
             [decode] specialization_graph: rustc_middle::traits::specialization_graph::Graph,
             [] crate_inherent_impls: rustc_middle::ty::CrateInherentImpls,
             [] hir_owner_nodes: rustc_hir::OwnerNodes<'tcx>,
+            [] pats2: rustc_middle::thir::Pat<'tcx>,
         ]);
     )
 }

compiler/rustc_mir_build/src/builder/matches/match_pair.rs

@@ -1,6 +1,7 @@
 use rustc_middle::mir::*;
 use rustc_middle::thir::{self, *};
 use rustc_middle::ty::{self, Ty, TypeVisitableExt};
+use rustc_span::Span;
 
 use crate::builder::Builder;
 use crate::builder::expr::as_place::{PlaceBase, PlaceBuilder};
@@ -27,12 +28,100 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             .collect()
     }
 
+    // We don't consider an empty subslice as a constant pattern subslice.
+    fn is_constant_pattern_subslice(&self, subslice: &[Box<Pat<'tcx>>]) -> bool {
+        !subslice.is_empty() && subslice.iter().all(|p| self.is_constant_pattern(p))
+    }
+
+    // TODO: expand beyond u8
+    fn is_constant_pattern(&self, pat: &Pat<'tcx>) -> bool {
+        if let PatKind::Constant { value } = pat.kind
+            && let Const::Ty(_, const_) = value
+            && let ty::ConstKind::Value(ty, valtree) = const_.kind()
+            && self.tcx.types.u8 == ty
+            && matches!(valtree, ty::ValTree::Leaf(_))
+        {
+            true
+        } else {
+            false
+        }
+    }
+
+    fn extract_leaf(&self, pat: &Pat<'tcx>) -> ty::ValTree<'tcx> {
+        if let PatKind::Constant { value } = pat.kind
+            && let Const::Ty(_, const_) = value
+            && let ty::ConstKind::Value(_, valtree) = const_.kind()
+            && matches!(valtree, ty::ValTree::Leaf(_))
+        {
+            valtree
+        } else {
+            unreachable!()
+        }
+    }
+
+    // This must only be called after ensuring that the subslice consists of only constant
+    // patterns, or it will panic.
+    fn simplify_const_pattern_slice_into_valtree(
+        &self,
+        subslice: &[Box<Pat<'tcx>>],
+    ) -> ty::ValTree<'tcx> {
+        let leaves = subslice.iter().map(|p| self.extract_leaf(p));
+        let interned = self.tcx.arena.alloc_from_iter(leaves);
+        ty::ValTree::Branch(interned)
+    }
+
+    fn valtree_to_match_pair<'pat>(
+        &mut self,
+        base_pat_ty: Ty<'tcx>,
+        span: Span,
+        subslice_len: u64,
+        valtree: ty::ValTree<'tcx>,
+        place: PlaceBuilder<'tcx>,
+        elem_ty: Ty<'tcx>,
+    ) -> MatchPairTree<'pat, 'tcx> {
+        let tcx = self.tcx;
+        let (const_ty, pat_ty) = if base_pat_ty.is_slice() {
+            (
+                Ty::new_imm_ref(
+                    tcx,
+                    tcx.lifetimes.re_erased,
+                    Ty::new_array(tcx, elem_ty, subslice_len),
+                ),
+                Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, Ty::new_slice(tcx, elem_ty)),
+            )
+        } else {
+            //let arr_ty = Ty::new_array(tcx, elem_ty, subslice_len);
+            let arr_ty = Ty::new_slice(tcx, elem_ty);
+            let pat_ty = Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, arr_ty);
+            (arr_ty, arr_ty)
+        };
+
+        let r#const = ty::Const::new(tcx, ty::ConstKind::Value(const_ty, valtree));
+        let r#const2 = Const::Ty(const_ty, r#const);
+
+        let pattern = tcx.arena.alloc(Pat {
+            ty: pat_ty,
+            span,
+            kind: PatKind::Constant { value: Const::Ty(const_ty, r#const) },
+        });
+
+        let test_case = TestCase::Constant { value: r#const2 };
+
+        MatchPairTree {
+            place: Some(place.to_place(self)),
+            test_case,
+            subpairs: Vec::new(),
+            pattern,
+        }
+    }
+
     /// Builds [`MatchPairTree`] subtrees for the prefix/middle/suffix parts of an
     /// array pattern or slice pattern, and adds those trees to `match_pairs`.
     ///
     /// Used internally by [`MatchPairTree::for_pattern`].
     fn prefix_slice_suffix<'pat>(
         &mut self,
+        base_pat: &'pat Pat<'tcx>,
         match_pairs: &mut Vec<MatchPairTree<'pat, 'tcx>>,
         place: &PlaceBuilder<'tcx>,
         prefix: &'pat [Box<Pat<'tcx>>],
@@ -54,11 +143,30 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             ((prefix.len() + suffix.len()).try_into().unwrap(), false)
         };
 
-        match_pairs.extend(prefix.iter().enumerate().map(|(idx, subpattern)| {
-            let elem =
-                ProjectionElem::ConstantIndex { offset: idx as u64, min_length, from_end: false };
-            MatchPairTree::for_pattern(place.clone_project(elem), subpattern, self)
-        }));
+        // Here we try to special-case constant pattern subslices into valtrees to generate
+        // nicer MIR.
+        // Try to perform simplification of a constant pattern slice `prefix` sequence into
+        // a valtree.
+        if self.is_constant_pattern_subslice(prefix) {
+            let elem_ty = prefix[0].ty;
+            let prefix_valtree = self.simplify_const_pattern_slice_into_valtree(prefix);
+            // FIXME(jieyouxu): triple check these place calculations!
+            let match_pair = self.valtree_to_match_pair(
+                base_pat.ty,
+                base_pat.span,
+                prefix.len() as u64,
+                prefix_valtree,
+                place.base().into(),
+                elem_ty,
+            );
+            match_pairs.push(match_pair);
+        } else {
+            match_pairs.extend(prefix.iter().enumerate().map(|(idx, subpattern)| {
+                let elem =
+                    ProjectionElem::ConstantIndex { offset: idx as u64, min_length, from_end: false };
+                MatchPairTree::for_pattern(place.clone_project(elem), subpattern, self)
+            }));
+        }
 
         if let Some(subslice_pat) = opt_slice {
             let suffix_len = suffix.len() as u64;
@@ -192,11 +300,11 @@ impl<'pat, 'tcx> MatchPairTree<'pat, 'tcx> {
             }
 
             PatKind::Array { ref prefix, ref slice, ref suffix } => {
-                cx.prefix_slice_suffix(&mut subpairs, &place_builder, prefix, slice, suffix);
+                cx.prefix_slice_suffix(pattern, &mut subpairs, &place_builder, prefix, slice, suffix);
                 default_irrefutable()
             }
             PatKind::Slice { ref prefix, ref slice, ref suffix } => {
-                cx.prefix_slice_suffix(&mut subpairs, &place_builder, prefix, slice, suffix);
+                cx.prefix_slice_suffix(pattern, &mut subpairs, &place_builder, prefix, slice, suffix);
 
                 if prefix.is_empty() && slice.is_some() && suffix.is_empty() {
                     default_irrefutable()