Remove linear searches of borrow stacks

Adding a HashMap from tags to stack locations lets us find a granting
item in amortized constant time, or with a linear search of the very
small array of untagged locations.

Additionally, we keep a range which denotes where there might be an item
granting Unique permission in the stack, so that when we invalidate
Uniques we do not need to scan much of the stack, and often scan nothing
at all.

Author: saethlin

src/stacked_borrows.rs

@@ -71,15 +71,35 @@ impl fmt::Debug for Item {
 }
 
 /// Extra per-location state.
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug)]
 pub struct Stack {
     /// Used *mostly* as a stack; never empty.
     /// Invariants:
     /// * Above a `SharedReadOnly` there can only be more `SharedReadOnly`.
     /// * Except for `Untagged`, no tag occurs in the stack more than once.
     borrows: Vec<Item>,
+    /// We often want to know what (if any) tag in the stack exists for some `SbTag`.
+    /// The tag we're looking for does not occur often enough at the end of the stack for a linear
+    /// search to be efficient, so we use the uniqueness guarantee to keep a map from tag to
+    /// position in the stack.
+    cache: FxHashMap<PtrId, usize>,
+    /// `Untagged` may occur multiple times so we store it outside of the map.
+    untagged: Vec<usize>,
+    /// On a read, we need to disable all `Unique` above the granting item. We can avoid most of
+    /// this scan by keeping track of the region of the borrow stack that may contain `Unique`s.
+    first_unique: usize,
+    last_unique: usize,
+}
+
+impl PartialEq for Stack {
+    fn eq(&self, other: &Self) -> bool {
+        // All the semantics of Stack are in self.borrows, everything else is caching
+        self.borrows == other.borrows
+    }
 }
 
+impl Eq for Stack {}
+
 /// Extra per-allocation state.
 #[derive(Clone, Debug)]
 pub struct Stacks {
@@ -256,17 +276,23 @@ impl<'tcx> Stack {
     /// Find the item granting the given kind of access to the given tag, and return where
     /// it is on the stack.
     fn find_granting(&self, access: AccessKind, tag: SbTag) -> Option<usize> {
-        self.borrows
-            .iter()
-            .enumerate() // we also need to know *where* in the stack
-            .rev() // search top-to-bottom
-            // Return permission of first item that grants access.
-            // We require a permission with the right tag, ensuring U3 and F3.
-            .find_map(
-                |(idx, item)| {
-                    if tag == item.tag && item.perm.grants(access) { Some(idx) } else { None }
-                },
-            )
+        match tag {
+            SbTag::Untagged => {
+                // Search top-to-bottom
+                for idx in self.untagged.iter().rev() {
+                    // Return permission of the first item that grants access.
+                    // We require a permission with the right tag, ensuring U3 and F3.
+                    if self.borrows[*idx].perm.grants(access) {
+                        return Some(*idx);
+                    }
+                }
+                return None;
+            }
+            SbTag::Tagged(id) =>
+                self.cache.get(&id).and_then(|idx| {
+                    if self.borrows[*idx].perm.grants(access) { Some(*idx) } else { None }
+                }),
+        }
     }
 
     /// Find the first write-incompatible item above the given one --
@@ -356,6 +382,23 @@ impl<'tcx> Stack {
             for item in self.borrows.drain(first_incompatible_idx..).rev() {
                 trace!("access: popping item {:?}", item);
                 Stack::check_protector(&item, Some((tag, access)), global)?;
+                // The drain removes elements from `borrows`, but we need to remove those elements
+                // from the lookup cache too.
+                if let SbTag::Tagged(id) = item.tag {
+                    assert!(self.cache.remove(&id).is_some());
+                }
+            }
+
+            while self.untagged.last() >= Some(&first_incompatible_idx) {
+                self.untagged.pop();
+            }
+            if first_incompatible_idx <= self.first_unique {
+                // We removed all the Unique items
+                self.first_unique = 0;
+                self.last_unique = 0;
+            } else {
+                // Ensure the range doesn't extend past the new top of the stack
+                self.last_unique = self.last_unique.min(first_incompatible_idx.saturating_sub(1));
             }
         } else {
             // On a read, *disable* all `Unique` above the granting item.  This ensures U2 for read accesses.
@@ -366,14 +409,26 @@ impl<'tcx> Stack {
             // This pattern occurs a lot in the standard library: create a raw pointer, then also create a shared
             // reference and use that.
             // We *disable* instead of removing `Unique` to avoid "connecting" two neighbouring blocks of SRWs.
-            for idx in ((granting_idx + 1)..self.borrows.len()).rev() {
-                let item = &mut self.borrows[idx];
-                if item.perm == Permission::Unique {
-                    trace!("access: disabling item {:?}", item);
-                    Stack::check_protector(item, Some((tag, access)), global)?;
-                    item.perm = Permission::Disabled;
+            if granting_idx < self.last_unique {
+                // add 1 so we don't disable the granting item
+                let lower = self.first_unique.max(granting_idx + 1);
+                //let upper = (self.last_unique + 1).min(self.borrows.len());
+                for item in &mut self.borrows[lower..=self.last_unique] {
+                    if item.perm == Permission::Unique {
+                        trace!("access: disabling item {:?}", item);
+                        Stack::check_protector(item, Some((tag, access)), global)?;
+                        item.perm = Permission::Disabled;
+                    }
                 }
             }
+            if granting_idx < self.first_unique {
+                // We disabled all Unique items
+                self.first_unique = 0;
+                self.last_unique = 0;
+            } else {
+                // Truncate the range to granting_idx
+                self.last_unique = self.last_unique.min(granting_idx);
+            }
         }
 
         // Done.
@@ -401,6 +456,11 @@ impl<'tcx> Stack {
             Stack::check_protector(&item, None, global)?;
         }
 
+        self.cache.clear();
+        self.untagged.clear();
+        self.first_unique = 0;
+        self.last_unique = 0;
+
         Ok(())
     }
 
