Refactor rustc_peek

We now use `DataflowResultsCursor` to get the dataflow state before
calls to `rustc_peek`. The original version used a custom implementation
that only looked at assignment statements. This also extends
`rustc_peek` to take arguments by-value as well as by-reference, and
allows *all* dataflow analyses, not just those dependent on `MoveData`,
to be inspected.

Author: ecstatic-morse

src/librustc_mir/transform/rustc_peek.rs

@@ -3,16 +3,17 @@ use syntax::ast;
 use syntax::symbol::sym;
 use syntax_pos::Span;
 
-use rustc::ty::{self, TyCtxt};
+use rustc::ty::{self, TyCtxt, Ty};
 use rustc::hir::def_id::DefId;
-use rustc::mir::{self, Body, Location};
+use rustc::mir::{self, Body, Location, Local};
 use rustc_index::bit_set::BitSet;
 use crate::transform::{MirPass, MirSource};
 
 use crate::dataflow::{do_dataflow, DebugFormatted};
 use crate::dataflow::MoveDataParamEnv;
 use crate::dataflow::BitDenotation;
 use crate::dataflow::DataflowResults;
+use crate::dataflow::DataflowResultsCursor;
 use crate::dataflow::{
     DefinitelyInitializedPlaces, MaybeInitializedPlaces, MaybeUninitializedPlaces
 };
@@ -88,151 +89,172 @@ pub fn sanity_check_via_rustc_peek<'tcx, O>(
     def_id: DefId,
     _attributes: &[ast::Attribute],
     results: &DataflowResults<'tcx, O>,
-) where
-    O: BitDenotation<'tcx, Idx = MovePathIndex> + HasMoveData<'tcx>,
-{
+) where O: RustcPeekAt<'tcx> {
     debug!("sanity_check_via_rustc_peek def_id: {:?}", def_id);
-    // FIXME: this is not DRY. Figure out way to abstract this and
-    // `dataflow::build_sets`. (But note it is doing non-standard
-    // stuff, so such generalization may not be realistic.)
 
-    for bb in body.basic_blocks().indices() {
-        each_block(tcx, body, results, bb);
-    }
-}
+    let mut cursor = DataflowResultsCursor::new(results, body);
 
