optionally only emit basic validation for functions containing unsafe block / unsafe function

7 files changed, 155 insertions, 58 deletions

diff --git a/src/librustc/hir/mod.rs b/src/librustc/hir/mod.rs
index cc0d49c1a36..85d9745246f 100644
--- a/src/librustc/hir/mod.rs
+++ b/src/librustc/hir/mod.rs
@@ -49,7 +49,7 @@ use rustc_data_structures::indexed_vec;
 use std::collections::BTreeMap;
 use std::fmt;
 
-/// HIR doesn't commit to a concrete storage type and have its own alias for a vector.
+/// HIR doesn't commit to a concrete storage type and has its own alias for a vector.
 /// It can be `Vec`, `P<[T]>` or potentially `Box<[T]>`, or some other container with similar
 /// behavior. Unlike AST, HIR is mostly a static structure, so we can use an owned slice instead
 /// of `Vec` to avoid keeping extra capacity.
@@ -76,14 +76,14 @@ pub mod pat_util;
 pub mod print;
 pub mod svh;
 
-/// A HirId uniquely identifies a node in the HIR of then current crate. It is
+/// A HirId uniquely identifies a node in the HIR of the current crate. It is
 /// composed of the `owner`, which is the DefIndex of the directly enclosing
 /// hir::Item, hir::TraitItem, or hir::ImplItem (i.e. the closest "item-like"),
 /// and the `local_id` which is unique within the given owner.
 ///
 /// This two-level structure makes for more stable values: One can move an item
 /// around within the source code, or add or remove stuff before it, without
-/// the local_id part of the HirId changing, which is a very useful property
+/// the local_id part of the HirId changing, which is a very useful property in
 /// incremental compilation where we have to persist things through changes to
 /// the code base.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug,
diff --git a/src/librustc/session/config.rs b/src/librustc/session/config.rs
index c5ddcb597cb..c8b9412c566 100644
--- a/src/librustc/session/config.rs
+++ b/src/librustc/session/config.rs
@@ -1025,8 +1025,9 @@ options! {DebuggingOptions, DebuggingSetter, basic_debugging_options,
           "the directory the MIR is dumped into"),
     dump_mir_exclude_pass_number: bool = (false, parse_bool, [UNTRACKED],
           "if set, exclude the pass number when dumping MIR (used in tests)"),
-    mir_emit_validate: bool = (false, parse_bool, [TRACKED],
-          "emit Validate MIR statements, interpreted e.g. by miri"),
+    mir_emit_validate: usize = (0, parse_uint, [TRACKED],
+          "emit Validate MIR statements, interpreted e.g. by miri (0: do not emit; 1: if function \
+           contains unsafe block, only validate arguments; 2: always emit full validation)"),
     perf_stats: bool = (false, parse_bool, [UNTRACKED],
           "print some performance-related statistics"),
     hir_stats: bool = (false, parse_bool, [UNTRACKED],
diff --git a/src/librustc_mir/transform/add_validation.rs b/src/librustc_mir/transform/add_validation.rs
index ee472c616f6..1329378fbef 100644
--- a/src/librustc_mir/transform/add_validation.rs
+++ b/src/librustc_mir/transform/add_validation.rs
@@ -14,6 +14,8 @@
 //! of MIR building, and only after this pass we think of the program has having the
 //! normal MIR semantics.
 
+use syntax_pos::Span;
+use syntax::ast::NodeId;
 use rustc::ty::{self, TyCtxt, RegionKind};
 use rustc::hir;
 use rustc::mir::*;
@@ -80,15 +82,78 @@ fn lval_context<'a, 'tcx, D>(
     }
 }
 
