CleanEndRegions: do not clean regions that occur in types in validation statements

2 files changed, 44 insertions, 9 deletions

diff --git a/src/librustc_mir/transform/clean_end_regions.rs b/src/librustc_mir/transform/clean_end_regions.rs
index 36125f94543..28311a5e68c 100644
--- a/src/librustc_mir/transform/clean_end_regions.rs
+++ b/src/librustc_mir/transform/clean_end_regions.rs
@@ -24,13 +24,14 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc::middle::region::CodeExtent;
 use rustc::mir::transform::{MirPass, MirSource};
 use rustc::mir::{BasicBlock, Location, Mir, Rvalue, Statement, StatementKind};
-use rustc::mir::visit::{MutVisitor, Visitor};
-use rustc::ty::{RegionKind, TyCtxt};
+use rustc::mir::visit::{MutVisitor, Visitor, Lookup};
+use rustc::ty::{Ty, RegionKind, TyCtxt};
 
 pub struct CleanEndRegions;
 
 struct GatherBorrowedRegions {
     seen_regions: FxHashSet<CodeExtent>,
+    in_validation_statement: bool,
 }
 
 struct DeleteTrivialEndRegions<'a> {
@@ -42,7 +43,7 @@ impl MirPass for CleanEndRegions {
                           _tcx: TyCtxt<'a, 'tcx, 'tcx>,
                           _source: MirSource,
                           mir: &mut Mir<'tcx>) {
-        let mut gather = GatherBorrowedRegions { seen_regions: FxHashSet() };
+        let mut gather = GatherBorrowedRegions { seen_regions: FxHashSet(), in_validation_statement: false };
         gather.visit_mir(mir);
 
         let mut delete = DeleteTrivialEndRegions { seen_regions: &mut gather.seen_regions };
@@ -54,6 +55,7 @@ impl<'tcx> Visitor<'tcx> for GatherBorrowedRegions {
     fn visit_rvalue(&mut self,
                     rvalue: &Rvalue<'tcx>,
                     location: Location) {
+        // Gather regions that are used for borrows
         if let Rvalue::Ref(r, _, _) = *rvalue {
             if let RegionKind::ReScope(ce) = *r {
                 self.seen_regions.insert(ce);
@@ -61,6 +63,31 @@ impl<'tcx> Visitor<'tcx> for GatherBorrowedRegions {
         }
         self.super_rvalue(rvalue, location);
     }
+
+    fn visit_statement(&mut self,
+                       block: BasicBlock,
+                       statement: &Statement<'tcx>,
+                       location: Location) {
+        self.in_validation_statement = match statement.kind {
+            StatementKind::Validate(..) => true,
+            _ => false,
+        };
+        self.super_statement(block, statement, location);
+        self.in_validation_statement = false;
+    }
+
+    fn visit_ty(&mut self, ty: &Ty<'tcx>, _: Lookup) {
+        // Gather regions that occur in types inside AcquireValid/ReleaseValid statements
+        if self.in_validation_statement {
+            for re in ty.walk().flat_map(|t| t.regions()) {
+                match *re {
+                    RegionKind::ReScope(ce) => { self.seen_regions.insert(ce); }
+                    _ => {},
+                }
+            }
+        }
+        self.super_ty(ty);
+    }
 }
 
 impl<'a, 'tcx> MutVisitor<'tcx> for DeleteTrivialEndRegions<'a> {
diff --git a/src/librustc_mir/transform/erase_regions.rs b/src/librustc_mir/transform/erase_regions.rs
index da9032685e0..12b1c549ffe 100644
--- a/src/librustc_mir/transform/erase_regions.rs
+++ b/src/librustc_mir/transform/erase_regions.rs
@@ -11,6 +11,8 @@
 //! This pass erases all early-bound regions from the types occuring in the MIR.
 //! We want to do this once just before trans, so trans does not have to take
 //! care erasing regions all over the place.
+//! NOTE:  We do NOT erase regions of statements that are relevant for "types-as-contracts"-validation,
+//! namely, AcquireValid, ReleaseValid, and EndRegion.
 
 use rustc::ty::subst::Substs;
 use rustc::ty::{Ty, TyCtxt, ClosureSubsts};
@@ -20,20 +22,24 @@ use rustc::mir::transform::{MirPass, MirSource};
 
 struct EraseRegionsVisitor<'a, 'tcx: 'a> {
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
+    in_validation_statement: bool,
 }
 
 impl<'a, 'tcx> EraseRegionsVisitor<'a, 'tcx> {
     pub fn new(tcx: TyCtxt<'a, 'tcx, 'tcx>) -> Self {
         EraseRegionsVisitor {
-            tcx: tcx
+            tcx: tcx,
+            in_validation_statement: false,
         }
     }
 }
 
 impl<'a, 'tcx> MutVisitor<'tcx> for EraseRegionsVisitor<'a, 'tcx> {
     fn visit_ty(&mut self, ty: &mut Ty<'tcx>, _: Lookup) {
-        let old_ty = *ty;
-        *ty = self.tcx.erase_regions(&old_ty);
+        if !self.in_validation_statement {
+            *ty = self.tcx.erase_regions(&{*ty});
+        }
+        self.super_ty(ty);
     }
 
     fn visit_substs(&mut self, substs: &mut &'tcx Substs<'tcx>, _: Location) {
@@ -71,10 +77,12 @@ impl<'a, 'tcx> MutVisitor<'tcx> for EraseRegionsVisitor<'a, 'tcx> {
                        block: BasicBlock,
                        statement: &mut Statement<'tcx>,
                        location: Location) {
-        if let StatementKind::EndRegion(_) = statement.kind {
-            statement.kind = StatementKind::Nop;
-        }
+        self.in_validation_statement = match statement.kind {
+            StatementKind::Validate(..) => true,
+            _ => false,
+        };
         self.super_statement(block, statement, location);
+        self.in_validation_statement = false;
     }
 }