refactor region value bitmatrix

4 files changed, 270 insertions, 110 deletions

diff --git a/src/librustc_mir/borrow_check/nll/region_infer/dump_mir.rs b/src/librustc_mir/borrow_check/nll/region_infer/dump_mir.rs
index 5477308bde9..69ecafa66ae 100644
--- a/src/librustc_mir/borrow_check/nll/region_infer/dump_mir.rs
+++ b/src/librustc_mir/borrow_check/nll/region_infer/dump_mir.rs
@@ -70,7 +70,7 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         with_msg: &mut FnMut(&str) -> io::Result<()>,
     ) -> io::Result<()> {
         for region in self.definitions.indices() {
-            let value = self.region_value_str_from_matrix(&self.liveness_constraints, region);
+            let value = self.liveness_constraints.region_value_str(region);
             if value != "{}" {
                 with_msg(&format!("{:?} live at {}", region, value))?;
             }
diff --git a/src/librustc_mir/borrow_check/nll/region_infer/mod.rs b/src/librustc_mir/borrow_check/nll/region_infer/mod.rs
index b2e2ccc5d0b..047b78dd55c 100644
--- a/src/librustc_mir/borrow_check/nll/region_infer/mod.rs
+++ b/src/librustc_mir/borrow_check/nll/region_infer/mod.rs
@@ -19,15 +19,15 @@ use rustc::mir::{ClosureOutlivesRequirement, ClosureRegionRequirements, Location
 use rustc::ty::{self, RegionVid};
 use rustc_data_structures::indexed_vec::IndexVec;
 use rustc_data_structures::fx::FxHashSet;
-use rustc_data_structures::bitvec::BitMatrix;
-use rustc_data_structures::indexed_vec::Idx;
-use std::collections::BTreeMap;
 use std::fmt;
+use std::rc::Rc;
 use syntax_pos::Span;
 
 mod annotation;
 mod dump_mir;
 mod graphviz;
+mod values;
+use self::values::{RegionValueElements, RegionValues};
 
 pub struct RegionInferenceContext<'tcx> {
     /// Contains the definition for every region variable.  Region
@@ -36,27 +36,22 @@ pub struct RegionInferenceContext<'tcx> {
     /// from as well as its final inferred value.
     definitions: IndexVec<RegionVid, RegionDefinition<'tcx>>,
 
+    /// Maps from points/universal-regions to a `RegionElementIndex`.
+    elements: Rc<RegionValueElements>,
+
     /// The liveness constraints added to each region. For most
     /// regions, these start out empty and steadily grow, though for
     /// each universally quantified region R they start out containing
     /// the entire CFG and `end(R)`.
-    ///
-    /// In this `BitMatrix` representation, the rows are the region
-    /// variables and the columns are the free regions and MIR locations.
-    liveness_constraints: BitMatrix,
+    liveness_constraints: RegionValues,
 
     /// The final inferred values of the inference variables; `None`
     /// until `solve` is invoked.
-    inferred_values: Option<BitMatrix>,
+    inferred_values: Option<RegionValues>,
 
     /// The constraints we have accumulated and used during solving.
     constraints: Vec<Constraint>,
 
-    /// A map from each MIR Location to its column index in
-    /// `liveness_constraints`/`inferred_values`. (The first N columns are
-    /// the free regions.)
-    point_indices: BTreeMap<Location, usize>,
-
     /// Information about the universally quantified regions in scope
     /// on this function and their (known) relations to one another.
     universal_regions: UniversalRegions<'tcx>,
@@ -112,19 +107,16 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         let num_region_variables = var_origins.len();
         let num_universal_regions = universal_regions.len();
 
-        let mut num_points = 0;
-        let mut point_indices = BTreeMap::new();
-
+        let mut points = Vec::new();
         for (block, block_data) in mir.basic_blocks().iter_enumerated() {
             for statement_index in 0..block_data.statements.len() + 1 {
-                let location = Location {
+                points.push(Location {
                     block,
                     statement_index,
-                };
-                point_indices.insert(location, num_universal_regions + num_points);
-                num_points += 1;
+                });
             }
         }
+        let elements = &Rc::new(RegionValueElements::new(points, num_universal_regions));
 
         // Create a RegionDefinition for each inference variable.
         let definitions = var_origins
@@ -134,13 +126,10 @@ impl<'tcx> RegionInferenceContext<'tcx> {
 
         let mut result = Self {
             definitions,
-            liveness_constraints: BitMatrix::new(
-                num_region_variables,
-                num_universal_regions + num_points,
-            ),
+            elements: elements.clone(),
+            liveness_constraints: RegionValues::new(elements, num_region_variables),
             inferred_values: None,
             constraints: Vec::new(),
-            point_indices,
             universal_regions,
         };
 
@@ -186,14 +175,12 @@ impl<'tcx> RegionInferenceContext<'tcx> {
             self.definitions[variable].is_universal = true;
 
             // Add all nodes in the CFG to liveness constraints
-            for (_location, point_index) in &self.point_indices {
-                self.liveness_constraints
-                    .add(variable.index(), *point_index);
+            for point_index in self.elements.all_point_indices() {
+                self.liveness_constraints.add(variable, point_index);
             }
 
             // Add `end(X)` into the set for X.
-            self.liveness_constraints
-                .add(variable.index(), variable.index());
+            self.liveness_constraints.add(variable, variable);
         }
     }
 
@@ -217,32 +204,7 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         let inferred_values = self.inferred_values
             .as_ref()
             .expect("region values not yet inferred");
-        self.region_contains_point_in_matrix(inferred_values, r, p)
-    }
-
-    /// True if given region `r` contains the point `p`, when
-    /// evaluated in the set of region values `matrix`.
-    fn region_contains_point_in_matrix(
-        &self,
-        matrix: &BitMatrix,
-        r: RegionVid,
-        p: Location,
-    ) -> bool {
-        let point_index = self.point_indices
-            .get(&p)
-            .expect("point index should be known");
-        matrix.contains(r.index(), *point_index)
-    }
-
-    /// True if given region `r` contains the `end(s)`, when
-    /// evaluated in the set of region values `matrix`.
-    fn region_contains_region_in_matrix(
-        &self,
-        matrix: &BitMatrix,
-        r: RegionVid,
-        s: RegionVid,
-    ) -> bool {
-        matrix.contains(r.index(), s.index())
+        inferred_values.contains(r, p)
     }
 
     /// Returns access to the value of `r` for debugging purposes.
@@ -251,43 +213,21 @@ impl<'tcx> RegionInferenceContext<'tcx> {
             .as_ref()
             .expect("region values not yet inferred");
 
-        self.region_value_str_from_matrix(inferred_values, r)
-    }
-
-    fn region_value_str_from_matrix(&self,
-                                    matrix: &BitMatrix,
-                                    r: RegionVid) -> String {
-        let mut result = String::new();
-        result.push_str("{");
-        let mut sep = "";
-
-        for &point in self.point_indices.keys() {
-            if self.region_contains_point_in_matrix(matrix, r, point) {
-                result.push_str(&format!("{}{:?}", sep, point));
-                sep = ", ";
-            }
-        }
-
-        for fr in (0..self.universal_regions.len()).map(RegionVid::new) {
-            if self.region_contains_region_in_matrix(matrix, r, fr) {
-                result.push_str(&format!("{}{:?}", sep, fr));
-                sep = ", ";
-            }
-        }
-
-        result.push_str("}");
-
-        result
+        inferred_values.region_value_str(r)
     }
 
     /// Indicates that the region variable `v` is live at the point `point`.
     pub(super) fn add_live_point(&mut self, v: RegionVid, point: Location) -> bool {
         debug!("add_live_point({:?}, {:?})", v, point);
         assert!(self.inferred_values.is_none(), "values already inferred");
-        let point_index = self.point_indices
-            .get(&point)
-            .expect("point index should be known");
-        self.liveness_constraints.add(v.index(), *point_index)
+        debug!("add_live_point: @{:?}", point);
+
+        let element = self.elements.index(point);
+        if self.liveness_constraints.add(v, element) {
+            true
+        } else {
+            false
+        }
     }
 
     /// Indicates that the region variable `sup` must outlive `sub` is live at the point `point`.
@@ -386,16 +326,12 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         outlives_requirements: &mut Vec<ClosureOutlivesRequirement>,
     ) {
         let inferred_values = self.inferred_values.as_ref().unwrap();
-        let longer_value = inferred_values.iter(longer_fr.index());
 
         debug!("check_universal_region(fr={:?})", longer_fr);
 
         // Find every region `o` such that `fr: o`
         // (because `fr` includes `end(o)`).
-        let shorter_frs = longer_value
-            .take_while(|&i| i < self.universal_regions.len())
-            .map(RegionVid::new);
-        for shorter_fr in shorter_frs {
+        for shorter_fr in inferred_values.universal_regions_outlived_by(longer_fr) {
             // If it is known that `fr: o`, carry on.
             if self.universal_regions.outlives(longer_fr, shorter_fr) {
                 continue;
@@ -512,20 +448,22 @@ impl<'tcx> RegionInferenceContext<'tcx> {
 
     fn copy(
         &self,
-        inferred_values: &mut BitMatrix,
+        inferred_values: &mut RegionValues,
         mir: &Mir<'tcx>,
         from_region: RegionVid,
         to_region: RegionVid,
-        start_point: Location,
+        constraint_point: Location,
     ) -> bool {
         let mut changed = false;
 
         let mut stack = vec![];
         let mut visited = FxHashSet();
 
-        stack.push(start_point);
+        stack.push(constraint_point);
         while let Some(p) = stack.pop() {
-            if !self.region_contains_point_in_matrix(inferred_values, from_region, p) {
+            let point_index = self.elements.index(p);
+
+            if !inferred_values.contains(from_region, point_index) {
                 debug!("            not in from-region");
                 continue;
             }
@@ -535,8 +473,8 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                 continue;
             }
 
-            let point_index = self.point_indices.get(&p).unwrap();
-            changed |= inferred_values.add(to_region.index(), *point_index);
+            let new = inferred_values.add(to_region, point_index);
+            changed |= new;
 
             let block_data = &mir[p.block];
             let successor_points = if p.statement_index < block_data.statements.len() {
@@ -564,13 +502,8 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                 // If we reach the END point in the graph, then copy
                 // over any skolemized end points in the `from_region`
                 // and make sure they are included in the `to_region`.
-                let universal_region_indices = inferred_values
-                    .iter(from_region.index())
-                    .take_while(|&i| i < self.universal_regions.len())
-                    .collect::<Vec<_>>();
-                for fr in &universal_region_indices {
-                    changed |= inferred_values.add(to_region.index(), *fr);
-                }
+                changed |=
+                    inferred_values.add_universal_regions_outlived_by(from_region, to_region);
             } else {
                 stack.extend(successor_points);
             }
diff --git a/src/librustc_mir/borrow_check/nll/region_infer/values.rs b/src/librustc_mir/borrow_check/nll/region_infer/values.rs
new file mode 100644
index 00000000000..cda199859e4
--- /dev/null
+++ b/src/librustc_mir/borrow_check/nll/region_infer/values.rs
@@ -0,0 +1,227 @@
+use std::rc::Rc;
+use rustc_data_structures::bitvec::BitMatrix;
+use rustc_data_structures::indexed_vec::Idx;
+use rustc::mir::Location;
+use rustc::ty::RegionVid;
+
+/// Maps between the various kinds of elements of a region value to
+/// the internal indices that w use.
+pub(super) struct RegionValueElements {
+    points: Vec<Location>,
+    num_universal_regions: usize,
+}
+
+impl RegionValueElements {
+    pub(super) fn new(points: Vec<Location>, num_universal_regions: usize) -> Self {
+        Self {
+            points,
+            num_universal_regions,
+        }
+    }
+
+    /// Converts an element of a region value into a `RegionElementIndex`.
+    pub(super) fn index<T: ToElementIndex>(&self, elem: T) -> RegionElementIndex {
+        elem.to_element_index(self)
+    }
+
+    /// Iterates over the `RegionElementIndex` for all points in the CFG.
+    pub(super) fn all_point_indices<'a>(&'a self) -> impl Iterator<Item = RegionElementIndex> + 'a {
+        (0..self.points.len()).map(move |i| RegionElementIndex::new(i + self.num_universal_regions))
+    }
+
+    /// Iterates over the `RegionElementIndex` for all points in the CFG.
+    pub(super) fn all_universal_region_indices(&self) -> impl Iterator<Item = RegionElementIndex> {
+        (0..self.num_universal_regions).map(move |i| RegionElementIndex::new(i))
+    }
+
+    /// Converts a particular `RegionElementIndex` to the `RegionElement` it represents.
+    pub(super) fn to_element(&self, i: RegionElementIndex) -> RegionElement {
+        if let Some(r) = self.to_universal_region(i) {
+            RegionElement::UniversalRegion(r)
+        } else {
+            RegionElement::Location(self.points[i.index() - self.num_universal_regions])
+        }
+    }
+
+    /// Converts a particular `RegionElementIndex` to a universal
+    /// region, if that is what it represents. Returns `None`
+    /// otherwise.
+    pub(super) fn to_universal_region(&self, i: RegionElementIndex) -> Option<RegionVid> {
+        if i.index() < self.num_universal_regions {
+            Some(RegionVid::new(i.index()))
+        } else {
+            None
+        }
+    }
+}
+
+/// A newtype for the integers that represent one of the possible
+/// elements in a region. These are the rows in the `BitMatrix` that
+/// is used to store the values of all regions. They have the following
+/// convention:
+///
+/// - The first N indices represent free regions (where N = universal_regions.len()).
+/// - The remainder represent the points in the CFG (see `point_indices` map).
+///
+/// You can convert a `RegionElementIndex` into a `RegionElement`
+/// using the `to_region_elem` method.
+newtype_index!(RegionElementIndex { DEBUG_FORMAT = "RegionElementIndex({})" });
+
+/// An individual element in a region value -- the value of a
+/// particular region variable consists of a set of these elements.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+pub(super) enum RegionElement {
+    /// A point in the control-flow graph.
+    Location(Location),
+
+    /// An in-scope, universally quantified region (e.g., a liftime parameter).
+    UniversalRegion(RegionVid),
+}
+
+
+pub(super) trait ToElementIndex {
+    fn to_element_index(self, elements: &RegionValueElements) -> RegionElementIndex;
+}
+
+impl ToElementIndex for Location {
+    fn to_element_index(self, elements: &RegionValueElements) -> RegionElementIndex {
+        let index = elements.points.binary_search(&self).unwrap();
+        RegionElementIndex::new(index + elements.num_universal_regions)
+    }
+}
+
+impl ToElementIndex for RegionVid {
+    fn to_element_index(self, elements: &RegionValueElements) -> RegionElementIndex {
+        assert!(self.index() < elements.num_universal_regions);
+        RegionElementIndex::new(self.index())
+    }
+}
+
+impl ToElementIndex for RegionElementIndex {
+    fn to_element_index(self, _elements: &RegionValueElements) -> RegionElementIndex {
+        self
+    }
+}
+
+/// Stores the values for a set of regions. These are stored in a
+/// compact `BitMatrix` representation, with one row per region
+/// variable. The columns consist of either universal regions or
+/// points in the CFG.
+#[derive(Clone)]
+pub(super) struct RegionValues {
+    elements: Rc<RegionValueElements>,
+    matrix: BitMatrix,
+}
+
+impl RegionValues {
+    pub(super) fn new(elements: &Rc<RegionValueElements>, num_region_variables: usize) -> Self {
+        assert!(
+            elements.num_universal_regions <= num_region_variables,
+            "universal regions are a subset of the region variables"
+        );
+
+        let num_columns = elements.points.len() + elements.num_universal_regions;
+
+        Self {
+            elements: elements.clone(),
+            matrix: BitMatrix::new(num_region_variables, num_columns),
+        }
+    }
+
+    /// Adds the given element to the value for the given region. Returns true if
+    /// the element is newly added (i.e., was not already present).
+    pub(super) fn add<E: ToElementIndex>(&mut self, r: RegionVid, elem: E) -> bool {
+        let i = self.elements.index(elem);
+        if self.matrix.add(r.index(), i.index()) {
+            debug!("add(r={:?}, i={:?})", r, self.elements.to_element(i));
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Adds all the universal regions outlived by `from_region` to
+    /// `to_region`.
+    pub(super) fn add_universal_regions_outlived_by(
+        &mut self,
+        from_region: RegionVid,
+        to_region: RegionVid,
+    ) -> bool {
+        // FIXME. We could optimize this by improving
+        // `BitMatrix::merge` so it does not always merge an entire
+        // row.
+        debug!("add_universal_regions_outlived_by(from_region={:?}, to_region={:?})",
+               from_region, to_region);
+        let mut changed = false;
+        for elem in self.elements.all_universal_region_indices() {
+            if self.contains(from_region, elem) {
+                changed |= self.add(to_region, elem);
+            }
+        }
+        changed
+    }
+
+    /// True if the region `r` contains the given element.
+    pub(super) fn contains<E: ToElementIndex>(&self, r: RegionVid, elem: E) -> bool {
+        let i = self.elements.index(elem);
+        self.matrix.contains(r.index(), i.index())
+    }
+
+    /// Iterate over the value of the region `r`, yielding up element
+    /// indices. You may prefer `universal_regions_outlived_by` or
+    /// `elements_contained_in`.
+    pub(super) fn element_indices_contained_in<'a>(
+        &'a self,
+        r: RegionVid,
+    ) -> impl Iterator<Item = RegionElementIndex> + 'a {
+        self.matrix
+            .iter(r.index())
+            .map(move |i| RegionElementIndex::new(i))
+    }
+
+    /// Returns just the universal regions that are contained in a given region's value.
+    pub(super) fn universal_regions_outlived_by<'a>(
+        &'a self,
+        r: RegionVid,
+    ) -> impl Iterator<Item = RegionVid> + 'a {
+        self.element_indices_contained_in(r)
+            .map(move |i| self.elements.to_universal_region(i))
+            .take_while(move |v| v.is_some()) // universal regions are a prefix
+            .map(move |v| v.unwrap())
+    }
+
+    /// Returns all the elements contained in a given region's value.
+    pub(super) fn elements_contained_in<'a>(
+        &'a self,
+        r: RegionVid,
+    ) -> impl Iterator<Item = RegionElement> + 'a {
+        self.element_indices_contained_in(r)
+            .map(move |r| self.elements.to_element(r))
+    }
+
+    /// Returns a "pretty" string value of the region. Meant for debugging.
+    pub(super) fn region_value_str(&self, r: RegionVid) -> String {
+        let mut result = String::new();
+        result.push_str("{");
+
+        for (index, element) in self.elements_contained_in(r).enumerate() {
+            if index > 0 {
+                result.push_str(", ");
+            }
+
+            match element {
+                RegionElement::Location(l) => {
+                    result.push_str(&format!("{:?}", l));
+                }
+
+                RegionElement::UniversalRegion(fr) => {
+                    result.push_str(&format!("{:?}", fr));
+                }
+            }
+        }
+
+        result.push_str("}");
+
+        result
+    }
+}
diff --git a/src/test/mir-opt/nll/named-lifetimes-basic.rs b/src/test/mir-opt/nll/named-lifetimes-basic.rs
index 0c42585a528..f3a57c08840 100644
--- a/src/test/mir-opt/nll/named-lifetimes-basic.rs
+++ b/src/test/mir-opt/nll/named-lifetimes-basic.rs
@@ -33,10 +33,10 @@ fn main() {
 // | '_#3r    | Local    | ['_#3r]
 // |
 // | Inferred Region Values
-// | '_#0r    | {bb0[0], bb0[1], '_#0r}
-// | '_#1r    | {bb0[0], bb0[1], '_#1r}
-// | '_#2r    | {bb0[0], bb0[1], '_#2r}
-// | '_#3r    | {bb0[0], bb0[1], '_#3r}
+// | '_#0r    | {'_#0r, bb0[0], bb0[1]}
+// | '_#1r    | {'_#1r, bb0[0], bb0[1]}
+// | '_#2r    | {'_#2r, bb0[0], bb0[1]}
+// | '_#3r    | {'_#3r, bb0[0], bb0[1]}
 // |
 // ...
 // fn use_x(_1: &'_#1r mut i32, _2: &'_#2r u32, _3: &'_#1r u32, _4: &'_#3r u32) -> bool {