Implement LazyBTreeMap and use it in a few places.

nnethercote · nnethercote · commit 259ae1813919 · 2018-04-26T11:55:06.000+10:00
This is a thin wrapper around BTreeMap that avoids allocating upon
creation. It speeds up some rustc-perf benchmarks by up to 3.6%.
diff --git a/src/librustc/infer/higher_ranked/mod.rs b/src/librustc/infer/higher_ranked/mod.rs
@@ -19,7 +19,7 @@ use super::{CombinedSnapshot,
 use super::combine::CombineFields;
 use super::region_constraints::{TaintDirections};
 
-use std::collections::BTreeMap;
+use rustc_data_structures::lazy_btree_map::LazyBTreeMap;
 use ty::{self, TyCtxt, Binder, TypeFoldable};
 use ty::error::TypeError;
 use ty::relate::{Relate, RelateResult, TypeRelation};
@@ -247,7 +247,8 @@ impl<'a, 'gcx, 'tcx> CombineFields<'a, 'gcx, 'tcx> {
                                              snapshot: &CombinedSnapshot<'a, 'tcx>,
                                              debruijn: ty::DebruijnIndex,
                                              new_vars: &[ty::RegionVid],
-                                             a_map: &BTreeMap<ty::BoundRegion, ty::Region<'tcx>>,
+                                             a_map: &LazyBTreeMap<ty::BoundRegion,
+                                                                  ty::Region<'tcx>>,
                                              r0: ty::Region<'tcx>)
                                              -> ty::Region<'tcx> {
             // Regions that pre-dated the LUB computation stay as they are.
@@ -343,7 +344,8 @@ impl<'a, 'gcx, 'tcx> CombineFields<'a, 'gcx, 'tcx> {
                                              snapshot: &CombinedSnapshot<'a, 'tcx>,
                                              debruijn: ty::DebruijnIndex,
                                              new_vars: &[ty::RegionVid],
-                                             a_map: &BTreeMap<ty::BoundRegion, ty::Region<'tcx>>,
+                                             a_map: &LazyBTreeMap<ty::BoundRegion,
+                                                                  ty::Region<'tcx>>,
                                              a_vars: &[ty::RegionVid],
                                              b_vars: &[ty::RegionVid],
                                              r0: ty::Region<'tcx>)
@@ -412,7 +414,7 @@ impl<'a, 'gcx, 'tcx> CombineFields<'a, 'gcx, 'tcx> {
 
         fn rev_lookup<'a, 'gcx, 'tcx>(infcx: &InferCtxt<'a, 'gcx, 'tcx>,
                                       span: Span,
-                                      a_map: &BTreeMap<ty::BoundRegion, ty::Region<'tcx>>,
+                                      a_map: &LazyBTreeMap<ty::BoundRegion, ty::Region<'tcx>>,
                                       r: ty::Region<'tcx>) -> ty::Region<'tcx>
         {
             for (a_br, a_r) in a_map {
@@ -435,7 +437,7 @@ impl<'a, 'gcx, 'tcx> CombineFields<'a, 'gcx, 'tcx> {
 }
 
 fn var_ids<'a, 'gcx, 'tcx>(fields: &CombineFields<'a, 'gcx, 'tcx>,
-                           map: &BTreeMap<ty::BoundRegion, ty::Region<'tcx>>)
+                           map: &LazyBTreeMap<ty::BoundRegion, ty::Region<'tcx>>)
                            -> Vec<ty::RegionVid> {
     map.iter()
        .map(|(_, &r)| match *r {
diff --git a/src/librustc/infer/mod.rs b/src/librustc/infer/mod.rs
@@ -28,9 +28,9 @@ use ty::error::{ExpectedFound, TypeError, UnconstrainedNumeric};
 use ty::fold::TypeFoldable;
 use ty::relate::RelateResult;
 use traits::{self, ObligationCause, PredicateObligations};
+use rustc_data_structures::lazy_btree_map::LazyBTreeMap;
 use rustc_data_structures::unify as ut;
 use std::cell::{Cell, RefCell, Ref, RefMut};
-use std::collections::BTreeMap;
 use std::fmt;
 use syntax::ast;
 use errors::DiagnosticBuilder;
@@ -187,7 +187,7 @@ pub struct InferCtxt<'a, 'gcx: 'a+'tcx, 'tcx: 'a> {
 
 /// A map returned by `skolemize_late_bound_regions()` indicating the skolemized
 /// region that each late-bound region was replaced with.
-pub type SkolemizationMap<'tcx> = BTreeMap<ty::BoundRegion, ty::Region<'tcx>>;
+pub type SkolemizationMap<'tcx> = LazyBTreeMap<ty::BoundRegion, ty::Region<'tcx>>;
 
 /// See `error_reporting` module for more details
 #[derive(Clone, Debug)]
@@ -1216,7 +1216,7 @@ impl<'a, 'gcx, 'tcx> InferCtxt<'a, 'gcx, 'tcx> {
         span: Span,
         lbrct: LateBoundRegionConversionTime,
         value: &ty::Binder<T>)
-        -> (T, BTreeMap<ty::BoundRegion, ty::Region<'tcx>>)
+        -> (T, LazyBTreeMap<ty::BoundRegion, ty::Region<'tcx>>)
         where T : TypeFoldable<'tcx>
     {
         self.tcx.replace_late_bound_regions(
diff --git a/src/librustc/ty/fold.rs b/src/librustc/ty/fold.rs
@@ -43,8 +43,8 @@ use middle::const_val::ConstVal;
 use hir::def_id::DefId;
 use ty::{self, Binder, Ty, TyCtxt, TypeFlags};
 
+use rustc_data_structures::lazy_btree_map::LazyBTreeMap;
 use std::fmt;
-use std::collections::BTreeMap;
 use util::nodemap::FxHashSet;
 
 /// The TypeFoldable trait is implemented for every type that can be folded.
@@ -334,7 +334,7 @@ struct RegionReplacer<'a, 'gcx: 'a+'tcx, 'tcx: 'a> {
     tcx: TyCtxt<'a, 'gcx, 'tcx>,
     current_depth: u32,
     fld_r: &'a mut (dyn FnMut(ty::BoundRegion) -> ty::Region<'tcx> + 'a),
-    map: BTreeMap<ty::BoundRegion, ty::Region<'tcx>>
+    map: LazyBTreeMap<ty::BoundRegion, ty::Region<'tcx>>
 }
 
 impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
@@ -349,7 +349,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
     pub fn replace_late_bound_regions<T,F>(self,
         value: &Binder<T>,
         mut f: F)
-        -> (T, BTreeMap<ty::BoundRegion, ty::Region<'tcx>>)
+        -> (T, LazyBTreeMap<ty::BoundRegion, ty::Region<'tcx>>)
         where F : FnMut(ty::BoundRegion) -> ty::Region<'tcx>,
               T : TypeFoldable<'tcx>,
     {
@@ -462,7 +462,7 @@ impl<'a, 'gcx, 'tcx> RegionReplacer<'a, 'gcx, 'tcx> {
             tcx,
             current_depth: 1,
             fld_r,
-            map: BTreeMap::default()
+            map: LazyBTreeMap::default()
         }
     }
 }
diff --git a/src/librustc_data_structures/lazy_btree_map.rs b/src/librustc_data_structures/lazy_btree_map.rs
@@ -0,0 +1,108 @@
+// Copyright 2018 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.
+
+use std::collections::btree_map;
+use std::collections::BTreeMap;
+
+/// A thin wrapper around BTreeMap that avoids allocating upon creation.
+///
+/// Vec, HashSet and HashMap all have the nice feature that they don't do any
+/// heap allocation when creating a new structure of the default size. In
+/// contrast, BTreeMap *does* allocate in that situation. The compiler uses
+/// B-Tree maps in some places such that many maps are created but few are
+/// inserted into, so having a BTreeMap alternative that avoids allocating on
+/// creation is a performance win.
+///
+/// Only a fraction of BTreeMap's functionality is currently supported.
+/// Additional functionality should be added on demand.
+#[derive(Debug)]
+pub struct LazyBTreeMap<K, V>(Option<BTreeMap<K, V>>);
+
+impl<K, V> LazyBTreeMap<K, V> {
+    pub fn new() -> LazyBTreeMap<K, V> {
+        LazyBTreeMap(None)
+    }
+
+    pub fn iter(&self) -> Iter<K, V> {
+        Iter(self.0.as_ref().map(|btm| btm.iter()))
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.as_ref().map_or(true, |btm| btm.is_empty())
+    }
+}
+
+impl<K: Ord, V> LazyBTreeMap<K, V> {
+    fn instantiate(&mut self) -> &mut BTreeMap<K, V> {
+        if let Some(ref mut btm) = self.0 {
+            btm
+        } else {
+            let btm = BTreeMap::new();
+            self.0 = Some(btm);
+            self.0.as_mut().unwrap()
+        }
+    }
+
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        self.instantiate().insert(key, value)
+    }
+
+    pub fn entry(&mut self, key: K) -> btree_map::Entry<K, V> {
+        self.instantiate().entry(key)
+    }
+
+    pub fn values<'a>(&'a self) -> Values<'a, K, V> {
+        Values(self.0.as_ref().map(|btm| btm.values()))
+    }
+}
+
+impl<K: Ord, V> Default for LazyBTreeMap<K, V> {
+    fn default() -> LazyBTreeMap<K, V> {
+        LazyBTreeMap::new()
+    }
+}
+
+impl<'a, K: 'a, V: 'a> IntoIterator for &'a LazyBTreeMap<K, V> {
+    type Item = (&'a K, &'a V);
+    type IntoIter = Iter<'a, K, V>;
+
+    fn into_iter(self) -> Iter<'a, K, V> {
+        self.iter()
+    }
+}
+
+pub struct Iter<'a, K: 'a, V: 'a>(Option<btree_map::Iter<'a, K, V>>);
+
+impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
+    type Item = (&'a K, &'a V);
+
+    fn next(&mut self) -> Option<(&'a K, &'a V)> {
+        self.0.as_mut().and_then(|iter| iter.next())
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.as_ref().map_or_else(|| (0, Some(0)), |iter| iter.size_hint())
+    }
+}
+
+pub struct Values<'a, K: 'a, V: 'a>(Option<btree_map::Values<'a, K, V>>);
+
+impl<'a, K, V> Iterator for Values<'a, K, V> {
+    type Item = &'a V;
+
+    fn next(&mut self) -> Option<&'a V> {
+        self.0.as_mut().and_then(|values| values.next())
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.as_ref().map_or_else(|| (0, Some(0)), |values| values.size_hint())
+    }
+}
+
diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs
@@ -56,6 +56,7 @@ pub mod bitvec;
 pub mod graph;
 pub mod indexed_set;
 pub mod indexed_vec;
+pub mod lazy_btree_map;
 pub mod obligation_forest;
 pub mod sip128;
 pub mod snapshot_map;