+/// Check if this function contains an unsafe block or is an unsafe function.
+fn fn_contains_unsafe<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource) -> bool {
+    use rustc::hir::intravisit::{self, Visitor};
+
+    let fn_node_id = match src {
+        MirSource::Fn(node_id) => node_id,
+        _ => return false, // only functions can have unsafe
+    };
+    let fn_item = tcx.hir.expect_item(fn_node_id);
+
+    struct FindUnsafe<'b, 'tcx> where 'tcx : 'b {
+        map: &'b hir::map::Map<'tcx>,
+        found_unsafe: bool,
+    }
+    let mut finder = FindUnsafe { map: &tcx.hir, found_unsafe: false };
+    finder.visit_item(fn_item);
+
+    impl<'b, 'tcx> Visitor<'tcx> for FindUnsafe<'b, 'tcx> {
+        fn nested_visit_map<'this>(&'this mut self) -> intravisit::NestedVisitorMap<'this, 'tcx> {
+            intravisit::NestedVisitorMap::OnlyBodies(self.map)
+        }
+
+        fn visit_fn(&mut self, fk: intravisit::FnKind<'tcx>, fd: &'tcx hir::FnDecl,
+                    b: hir::BodyId, s: Span, id: NodeId)
+        {
+            assert!(!self.found_unsafe, "We should never see more than one fn");
+            let is_unsafe = match fk {
+                intravisit::FnKind::ItemFn(_, _, unsafety, ..) => unsafety == hir::Unsafety::Unsafe,
+                intravisit::FnKind::Method(_, sig, ..) => sig.unsafety == hir::Unsafety::Unsafe,
+                intravisit::FnKind::Closure(_) => false,
+            };
+            if is_unsafe {
+                // This is unsafe, and we are done.
+                self.found_unsafe = true;
+            } else {
+                // Go on searching.
+                intravisit::walk_fn(self, fk, fd, b, s, id)
+            }
+        }
+
+        fn visit_block(&mut self, b: &'tcx hir::Block) {
+            use rustc::hir::BlockCheckMode::*;
+
+            if self.found_unsafe { return; } // short-circuit
+
+            match b.rules {
+                UnsafeBlock(_) | PushUnsafeBlock(_) => {
+                    // We found an unsafe block.
+                    self.found_unsafe = true;
+                }
+                DefaultBlock | PopUnsafeBlock(_) => {
+                    // No unsafe block here, go on searching.
+                    intravisit::walk_block(self, b);
+                }
+            };
+        }
+    }
+
+    finder.found_unsafe
+}
+
 impl MirPass for AddValidation {
     fn run_pass<'a, 'tcx>(&self,
                           tcx: TyCtxt<'a, 'tcx, 'tcx>,
-                          _: MirSource,
-                          mir: &mut Mir<'tcx>) {
-        if !tcx.sess.opts.debugging_opts.mir_emit_validate {
+                          src: MirSource,
+                          mir: &mut Mir<'tcx>)
+    {
+        let emit_validate = tcx.sess.opts.debugging_opts.mir_emit_validate;
+        if emit_validate == 0 {
             return;
         }
-
+        let restricted_validation = emit_validate == 1 && fn_contains_unsafe(tcx, src);
         let local_decls = mir.local_decls.clone(); // FIXME: Find a way to get rid of this clone.
 
         // Convert an lvalue to a validation operand.
@@ -98,22 +163,40 @@ impl MirPass for AddValidation {
             ValidationOperand { lval, ty, re, mutbl }
         };
 
+        // Emit an Acquire at the beginning of the given block.  If we are in restricted emission mode
+        // (mir_emit_validate=1), also emit a Release immediately after the Acquire.
+        let emit_acquire = |block: &mut BasicBlockData<'tcx>, source_info, operands: Vec<_>| {
+            if operands.len() == 0 {
+                return; // Nothing to do
+            }
+            // Emit the release first, to avoid cloning if we do not emit it
+            if restricted_validation {
+                let release_stmt = Statement {
+                    source_info,
+                    kind: StatementKind::Validate(ValidationOp::Release, operands.clone()),
+                };
+                block.statements.insert(0, release_stmt);
+            }
+            // Now, the acquire
+            let acquire_stmt = Statement {
+                source_info,
+                kind: StatementKind::Validate(ValidationOp::Acquire, operands),
+            };
+            block.statements.insert(0, acquire_stmt);
+        };
+
         // PART 1
         // Add an AcquireValid at the beginning of the start block.
-        if mir.arg_count > 0 {
-            let acquire_stmt = Statement {
-                source_info: SourceInfo {
-                    scope: ARGUMENT_VISIBILITY_SCOPE,
-                    span: mir.span, // FIXME: Consider using just the span covering the function
-                                    // argument declaration.
-                },
-                kind: StatementKind::Validate(ValidationOp::Acquire,
-                    // Skip return value, go over all the arguments
-                    mir.local_decls.iter_enumerated().skip(1).take(mir.arg_count)
-                    .map(|(local, _)| lval_to_operand(Lvalue::Local(local))).collect()
-                )
+        {
+            let source_info = SourceInfo {
+                scope: ARGUMENT_VISIBILITY_SCOPE,
+                span: mir.span, // FIXME: Consider using just the span covering the function
+                                // argument declaration.
             };
-            mir.basic_blocks_mut()[START_BLOCK].statements.insert(0, acquire_stmt);
+            // Gather all arguments, skip return value.
+            let operands = mir.local_decls.iter_enumerated().skip(1).take(mir.arg_count)
+                    .map(|(local, _)| lval_to_operand(Lvalue::Local(local))).collect();
+            emit_acquire(&mut mir.basic_blocks_mut()[START_BLOCK], source_info, operands);
         }
 
         // PART 2
@@ -125,18 +208,20 @@ impl MirPass for AddValidation {
                 Some(Terminator { kind: TerminatorKind::Call { ref args, ref destination, .. },
                                   source_info }) => {
                     // Before the call: Release all arguments
-                    let release_stmt = Statement {
-                        source_info,
-                        kind: StatementKind::Validate(ValidationOp::Release,
-                            args.iter().filter_map(|op| {
-                                match op {
-                                    &Operand::Consume(ref lval) =>
-                                        Some(lval_to_operand(lval.clone())),
-                                    &Operand::Constant(..) => { None },
-                                }
-                            }).collect())
-                    };
-                    block_data.statements.push(release_stmt);
+                    if !restricted_validation {
+                        let release_stmt = Statement {
+                            source_info,
+                            kind: StatementKind::Validate(ValidationOp::Release,
+                                args.iter().filter_map(|op| {
+                                    match op {
+                                        &Operand::Consume(ref lval) =>
+                                            Some(lval_to_operand(lval.clone())),
+                                        &Operand::Constant(..) => { None },
+                                    }
+                                }).collect())
+                        };
+                        block_data.statements.push(release_stmt);
+                    }
                     // Remember the return destination for later
                     if let &Some(ref destination) = destination {
                         returns.push((source_info, destination.0.clone(), destination.1));
@@ -147,12 +232,14 @@ impl MirPass for AddValidation {
                 Some(Terminator { kind: TerminatorKind::DropAndReplace { location: ref lval, .. },
                                   source_info }) => {
                     // Before the call: Release all arguments
-                    let release_stmt = Statement {
-                        source_info,
-                        kind: StatementKind::Validate(ValidationOp::Release,
-                                vec![lval_to_operand(lval.clone())]),
-                    };
-                    block_data.statements.push(release_stmt);
+                    if !restricted_validation {
+                        let release_stmt = Statement {
+                            source_info,
+                            kind: StatementKind::Validate(ValidationOp::Release,
+                                    vec![lval_to_operand(lval.clone())]),
+                        };
+                        block_data.statements.push(release_stmt);
+                    }
                     // drop doesn't return anything, so we need no acquire.
                 }
                 _ => {
@@ -162,18 +249,21 @@ impl MirPass for AddValidation {
         }
         // Now we go over the returns we collected to acquire the return values.
         for (source_info, dest_lval, dest_block) in returns {
-            let acquire_stmt = Statement {
+            emit_acquire(
+                &mut mir.basic_blocks_mut()[dest_block],
                 source_info,
-                kind: StatementKind::Validate(ValidationOp::Acquire,
-                        vec![lval_to_operand(dest_lval)]),
-            };
-            mir.basic_blocks_mut()[dest_block].statements.insert(0, acquire_stmt);
+                vec![lval_to_operand(dest_lval)]
+            );
+        }
+
+        if restricted_validation {
+            // No part 3 for us.
+            return;
         }
 
         // PART 3
         // Add ReleaseValid/AcquireValid around Ref and Cast.  Again an iterator does not seem very
-        // suited
-        // as we need to add new statements before and after each Ref.
+        // suited as we need to add new statements before and after each Ref.
         for block_data in mir.basic_blocks_mut() {
             // We want to insert statements around Ref commands as we iterate.  To this end, we
             // iterate backwards using indices.
diff --git a/src/librustc_mir/transform/erase_regions.rs b/src/librustc_mir/transform/erase_regions.rs
index f01d71fde26..baf0522896c 100644
--- a/src/librustc_mir/transform/erase_regions.rs
+++ b/src/librustc_mir/transform/erase_regions.rs
@@ -77,7 +77,9 @@ impl<'a, 'tcx> MutVisitor<'tcx> for EraseRegionsVisitor<'a, 'tcx> {
                        block: BasicBlock,
                        statement: &mut Statement<'tcx>,
                        location: Location) {
-        if !self.tcx.sess.opts.debugging_opts.mir_emit_validate {
+        // Do NOT delete EndRegion if validation statements are emitted.
+        // Validation needs EndRegion.
+        if self.tcx.sess.opts.debugging_opts.mir_emit_validate == 0 {
             if let StatementKind::EndRegion(_) = statement.kind {
                 statement.kind = StatementKind::Nop;
             }
diff --git a/src/test/mir-opt/validate_1.rs b/src/test/mir-opt/validate_1.rs
index 868d23b03c2..558426fcde1 100644
--- a/src/test/mir-opt/validate_1.rs
+++ b/src/test/mir-opt/validate_1.rs
@@ -9,7 +9,7 @@
 // except according to those terms.
 
 // ignore-tidy-linelength
-// compile-flags: -Z verbose -Z mir-emit-validate
+// compile-flags: -Z verbose -Z mir-emit-validate=1
 
 fn foo(_x: &mut i32) {}
 
diff --git a/src/test/mir-opt/validate_2.rs b/src/test/mir-opt/validate_2.rs
index a219c5fc78e..21723739ca1 100644
--- a/src/test/mir-opt/validate_2.rs
+++ b/src/test/mir-opt/validate_2.rs
@@ -9,7 +9,7 @@
 // except according to those terms.
 
 // ignore-tidy-linelength
-// compile-flags: -Z verbose -Z mir-emit-validate
+// compile-flags: -Z verbose -Z mir-emit-validate=1
 
 fn main() {
     let _x : Box<[i32]> = Box::new([1, 2, 3]);
diff --git a/src/test/mir-opt/validate_3.rs b/src/test/mir-opt/validate_3.rs
index 78957115f50..88ae114c579 100644
--- a/src/test/mir-opt/validate_3.rs
+++ b/src/test/mir-opt/validate_3.rs
@@ -9,7 +9,7 @@
 // except according to those terms.
 
 // ignore-tidy-linelength
-// compile-flags: -Z verbose -Z mir-emit-validate
+// compile-flags: -Z verbose -Z mir-emit-validate=1
 
 struct Test {
     x: i32
@@ -18,6 +18,10 @@ struct Test {
 fn foo(_x: &i32) {}
 
 fn main() {
+    // These internal unsafe functions should have no effect on the code generation.
+    unsafe fn _unused1() {}
+    fn _unused2(x: *const i32) -> i32 { unsafe { *x }}
+
     let t = Test { x: 0 };
     let t = &t;
     foo(&t.x);
@@ -28,18 +32,18 @@ fn main() {
 // fn main() -> () {
 //     let mut _5: &ReErased i32;
 //     bb0: {
-//         Validate(Suspend(ReScope(Misc(NodeId(31)))), [((*_2).0: i32)@i32/ReScope(Remainder(BlockRemainder { block: NodeId(18), first_statement_index: 1 })) (imm)]);
+//         Validate(Suspend(ReScope(Misc(NodeId(46)))), [((*_2).0: i32)@i32/ReScope(Remainder(BlockRemainder { block: NodeId(18), first_statement_index: 3 })) (imm)]);
 //         _5 = &ReErased ((*_2).0: i32);
-//         Validate(Acquire, [(*_5)@i32/ReScope(Misc(NodeId(31))) (imm)]);
-//         Validate(Suspend(ReScope(Misc(NodeId(31)))), [(*_5)@i32/ReScope(Misc(NodeId(31))) (imm)]);
+//         Validate(Acquire, [(*_5)@i32/ReScope(Misc(NodeId(46))) (imm)]);
+//         Validate(Suspend(ReScope(Misc(NodeId(46)))), [(*_5)@i32/ReScope(Misc(NodeId(46))) (imm)]);
 //         _4 = &ReErased (*_5);
-//         Validate(Acquire, [(*_4)@i32/ReScope(Misc(NodeId(31))) (imm)]);
-//         Validate(Release, [_4@&ReScope(Misc(NodeId(31))) i32]);
+//         Validate(Acquire, [(*_4)@i32/ReScope(Misc(NodeId(46))) (imm)]);
+//         Validate(Release, [_4@&ReScope(Misc(NodeId(46))) i32]);
 //         _3 = const foo(_4) -> bb1;
 //     }
 //     bb1: {
-//         EndRegion(ReScope(Misc(NodeId(31))));
-//         EndRegion(ReScope(Remainder(BlockRemainder { block: NodeId(18), first_statement_index: 1 })));
+//         EndRegion(ReScope(Misc(NodeId(46))));
+//         EndRegion(ReScope(Remainder(BlockRemainder { block: NodeId(18), first_statement_index: 3 })));
 //         return;
 //     }
 // }