From 53fbb4507ef6710296d974117808079dd6610d41 Mon Sep 17 00:00:00 2001
From: Kuba Jaroszewski <jakub.jaroszewski@gmail.com>
Date: Tue, 28 Feb 2023 21:03:19 +0100
Subject: [PATCH] Add benchmarks to instant-distance-py

---
 .github/workflows/rust.yml                    |  2 +-
 Makefile                                      |  9 ++--
 instant-distance-py/Cargo.toml                | 13 +++++-
 instant-distance-py/benches/all.rs            | 41 +++++++++++++++++++
 instant-distance-py/pyproject.toml            |  3 ++
 .../python/instant_distance/__init__.py       |  5 +++
 instant-distance-py/src/lib.rs                |  2 +-
 7 files changed, 69 insertions(+), 6 deletions(-)
 create mode 100644 instant-distance-py/benches/all.rs
 create mode 100644 instant-distance-py/python/instant_distance/__init__.py

diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 9edcdb8..bf86695 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -53,7 +53,7 @@ jobs:
       - name: Test Python bindings
         run: |
           sudo apt-get install -y wamerican
-          cp target/release/libinstant_distance.so instant-distance-py/test/instant_distance.so
+          cp target/release/libinstant_distance_py.so instant-distance-py/test/instant_distance.so
           PYTHONPATH=instant-distance-py/test/ python3 -m test
 
   lint:
diff --git a/Makefile b/Makefile
index ef56aec..6b0105e 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,9 @@
-test-python:
-	cargo build --release
-	cp target/release/libinstant_distance.dylib instant-distance-py/test/instant_distance.so
+instant-distance-py/test/instant_distance.so: instant-distance-py/src/lib.rs
+	RUSTFLAGS="-C target-cpu=native" cargo build --release
+	([ -f target/release/libinstant_distance_py.dylib ] && cp target/release/libinstant_distance_py.dylib instant-distance-py/test/instant_distance.so) || \
+	([ -f target/release/libinstant_distance_py.so ] && cp target/release/libinstant_distance_py.so instant-distance-py/test/instant_distance.so)
+
+test-python: instant-distance-py/test/instant_distance.so
 	PYTHONPATH=instant-distance-py/test/ python3 -m test
 
 clean:
diff --git a/instant-distance-py/Cargo.toml b/instant-distance-py/Cargo.toml
index e8e3357..ad45933 100644
--- a/instant-distance-py/Cargo.toml
+++ b/instant-distance-py/Cargo.toml
@@ -11,8 +11,11 @@ repository = "https://github.com/InstantDomain/instant-distance"
 readme = "../README.md"
 
 [lib]
+name = "instant_distance_py"
+crate-type = ["cdylib", "lib"]
+
+[package.metadata.maturin]
 name = "instant_distance"
-crate-type = ["cdylib"]
 
 [dependencies]
 bincode = "1.3.1"
@@ -20,3 +23,11 @@ instant-distance = { version = "0.6", path = "../instant-distance", features = [
 pyo3 = { version = "0.18.0", features = ["extension-module"] }
 serde = { version = "1", features = ["derive"] }
 serde-big-array = "0.5.0"
+
+[dev-dependencies]
+bencher = "0.1.5"
+rand = { version = "0.8", features = ["small_rng"] }
+
+[[bench]]
+name = "all"
+harness = false
diff --git a/instant-distance-py/benches/all.rs b/instant-distance-py/benches/all.rs
new file mode 100644
index 0000000..a727b75
--- /dev/null
+++ b/instant-distance-py/benches/all.rs
@@ -0,0 +1,41 @@
+use bencher::{benchmark_group, benchmark_main, Bencher};
+
+use instant_distance::{Builder, Point, Search};
+use instant_distance_py::FloatArray;
+use rand::{rngs::StdRng, Rng, SeedableRng};
+
+benchmark_main!(benches);
+benchmark_group!(benches, distance, build, query);
+
+fn distance(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let point_a = FloatArray([rng.gen(); 300]);
+    let point_b = FloatArray([rng.gen(); 300]);
+
+    bench.iter(|| point_a.distance(&point_b));
+}
+
+fn build(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let points = (0..1024)
+        .map(|_| FloatArray([rng.gen(); 300]))
+        .collect::<Vec<_>>();
+
+    bench.iter(|| Builder::default().seed(SEED).build_hnsw(points.clone()));
+}
+
+fn query(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let points = (0..1024)
+        .map(|_| FloatArray([rng.gen(); 300]))
+        .collect::<Vec<_>>();
+    let (hnsw, _) = Builder::default().seed(SEED).build_hnsw(points);
+    let point = FloatArray([rng.gen(); 300]);
+
+    bench.iter(|| {
+        let mut search = Search::default();
+        let _ = hnsw.search(&point, &mut search);
+    });
+}
+
+const SEED: u64 = 123456789;
diff --git a/instant-distance-py/pyproject.toml b/instant-distance-py/pyproject.toml
index 409cabc..16ce725 100644
--- a/instant-distance-py/pyproject.toml
+++ b/instant-distance-py/pyproject.toml
@@ -4,3 +4,6 @@ name = "instant-distance"
 [build-system]
 requires = ["maturin >= 0.14, < 0.15"]
 build-backend = "maturin"
+
+[tool.maturin]
+python-source = "python"
diff --git a/instant-distance-py/python/instant_distance/__init__.py b/instant-distance-py/python/instant_distance/__init__.py
new file mode 100644
index 0000000..0e28f6c
--- /dev/null
+++ b/instant-distance-py/python/instant_distance/__init__.py
@@ -0,0 +1,5 @@
+from .instant_distance import *
+
+__doc__ = instant_distance.__doc__
+if hasattr(instant_distance, "__all__"):
+    __all__ = instant_distance.__all__
diff --git a/instant-distance-py/src/lib.rs b/instant-distance-py/src/lib.rs
index bc35090..455641a 100644
--- a/instant-distance-py/src/lib.rs
+++ b/instant-distance-py/src/lib.rs
@@ -348,7 +348,7 @@ impl Neighbor {
 
 #[repr(align(32))]
 #[derive(Clone, Deserialize, Serialize)]
-struct FloatArray(#[serde(with = "BigArray")] [f32; DIMENSIONS]);
+pub struct FloatArray(#[serde(with = "BigArray")] pub [f32; DIMENSIONS]);
 
 impl TryFrom<&PyAny> for FloatArray {
     type Error = PyErr;