Simplify construction preparation

instant-distance/src/lib.rs

@@ -235,6 +235,7 @@ where
         }
         sizes.push((num, num));
         sizes.reverse();
+        let top = LayerId(sizes.len() - 1);
 
         // Give all points a random layer and sort the list of nodes by descending order for
         // construction. This allows us to copy higher layers to lower layers as construction
@@ -246,34 +247,15 @@ where
             .collect::<Vec<_>>();
         shuffled.sort_unstable();
 
-        let mut new_points = Vec::with_capacity(points.len());
-        let mut new_nodes = Vec::with_capacity(points.len());
         let mut out = vec![INVALID; points.len()];
-        for (_, idx) in shuffled {
+        let points = shuffled
-            let pid = PointId(new_nodes.len() as u32);
+            .into_iter()
-            let layer = sizes
+            .enumerate()
-                .iter()
+            .map(|(i, (_, idx))| {
-                .enumerate()
+                out[idx] = PointId(i as u32);
-                .find_map(|(i, &size)| match (pid.0 as usize) < size.1 {
+                points[idx].clone()
-                    true => Some(i),
+            })
-                    false => None,
+            .collect::<Vec<_>>();
-                })
-                .unwrap();
-
-            new_points.push(points[idx].clone());
-            new_nodes.push((LayerId(sizes.len() - layer - 1), pid));
-            out[idx] = pid;
-        }
-        let (points, nodes) = (new_points, new_nodes);
-        debug_assert_eq!(nodes.last().unwrap().0, LayerId(0));
-        debug_assert_eq!(nodes.first().unwrap().0, LayerId(sizes.len() - 1));
-
-        // The layer from the first node is our top layer, or the zero layer if we have no nodes.
-
-        let top = match nodes.first() {
-            Some((top, _)) => *top,
-            None => LayerId(0),
-        };
 
         // Figure out how many nodes will go on each layer. This helps us allocate memory capacity
         // for each layer in advance, and also helps enable batch insertion of points.
@@ -313,17 +295,15 @@ where
                 bar.set_message(format!("Building index (layer {})", layer.0));
             }
 
+            let inserter = |pid| state.insert(pid, layer, &layers);
+
             let end = range.end;
             if layer == top {
-                nodes[range].into_iter().for_each(|(_, pid)| {
+                range.into_iter().for_each(|i| inserter(PointId(i as u32)))
-                    let node = state.zero[*pid].write();
-                    state.insert(*pid, node, layer, &layers);
-                })
             } else {
-                nodes[range].into_par_iter().for_each(|(_, pid)| {
+                range
-                    let node = state.zero[*pid].write();
+                    .into_par_iter()
-                    state.insert(*pid, node, layer, &layers);
+                    .for_each(|i| inserter(PointId(i as u32)));
-                });
             }
 
             // For layers above the zero layer, make a copy of the current state of the zero layer
@@ -431,21 +411,14 @@ struct Construction<'a, P: Point> {
 impl<'a, P: Point> Construction<'a, P> {
     /// Insert new node in the zero layer
     ///
-    /// * `new`: the `PointId` for the new node
+    /// * `new` is the `PointId` for the new node
-    /// * `insertion`: a `Search` for shrinking a neighbor set (only used with heuristic neighbor selection)
+    /// * `layer` contains all the nodes at the current layer
-    /// * `search`: the result for searching potential neighbors for the new node
+    /// * `layers` refers to the existing higher-level layers
-    /// *  `layer` contains all the nodes at the current layer
-    /// * `points` is a slice of all the points in the index
     ///
     /// Creates the new node, initializing its `nearest` array and updates the nearest neighbors
     /// for the new node's neighbors if necessary before appending the new node to the layer.
-    fn insert(
+    fn insert(&self, new: PointId, layer: LayerId, layers: &[Vec<UpperNode>]) {
-        &self,
+        let mut node = self.zero[new].write();
-        new: PointId,
-        mut node: parking_lot::RwLockWriteGuard<ZeroNode>,
-        layer: LayerId,
-        layers: &[Vec<UpperNode>],
-    ) {
         let (mut search, mut insertion) = self.pool.pop();
         insertion.ef = self.ef_construction;