Use AHash for borrow checking hash tables.

Author: adamreichold

Cargo.toml

@@ -15,6 +15,7 @@ keywords = ["python", "numpy", "ffi", "pyo3"]
 license = "BSD-2-Clause"
 
 [dependencies]
+ahash = "0.7"
 libc = "0.2"
 num-complex = ">= 0.2, < 0.5"
 num-integer = "0.1"
@@ -23,7 +24,7 @@ ndarray = ">= 0.13, < 0.16"
 pyo3 = { version = "0.16", default-features = false, features = ["macros"] }
 
 [dev-dependencies]
-pyo3 = { version = "0.16", features = ["auto-initialize"] }
+pyo3 = { version = "0.16", default-features = false, features = ["auto-initialize"] }
 
 [workspace]
 members = ["examples/*"]

benches/borrow.rs

@@ -9,7 +9,7 @@ use pyo3::Python;
 #[bench]
 fn initial_shared_borrow(bencher: &mut Bencher) {
     Python::with_gil(|py| {
-        let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
+        let array = PyArray::<f64, _>::zeros(py, (6, 5, 4, 3, 2, 1), false);
 
         bencher.iter(|| {
             let array = black_box(array);
@@ -22,7 +22,7 @@ fn initial_shared_borrow(bencher: &mut Bencher) {
 #[bench]
 fn additional_shared_borrow(bencher: &mut Bencher) {
     Python::with_gil(|py| {
-        let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
+        let array = PyArray::<f64, _>::zeros(py, (6, 5, 4, 3, 2, 1), false);
 
         let _shared = (0..128).map(|_| array.readonly()).collect::<Vec<_>>();
 
@@ -37,7 +37,7 @@ fn additional_shared_borrow(bencher: &mut Bencher) {
 #[bench]
 fn exclusive_borrow(bencher: &mut Bencher) {
     Python::with_gil(|py| {
-        let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
+        let array = PyArray::<f64, _>::zeros(py, (6, 5, 4, 3, 2, 1), false);
 
         bencher.iter(|| {
             let array = black_box(array);

src/borrow.rs

@@ -163,11 +163,12 @@
 
 use std::any::type_name;
 use std::cell::UnsafeCell;
-use std::collections::hash_map::{Entry, HashMap};
+use std::collections::hash_map::Entry;
 use std::fmt;
 use std::mem::size_of;
 use std::ops::Deref;
 
+use ahash::AHashMap;
 use ndarray::{ArrayView, ArrayViewMut, Dimension, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
 use num_integer::gcd;
 use pyo3::{FromPyObject, PyAny, PyResult};
@@ -235,7 +236,7 @@ impl BorrowKey {
     }
 }
 
-type BorrowFlagsInner = HashMap<usize, HashMap<BorrowKey, isize>>;
+type BorrowFlagsInner = AHashMap<usize, AHashMap<BorrowKey, isize>>;
 
 struct BorrowFlags(UnsafeCell<Option<BorrowFlagsInner>>);
 
@@ -248,7 +249,7 @@ impl BorrowFlags {
 
     #[allow(clippy::mut_from_ref)]
     unsafe fn get(&self) -> &mut BorrowFlagsInner {
-        (*self.0.get()).get_or_insert_with(HashMap::new)
+        (*self.0.get()).get_or_insert_with(AHashMap::new)
     }
 
     fn acquire<T, D>(&self, array: &PyArray<T, D>) -> Result<(), BorrowError>
@@ -292,7 +293,7 @@ impl BorrowFlags {
                 }
             }
             Entry::Vacant(entry) => {
-                let mut same_base_arrays = HashMap::with_capacity(1);
+                let mut same_base_arrays = AHashMap::with_capacity(1);
                 same_base_arrays.insert(key, 1);
                 entry.insert(same_base_arrays);
             }
@@ -363,7 +364,7 @@ impl BorrowFlags {
                 }
             }
             Entry::Vacant(entry) => {
-                let mut same_base_arrays = HashMap::with_capacity(1);
+                let mut same_base_arrays = AHashMap::with_capacity(1);
                 same_base_arrays.insert(key, -1);
                 entry.insert(same_base_arrays);
             }