Refactor lower_stmts

Author: camsteffen

compiler/rustc_ast_lowering/src/block.rs

@@ -2,7 +2,7 @@ use crate::{ImplTraitContext, ImplTraitPosition, LoweringContext};
 use rustc_ast::{Block, BlockCheckMode, Local, Stmt, StmtKind};
 use rustc_hir as hir;
 
-use smallvec::{smallvec, SmallVec};
+use smallvec::SmallVec;
 
 impl<'a, 'hir> LoweringContext<'a, 'hir> {
     pub(super) fn lower_block(
@@ -18,66 +18,67 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         b: &Block,
         targeted_by_break: bool,
     ) -> hir::Block<'hir> {
-        let (stmts, expr) = match &*b.stmts {
-            [stmts @ .., Stmt { kind: StmtKind::Expr(e), .. }] => (stmts, Some(&*e)),
-            stmts => (stmts, None),
-        };
-        let stmts = self.arena.alloc_from_iter(stmts.iter().flat_map(|stmt| self.lower_stmt(stmt)));
-        let expr = expr.map(|e| self.lower_expr(e));
+        let (stmts, expr) = self.lower_stmts(&b.stmts);
         let rules = self.lower_block_check_mode(&b.rules);
         let hir_id = self.lower_node_id(b.id);
-
         hir::Block { hir_id, stmts, expr, rules, span: self.lower_span(b.span), targeted_by_break }
     }
 
-    fn lower_stmt(&mut self, s: &Stmt) -> SmallVec<[hir::Stmt<'hir>; 1]> {
-        let (hir_id, kind) = match s.kind {
-            StmtKind::Local(ref l) => {
-                let l = self.lower_local(l);
-                let hir_id = self.lower_node_id(s.id);
-                self.alias_attrs(hir_id, l.hir_id);
-                return smallvec![hir::Stmt {
-                    hir_id,
-                    kind: hir::StmtKind::Local(self.arena.alloc(l)),
-                    span: self.lower_span(s.span),
-                }];
-            }
-            StmtKind::Item(ref it) => {
-                // Can only use the ID once.
-                let mut id = Some(s.id);
-                return self
-                    .lower_item_id(it)
-                    .into_iter()
-                    .map(|item_id| {
-                        let hir_id = id
-                            .take()
-                            .map(|id| self.lower_node_id(id))
-                            .unwrap_or_else(|| self.next_id());
-
-                        hir::Stmt {
-                            hir_id,
-                            kind: hir::StmtKind::Item(item_id),
-                            span: self.lower_span(s.span),
-                        }
-                    })
-                    .collect();
-            }
-            StmtKind::Expr(ref e) => {
-                let e = self.lower_expr(e);
-                let hir_id = self.lower_node_id(s.id);
-                self.alias_attrs(hir_id, e.hir_id);
-                (hir_id, hir::StmtKind::Expr(e))
-            }
-            StmtKind::Semi(ref e) => {
-                let e = self.lower_expr(e);
-                let hir_id = self.lower_node_id(s.id);
-                self.alias_attrs(hir_id, e.hir_id);
-                (hir_id, hir::StmtKind::Semi(e))
+    fn lower_stmts(
+        &mut self,
+        mut ast_stmts: &[Stmt],
+    ) -> (&'hir [hir::Stmt<'hir>], Option<&'hir hir::Expr<'hir>>) {
+        let mut stmts = SmallVec::<[hir::Stmt<'hir>; 8]>::new();
+        let mut expr = None;
+        while let [s, tail @ ..] = ast_stmts {
+            match s.kind {
+                StmtKind::Local(ref l) => {
+                    let l = self.lower_local(l);
+                    let hir_id = self.lower_node_id(s.id);
+                    self.alias_attrs(hir_id, l.hir_id);
+                    let kind = hir::StmtKind::Local(self.arena.alloc(l));
+                    let span = self.lower_span(s.span);
+                    stmts.push(hir::Stmt { hir_id, kind, span });
+                }
+                StmtKind::Item(ref it) => {
+                    stmts.extend(self.lower_item_id(it).into_iter().enumerate().map(
+                        |(i, item_id)| {
+                            let hir_id = match i {
+                                0 => self.lower_node_id(s.id),
+                                _ => self.next_id(),
+                            };
+                            let kind = hir::StmtKind::Item(item_id);
+                            let span = self.lower_span(s.span);
+                            hir::Stmt { hir_id, kind, span }
+                        },
+                    ));
+                }
+                StmtKind::Expr(ref e) => {
+                    let e = self.lower_expr(e);
+                    if tail.is_empty() {
+                        expr = Some(e);
+                    } else {
+                        let hir_id = self.lower_node_id(s.id);
+                        self.alias_attrs(hir_id, e.hir_id);
+                        let kind = hir::StmtKind::Expr(e);
+                        let span = self.lower_span(s.span);
+                        stmts.push(hir::Stmt { hir_id, kind, span });
+                    }
+                }
+                StmtKind::Semi(ref e) => {
+                    let e = self.lower_expr(e);
+                    let hir_id = self.lower_node_id(s.id);
+                    self.alias_attrs(hir_id, e.hir_id);
+                    let kind = hir::StmtKind::Semi(e);
+                    let span = self.lower_span(s.span);
+                    stmts.push(hir::Stmt { hir_id, kind, span });
+                }
+                StmtKind::Empty => {}
+                StmtKind::MacCall(..) => panic!("shouldn't exist here"),
             }
-            StmtKind::Empty => return smallvec![],
-            StmtKind::MacCall(..) => panic!("shouldn't exist here"),
-        };
-        smallvec![hir::Stmt { hir_id, kind, span: self.lower_span(s.span) }]
+            ast_stmts = &ast_stmts[1..];
+        }
+        (self.arena.alloc_from_iter(stmts), expr)
     }
 
     fn lower_local(&mut self, l: &Local) -> hir::Local<'hir> {