@@ -459,6 +519,43 @@ impl<'tcx> Stack {
         } else {
             trace!("reborrow: adding item {:?}", new);
             self.borrows.insert(new_idx, new);
+            // The above insert changes the meaning of every index in the cache >= new_idx, so now
+            // we need to find every one of those indexes and increment it.
+
+            // Adjust the possibly-unique range if an insert occurs before or within it
+            if self.first_unique >= new_idx {
+                self.first_unique += 1;
+            }
+            if self.last_unique >= new_idx {
+                self.last_unique += 1;
+            }
+            if new.perm == Permission::Unique {
+                // Make sure the possibly-unique range contains the new borrow
+                self.first_unique = self.first_unique.min(new_idx);
+                self.last_unique = self.last_unique.max(new_idx);
+            }
+            // Repair the lookup cache; indices have shifted if an insert was not at the end
+            if new_idx != self.borrows.len() - 1 {
+                for item in &self.borrows[new_idx + 1..] {
+                    if let SbTag::Tagged(id) = item.tag {
+                        *self.cache.get_mut(&id).unwrap() += 1;
+                    }
+                }
+                for idx in self.untagged.iter_mut().rev().take_while(|idx| **idx >= new_idx) {
+                    *idx += 1;
+                }
+            }
+            // We've now made a (conceptual) hole in the tag lookup cache.
+            // Nothing maps to new_idx, because we've adjusted everything >= it.
+            match new.tag {
+                SbTag::Untagged => {
+                    self.untagged.push(new_idx);
+                    self.untagged.sort();
+                }
+                SbTag::Tagged(id) => {
+                    self.cache.insert(id, new_idx);
+                }
+            }
         }
 
         Ok(())
@@ -471,7 +568,17 @@ impl<'tcx> Stacks {
     /// Creates new stack with initial tag.
     fn new(size: Size, perm: Permission, tag: SbTag) -> Self {
         let item = Item { perm, tag, protector: None };
-        let stack = Stack { borrows: vec![item] };
+        let mut cache = FxHashMap::default();
+        let mut untagged = Vec::new();
+        match item.tag {
+            SbTag::Untagged => {
+                untagged.push(0);
+            }
+            SbTag::Tagged(id) => {
+                cache.insert(id, 0);
+            }
+        }
+        let stack = Stack { borrows: vec![item], cache, untagged, first_unique: 0, last_unique: 0 };
 
         Stacks { stacks: RefCell::new(RangeMap::new(size, stack)) }
     }