-fn each_block<'tcx, O>(
-    tcx: TyCtxt<'tcx>,
-    body: &Body<'tcx>,
-    results: &DataflowResults<'tcx, O>,
-    bb: mir::BasicBlock,
-) where
-    O: BitDenotation<'tcx, Idx = MovePathIndex> + HasMoveData<'tcx>,
-{
-    let move_data = results.0.operator.move_data();
-    let mir::BasicBlockData { ref statements, ref terminator, is_cleanup: _ } = body[bb];
-
-    let (args, span) = match is_rustc_peek(tcx, terminator) {
-        Some(args_and_span) => args_and_span,
-        None => return,
-    };
-    assert!(args.len() == 1);
-    let peek_arg_place = match args[0] {
-        mir::Operand::Copy(ref place @ mir::Place {
-            base: mir::PlaceBase::Local(_),
-            projection: box [],
-        }) |
-        mir::Operand::Move(ref place @ mir::Place {
-            base: mir::PlaceBase::Local(_),
-            projection: box [],
-        }) => Some(place),
-        _ => None,
-    };
-
-    let peek_arg_place = match peek_arg_place {
-        Some(arg) => arg,
-        None => {
-            tcx.sess.diagnostic().span_err(
-                span, "dataflow::sanity_check cannot feed a non-temp to rustc_peek.");
-            return;
-        }
-    };
-
-    let mut on_entry = results.0.sets.entry_set_for(bb.index()).to_owned();
-    let mut trans = results.0.sets.trans_for(bb.index()).clone();
-
-    // Emulate effect of all statements in the block up to (but not
-    // including) the borrow within `peek_arg_place`. Do *not* include
-    // call to `peek_arg_place` itself (since we are peeking the state
-    // of the argument at time immediate preceding Call to
-    // `rustc_peek`).
-
-    for (j, stmt) in statements.iter().enumerate() {
-        debug!("rustc_peek: ({:?},{}) {:?}", bb, j, stmt);
-        let (place, rvalue) = match stmt.kind {
-            mir::StatementKind::Assign(box(ref place, ref rvalue)) => {
-                (place, rvalue)
+    let peek_calls = body
+            .basic_blocks()
+            .iter_enumerated()
+            .filter_map(|(bb, block_data)| {
+                PeekCall::from_terminator(tcx, block_data.terminator())
+                    .map(|call| (bb, block_data, call))
+            });
+
+    for (bb, block_data, call) in peek_calls {
+        // Look for a sequence like the following to indicate that we should be peeking at `_1`:
+        //    _2 = &_1;
+        //    rustc_peek(_2);
+        //
+        //    /* or */
+        //
+        //    _2 = _1;
+        //    rustc_peek(_2);
+        let (statement_index, peek_rval) = block_data
+            .statements
+            .iter()
+            .enumerate()
+            .filter_map(|(i, stmt)| value_assigned_to_local(stmt, call.arg).map(|rval| (i, rval)))
+            .next()
+            .expect("call to rustc_peek should be preceded by \
+                    assignment to temporary holding its argument");
+
+        match (call.kind, peek_rval) {
+            | (PeekCallKind::ByRef, mir::Rvalue::Ref(_, _, place))
+            | (PeekCallKind::ByVal, mir::Rvalue::Use(mir::Operand::Move(place)))
+            | (PeekCallKind::ByVal, mir::Rvalue::Use(mir::Operand::Copy(place)))
+            => {
+                let loc = Location { block: bb, statement_index };
+                cursor.seek(loc);
+                let state = cursor.get();
+                results.operator().peek_at(tcx, place, state, call);
             }
-            mir::StatementKind::FakeRead(..) |
-            mir::StatementKind::StorageLive(_) |
-            mir::StatementKind::StorageDead(_) |
-            mir::StatementKind::InlineAsm { .. } |
-            mir::StatementKind::Retag { .. } |
-            mir::StatementKind::AscribeUserType(..) |
-            mir::StatementKind::Nop => continue,
-            mir::StatementKind::SetDiscriminant{ .. } =>
-                span_bug!(stmt.source_info.span,
-                          "sanity_check should run before Deaggregator inserts SetDiscriminant"),
-        };
-
-        if place == peek_arg_place {
-            if let mir::Rvalue::Ref(_, mir::BorrowKind::Shared, ref peeking_at_place) = *rvalue {
-                // Okay, our search is over.
-                match move_data.rev_lookup.find(peeking_at_place.as_ref()) {
-                    LookupResult::Exact(peek_mpi) => {
-                        let bit_state = on_entry.contains(peek_mpi);
-                        debug!("rustc_peek({:?} = &{:?}) bit_state: {}",
-                               place, peeking_at_place, bit_state);
-                        if !bit_state {
-                            tcx.sess.span_err(span, "rustc_peek: bit not set");
-                        }
-                    }
-                    LookupResult::Parent(..) => {
-                        tcx.sess.span_err(span, "rustc_peek: argument untracked");
-                    }
-                }
-                return;
-            } else {
-                // Our search should have been over, but the input
-                // does not match expectations of `rustc_peek` for
-                // this sanity_check.
+
+            _ => {
                 let msg = "rustc_peek: argument expression \
-                           must be immediate borrow of form `&expr`";
-                tcx.sess.span_err(span, msg);
+                           must be either `place` or `&place`";
+                tcx.sess.span_err(call.span, msg);
             }
         }
+    }
+}
 
-        let lhs_mpi = move_data.rev_lookup.find(place.as_ref());
-
-        debug!("rustc_peek: computing effect on place: {:?} ({:?}) in stmt: {:?}",
-               place, lhs_mpi, stmt);
-        // reset GEN and KILL sets before emulating their effect.
-        trans.clear();
-        results.0.operator.before_statement_effect(
-            &mut trans,
-            Location { block: bb, statement_index: j });
-        results.0.operator.statement_effect(
-            &mut trans,
-            Location { block: bb, statement_index: j });
-        trans.apply(&mut on_entry);
+/// If `stmt` is an assignment where the LHS is the given local (with no projections), returns the
+/// RHS of the assignment.
+fn value_assigned_to_local<'a, 'tcx>(
+    stmt: &'a mir::Statement<'tcx>,
+    local: Local,
+) -> Option<&'a mir::Rvalue<'tcx>> {
+    if let mir::StatementKind::Assign(box (place, rvalue)) = &stmt.kind {
+        if let mir::Place { base: mir::PlaceBase::Local(l), projection: box [] } = place {
+            if local == *l {
+                return Some(&*rvalue);
+            }
+        }
     }
 
-    results.0.operator.before_terminator_effect(
-        &mut trans,
-        Location { block: bb, statement_index: statements.len() });
+    None
+}
 
-    tcx.sess.span_err(span, &format!("rustc_peek: MIR did not match \
-                                      anticipated pattern; note that \
-                                      rustc_peek expects input of \
-                                      form `&expr`"));
+#[derive(Clone, Copy, Debug)]
+enum PeekCallKind {
+    ByVal,
+    ByRef,
 }
 
-fn is_rustc_peek<'a, 'tcx>(
-    tcx: TyCtxt<'tcx>,
-    terminator: &'a Option<mir::Terminator<'tcx>>,
-) -> Option<(&'a [mir::Operand<'tcx>], Span)> {
-    if let Some(mir::Terminator { ref kind, source_info, .. }) = *terminator {
-        if let mir::TerminatorKind::Call { func: ref oper, ref args, .. } = *kind {
-            if let mir::Operand::Constant(ref func) = *oper {
-                if let ty::FnDef(def_id, _) = func.literal.ty.kind {
-                    let abi = tcx.fn_sig(def_id).abi();
-                    let name = tcx.item_name(def_id);
-                    if abi == Abi::RustIntrinsic && name == sym::rustc_peek {
-                        return Some((args, source_info.span));
+impl PeekCallKind {
+    fn from_arg_ty(arg: Ty<'_>) -> Self {
+        match arg.kind {
+            ty::Ref(_, _, _) => PeekCallKind::ByRef,
+            _ => PeekCallKind::ByVal,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct PeekCall {
+    arg: Local,
+    kind: PeekCallKind,
+    span: Span,
+}
+
+impl PeekCall {
+    fn from_terminator<'tcx>(
+        tcx: TyCtxt<'tcx>,
+        terminator: &mir::Terminator<'tcx>,
+    ) -> Option<Self> {
+        use mir::{Operand, Place, PlaceBase};
+
+        let span = terminator.source_info.span;
+        if let mir::TerminatorKind::Call { func: Operand::Constant(func), args, .. } =
+            &terminator.kind
+        {
+            if let ty::FnDef(def_id, substs) = func.literal.ty.kind {
+                let sig = tcx.fn_sig(def_id);
+                let name = tcx.item_name(def_id);
+                if sig.abi() != Abi::RustIntrinsic || name != sym::rustc_peek {
+                    return None;
+                }
+
+                assert_eq!(args.len(), 1);
+                let kind = PeekCallKind::from_arg_ty(substs.type_at(0));
+                let arg = match args[0] {
+                    | Operand::Copy(Place { base: PlaceBase::Local(local), projection: box [] })
+                    | Operand::Move(Place { base: PlaceBase::Local(local), projection: box [] })
+                    => local,
+
+                    _ => {
+                        tcx.sess.diagnostic().span_err(
+                            span, "dataflow::sanity_check cannot feed a non-temp to rustc_peek.");
+                        return None;
                     }
+                };
+
+                return Some(PeekCall {
+                    arg,
+                    kind,
+                    span,
+                });
+            }
+        }
+
+        None
+    }
+}
+
+pub trait RustcPeekAt<'tcx>: BitDenotation<'tcx> {
+    fn peek_at(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        place: &mir::Place<'tcx>,
+        flow_state: &BitSet<Self::Idx>,
+        call: PeekCall,
+    );
+}
+
+impl<'tcx, O> RustcPeekAt<'tcx> for O
+    where O: BitDenotation<'tcx, Idx = MovePathIndex> + HasMoveData<'tcx>,
+{
+    fn peek_at(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        place: &mir::Place<'tcx>,
+        flow_state: &BitSet<Self::Idx>,
+        call: PeekCall,
+    ) {
+        match self.move_data().rev_lookup.find(place.as_ref()) {
+            LookupResult::Exact(peek_mpi) => {
+                let bit_state = flow_state.contains(peek_mpi);
+                debug!("rustc_peek({:?} = &{:?}) bit_state: {}",
+                       call.arg, place, bit_state);
+                if !bit_state {
+                    tcx.sess.span_err(call.span, "rustc_peek: bit not set");
                 }
             }
+            LookupResult::Parent(..) => {
+                tcx.sess.span_err(call.span, "rustc_peek: argument untracked");
+            }
         }
     }
-    return None;
 }