Handle closures. Add some more tests.

4 files changed, 156 insertions, 14 deletions

diff --git a/src/librustc_mir/transform/add_validation.rs b/src/librustc_mir/transform/add_validation.rs
index a3ec6af7603..6f136624f0a 100644
--- a/src/librustc_mir/transform/add_validation.rs
+++ b/src/librustc_mir/transform/add_validation.rs
@@ -87,6 +87,17 @@ fn fn_contains_unsafe<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource) ->
     use rustc::hir::intravisit::{self, Visitor};
     use rustc::hir::map::Node;
 
+    fn block_is_unsafe(block: &hir::Block) -> bool {
+        use rustc::hir::BlockCheckMode::*;
+
+        match block.rules {
+            UnsafeBlock(_) | PushUnsafeBlock(_) => true,
+            // For PopUnsafeBlock, we don't actually know -- but we will always also check all
+            // parent blocks, so we can safely declare the PopUnsafeBlock to not be unsafe.
+            DefaultBlock | PopUnsafeBlock(_) => false,
+        }
+    }
+
     let fn_node_id = match src {
         MirSource::Fn(node_id) => node_id,
         _ => return false, // only functions can have unsafe
@@ -101,8 +112,35 @@ fn fn_contains_unsafe<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource) ->
     match tcx.hir.find(fn_node_id) {
         Some(Node::NodeItem(item)) => finder.visit_item(item),
         Some(Node::NodeImplItem(item)) => finder.visit_impl_item(item),
+        Some(Node::NodeExpr(item)) => {
+            // This is a closure.
+            // We also have to walk up the parents and check that there is no unsafe block
+            // there.
+            let mut cur = fn_node_id;
+            loop {
+                // Go further upwards.
+                let parent = tcx.hir.get_parent_node(cur);
+                if cur == parent {
+                    break;
+                }
+                cur = parent;
+                // Check if this is a a block
+                match tcx.hir.find(cur) {
+                    Some(Node::NodeExpr(&hir::Expr { node: hir::ExprBlock(ref block), ..})) => {
+                        if block_is_unsafe(&*block) {
+                            // We can bail out here.
+                            return true;
+                        }
+                    }
+                    _ => {},
+                }
+            }
+            // Finally, visit the closure itself.
+            finder.visit_expr(item);
+        }
         Some(_) | None =>
-            bug!("Expected method or function, found {}", tcx.hir.node_to_string(fn_node_id)),
+            bug!("Expected function, method or closure, found {}",
+                 tcx.hir.node_to_string(fn_node_id)),
     };
 
     impl<'b, 'tcx> Visitor<'tcx> for FindUnsafe<'b, 'tcx> {
@@ -113,7 +151,7 @@ fn fn_contains_unsafe<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource) ->
         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");
+            assert!(!self.found_unsafe, "We should never see a fn when we already saw unsafe");
             let is_unsafe = match fk {
                 intravisit::FnKind::ItemFn(_, _, unsafety, ..) => unsafety == hir::Unsafety::Unsafe,
                 intravisit::FnKind::Method(_, sig, ..) => sig.unsafety == hir::Unsafety::Unsafe,
@@ -129,20 +167,15 @@ fn fn_contains_unsafe<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource) ->
         }
 
         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);
-                }
-            };
+            if block_is_unsafe(b) {
+                // We found an unsafe block.  We can stop searching.
+                self.found_unsafe = true;
+            } else {
+                // No unsafe block here, go on searching.
+                intravisit::walk_block(self, b);
+            }
         }
     }
 
