From 2ab15e0c661eb97230adbb53139ddce70ac595e9 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Wed, 20 Jan 2021 18:07:03 +0100
Subject: [PATCH] Improve ZeroNode abstraction some more

---
 src/lib.rs   | 29 +++--------------------------
 src/types.rs | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index ab2598b..1738d49 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -412,22 +412,12 @@ fn insert<P: Point>(
 
             let found =
                 insertion.select_heuristic(&layer, M * 2, candidate_point, points, *heuristic);
-            for (i, slot) in layer[pid].0.iter_mut().enumerate() {
-                if let Some(&Candidate { pid, .. }) = found.get(i) {
-                    *slot = pid;
-                } else if *slot != PointId::invalid() {
-                    *slot = PointId::invalid();
-                } else {
-                    break;
-                }
-            }
-
+            layer[pid].rewrite(found.iter().map(|candidate| candidate.pid));
             layer[new].set(i, pid);
         } else {
             // Find the correct index to insert at to keep the neighbor's neighbors sorted
             let old = &points[pid];
-            let nearest = &layer[pid].0;
-            let idx = nearest
+            let idx = layer[pid]
                 .binary_search_by(|third| {
                     // `third` here is one of the neighbors of the new node's neighbor.
                     let third = match third {
@@ -440,20 +430,7 @@ fn insert<P: Point>(
                 })
                 .unwrap_or_else(|e| e);
 
-            // It might be possible for all the neighbor's current neighbors to be closer to our
-            // neighbor than to the new node, in which case we skip insertion of our new node's ID.
-            if idx >= nearest.len() {
-                layer[new].set(i, pid);
-                continue;
-            }
-
-            let nearest = &mut layer[pid].0;
-            if nearest[idx].is_valid() {
-                let end = (M * 2) - 1;
-                nearest.copy_within(idx..end, idx + 1);
-            }
-
-            nearest[idx] = new;
+            layer[pid].insert(idx, new);
             layer[new].set(i, pid);
         }
     }
diff --git a/src/types.rs b/src/types.rs
index 4d0a224..0682530 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -1,5 +1,5 @@
 use std::hash::Hash;
-use std::ops::{Index, IndexMut};
+use std::ops::{Deref, Index, IndexMut};
 
 use ordered_float::OrderedFloat;
 use rand::rngs::SmallRng;
@@ -82,6 +82,33 @@ impl Layer for &Vec<UpperNode> {
 pub(crate) struct ZeroNode(pub(crate) [PointId; M * 2]);
 
 impl ZeroNode {
+    pub(crate) fn rewrite(&mut self, mut iter: impl Iterator<Item = PointId>) {
+        for slot in self.0.iter_mut() {
+            if let Some(pid) = iter.next() {
+                *slot = pid;
+            } else if *slot != PointId::invalid() {
+                *slot = PointId::invalid();
+            } else {
+                break;
+            }
+        }
+    }
+
+    pub(crate) fn insert(&mut self, idx: usize, pid: PointId) {
+        // It might be possible for all the neighbor's current neighbors to be closer to our
+        // neighbor than to the new node, in which case we skip insertion of our new node's ID.
+        if idx >= self.0.len() {
+            return;
+        }
+
+        if self.0[idx].is_valid() {
+            let end = (M * 2) - 1;
+            self.0.copy_within(idx..end, idx + 1);
+        }
+
+        self.0[idx] = pid;
+    }
+
     pub(crate) fn set(&mut self, idx: usize, pid: PointId) {
         self.0[idx] = pid;
     }
@@ -93,6 +120,14 @@ impl Default for ZeroNode {
     }
 }
 
+impl Deref for ZeroNode {
+    type Target = [PointId];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
 impl Layer for &Vec<ZeroNode> {
     fn nearest_iter(&self, pid: PointId) -> NearestIter<'_> {
         NearestIter {