add code comments, close #14

bench.js

@@ -12,7 +12,7 @@ for (let i = 0; i < 1000000; i++) points.push(randomPoint(1000));
 console.log(`memory: ${  heapSize()}`);
 
 console.time(`index ${  points.length  } points`);
-const index = new KDBush(points, p => p.x, p => p.y, 64, Int32Array);
+const index = new KDBush(points, p => p.x, p => p.y, 64, Uint32Array);
 console.timeEnd(`index ${  points.length  } points`);
 
 console.log(`memory: ${  heapSize()}`);

src/index.js

@@ -13,6 +13,7 @@ export default class KDBush {
 
         const IndexArrayType = points.length < 65536 ? Uint16Array : Uint32Array;
 
+        // store indices to the input array and coordinates in separate typed arrays
         const ids = this.ids = new IndexArrayType(points.length);
         const coords = this.coords = new ArrayType(points.length * 2);
 
@@ -22,6 +23,7 @@ export default class KDBush {
             coords[2 * i + 1] = getY(points[i]);
         }
 
+        // kd-sort both arrays for efficient search (see comments in sort.js)
         sort(ids, coords, nodeSize, 0, ids.length - 1, 0);
     }
 

src/range.js

@@ -2,40 +2,41 @@
 export default function range(ids, coords, minX, minY, maxX, maxY, nodeSize) {
     const stack = [0, ids.length - 1, 0];
     const result = [];
-    let x, y;
 
+    // recursively search for items in range in the kd-sorted arrays
     while (stack.length) {
         const axis = stack.pop();
         const right = stack.pop();
         const left = stack.pop();
 
+        // if we reached "tree node", search linearly
         if (right - left <= nodeSize) {
             for (let i = left; i <= right; i++) {
-                x = coords[2 * i];
-                y = coords[2 * i + 1];
+                const x = coords[2 * i];
+                const y = coords[2 * i + 1];
                 if (x >= minX && x <= maxX && y >= minY && y <= maxY) result.push(ids[i]);
             }
             continue;
         }
 
-        const m = Math.floor((left + right) / 2);
-
-        x = coords[2 * m];
-        y = coords[2 * m + 1];
+        // otherwise find the middle index
+        const m = (left + right) >> 1;
 
+        // include the middle item if it's in range
+        const x = coords[2 * m];
+        const y = coords[2 * m + 1];
         if (x >= minX && x <= maxX && y >= minY && y <= maxY) result.push(ids[m]);
 
-        const nextAxis = (axis + 1) % 2;
-
+        // queue search in halves that intersect the query
         if (axis === 0 ? minX <= x : minY <= y) {
             stack.push(left);
             stack.push(m - 1);
-            stack.push(nextAxis);
+            stack.push(1 - axis);
         }
         if (axis === 0 ? maxX >= x : maxY >= y) {
             stack.push(m + 1);
             stack.push(right);
-            stack.push(nextAxis);
+            stack.push(1 - axis);
         }
     }
 

src/sort.js

@@ -1,16 +1,21 @@
 
-export default function sortKD(ids, coords, nodeSize, left, right, depth) {
+export default function sortKD(ids, coords, nodeSize, left, right, axis) {
     if (right - left <= nodeSize) return;
 
-    const m = (left + right) >> 1;
+    const m = (left + right) >> 1; // middle index
 
-    select(ids, coords, m, left, right, depth % 2);
+    // sort ids and coords around the middle index so that the halves lie
+    // either left/right or top/bottom correspondingly (taking turns)
+    select(ids, coords, m, left, right, axis);
 
-    sortKD(ids, coords, nodeSize, left, m - 1, depth + 1);
-    sortKD(ids, coords, nodeSize, m + 1, right, depth + 1);
+    // recursively kd-sort first half and second half on the opposite axis
+    sortKD(ids, coords, nodeSize, left, m - 1, 1 - axis);
+    sortKD(ids, coords, nodeSize, m + 1, right, 1 - axis);
 }
 
-function select(ids, coords, k, left, right, inc) {
+// custom Floyd-Rivest selection algorithm: sort ids and coords so that
+// [left..k-1] items are smaller than k-th item (on either x or y axis)
+function select(ids, coords, k, left, right, axis) {
 
     while (right > left) {
         if (right - left > 600) {
@@ -21,25 +26,25 @@ function select(ids, coords, k, left, right, inc) {
             const sd = 0.5 * Math.sqrt(z * s * (n - s) / n) * (m - n / 2 < 0 ? -1 : 1);
             const newLeft = Math.max(left, Math.floor(k - m * s / n + sd));
             const newRight = Math.min(right, Math.floor(k + (n - m) * s / n + sd));
-            select(ids, coords, k, newLeft, newRight, inc);
+            select(ids, coords, k, newLeft, newRight, axis);
         }
 
-        const t = coords[2 * k + inc];
+        const t = coords[2 * k + axis];
         let i = left;
         let j = right;
 
         swapItem(ids, coords, left, k);
-        if (coords[2 * right + inc] > t) swapItem(ids, coords, left, right);
+        if (coords[2 * right + axis] > t) swapItem(ids, coords, left, right);
 
         while (i < j) {
             swapItem(ids, coords, i, j);
             i++;
             j--;
-            while (coords[2 * i + inc] < t) i++;
-            while (coords[2 * j + inc] > t) j--;
+            while (coords[2 * i + axis] < t) i++;
+            while (coords[2 * j + axis] > t) j--;
         }
 
-        if (coords[2 * left + inc] === t) swapItem(ids, coords, left, j);
+        if (coords[2 * left + axis] === t) swapItem(ids, coords, left, j);
         else {
             j++;
             swapItem(ids, coords, j, right);

src/within.js

@@ -4,36 +4,38 @@ export default function within(ids, coords, qx, qy, r, nodeSize) {
     const result = [];
     const r2 = r * r;
 
+    // recursively search for items within radius in the kd-sorted arrays
     while (stack.length) {
         const axis = stack.pop();
         const right = stack.pop();
         const left = stack.pop();
 
+        // if we reached "tree node", search linearly
         if (right - left <= nodeSize) {
             for (let i = left; i <= right; i++) {
                 if (sqDist(coords[2 * i], coords[2 * i + 1], qx, qy) <= r2) result.push(ids[i]);
             }
             continue;
         }
 
-        const m = Math.floor((left + right) / 2);
+        // otherwise find the middle index
+        const m = (left + right) >> 1;
 
+        // include the middle item if it's in range
         const x = coords[2 * m];
         const y = coords[2 * m + 1];
-
         if (sqDist(x, y, qx, qy) <= r2) result.push(ids[m]);
 
-        const nextAxis = (axis + 1) % 2;
-
+        // queue search in halves that intersect the query
         if (axis === 0 ? qx - r <= x : qy - r <= y) {
             stack.push(left);
             stack.push(m - 1);
-            stack.push(nextAxis);
+            stack.push(1 - axis);
         }
         if (axis === 0 ? qx + r >= x : qy + r >= y) {
             stack.push(m + 1);
             stack.push(right);
-            stack.push(nextAxis);
+            stack.push(1 - axis);
         }
     }