diff --git a/src/test/mir-opt/validate_1.rs b/src/test/mir-opt/validate_1.rs
index c8ea2bc2544..b85d9261e4a 100644
--- a/src/test/mir-opt/validate_1.rs
+++ b/src/test/mir-opt/validate_1.rs
@@ -21,8 +21,15 @@ impl Test {
 fn main() {
     let mut x = 0;
     Test.foo(&mut x);
+
+    // Also test closures
+    let c = |x: &mut i32| { let y = &*x; *y };
+    c(&mut x);
 }
 
+// FIXME: Also test code generated inside the closure, make sure it has validation.  Unfortunately,
+// the interesting lines of code also contain name of the source file, so we cannot test for it.
+
 // END RUST SOURCE
 // START rustc.node10.EraseRegions.after.mir
 //     bb0: {
diff --git a/src/test/mir-opt/validate_4.rs b/src/test/mir-opt/validate_4.rs
new file mode 100644
index 00000000000..49acaccd86a
--- /dev/null
+++ b/src/test/mir-opt/validate_4.rs
@@ -0,0 +1,58 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// ignore-tidy-linelength
+// compile-flags: -Z verbose -Z mir-emit-validate=1
+
+// Make sure unsafe fns and fns with an unsafe block only get restricted validation.
+
+unsafe fn write_42(x: *mut i32) -> bool {
+    *x = 42;
+    true
+}
+
+fn test(x: &mut i32) {
+    unsafe { write_42(x) };
+}
+
+fn main() {
+    test(&mut 0);
+
+    let test_closure = unsafe { |x: &mut i32| write_42(x) };
+    test_closure(&mut 0);
+}
+
+// FIXME: Also test code generated inside the closure, make sure it only does restricted validation
+// because it is entirely inside an unsafe block.  Unfortunately, the interesting lines of code also
+// contain name of the source file, so we cannot test for it.
+
+// END RUST SOURCE
+// START rustc.node4.EraseRegions.after.mir
+// fn write_42(_1: *mut i32) -> bool {
+//     bb0: {
+//         Validate(Acquire, [_1: *mut i32]);
+//         Validate(Release, [_1: *mut i32]);
+//         return;
+//     }
+// }
+// END rustc.node4.EraseRegions.after.mir
+// START rustc.node17.EraseRegions.after.mir
+// fn test(_1: &ReErased mut i32) -> () {
+//     bb0: {
+//         Validate(Acquire, [_1: &ReFree(DefId { krate: CrateNum(0), node: DefIndex(4) => validate_4/8cd878b::test[0] }, BrAnon(0)) mut i32]);
+//         Validate(Release, [_1: &ReFree(DefId { krate: CrateNum(0), node: DefIndex(4) => validate_4/8cd878b::test[0] }, BrAnon(0)) mut i32]);
+//         _3 = const write_42(_4) -> bb1;
+//     }
+//     bb1: {
+//         Validate(Acquire, [_3: bool]);
+//         Validate(Release, [_3: bool]);
+//     }
+// }
+// END rustc.node17.EraseRegions.after.mir
diff --git a/src/test/mir-opt/validate_5.rs b/src/test/mir-opt/validate_5.rs
new file mode 100644
index 00000000000..1831f9dd713
--- /dev/null
+++ b/src/test/mir-opt/validate_5.rs
@@ -0,0 +1,44 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// ignore-tidy-linelength
+// compile-flags: -Z verbose -Z mir-emit-validate=2
+
+// Make sure unsafe fns and fns with an unsafe block only get full validation.
+
+unsafe fn write_42(x: *mut i32) -> bool {
+    *x = 42;
+    true
+}
+
+fn test(x: &mut i32) {
+    unsafe { write_42(x) };
+}
+
+fn main() {
+    test(&mut 0);
+
+    let test_closure = unsafe { |x: &mut i32| write_42(x) };
+    test_closure(&mut 0);
+}
+
+// FIXME: Also test code generated inside the closure, make sure it has validation.  Unfortunately,
+// the interesting lines of code also contain name of the source file, so we cannot test for it.
+
+// END RUST SOURCE
+// START rustc.node17.EraseRegions.after.mir
+// fn test(_1: &ReErased mut i32) -> () {
+//     bb0: {
+//         Validate(Acquire, [_1: &ReFree(DefId { krate: CrateNum(0), node: DefIndex(4) => validate_5/8cd878b::test[0] }, BrAnon(0)) mut i32]);
+//         Validate(Release, [_4: *mut i32]);
+//         _3 = const write_42(_4) -> bb1;
+//     }
+// }
+// END rustc.node17.EraseRegions.after.mir