From 6cf9bfefbca7853a55e34a0929d15c0420d434b8 Mon Sep 17 00:00:00 2001
From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com>
Date: Fri, 22 May 2026 13:15:22 +0200
Subject: [PATCH 1/2] Parallelize sorting of `PrecomputedNeighborhoodSearch`
 and preserve across `copy_neighborhood_search`

---
 src/nhs_precomputed.jl   |  5 +++--
 src/vector_of_vectors.jl | 14 ++++++++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/nhs_precomputed.jl b/src/nhs_precomputed.jl
index c7ddceca..cf23a23e 100644
--- a/src/nhs_precomputed.jl
+++ b/src/nhs_precomputed.jl
@@ -192,7 +192,7 @@ function initialize_neighbor_lists!(neighbor_lists::DynamicVectorOfVectors,
     # Fill neighbor lists
     foreach_point_neighbor(x, y, neighborhood_search;
                            parallelization_backend) do point, neighbor, _, _
-        pushat!(neighbor_lists, point, neighbor)
+        @inbounds pushat!(neighbor_lists, point, neighbor)
     end
 
     if sort_neighbor_lists
@@ -250,7 +250,8 @@ function copy_neighborhood_search(nhs::PrecomputedNeighborhoodSearch,
                                                      update_neighborhood_search,
                                                      backend = typeof(nhs.neighbor_lists),
                                                      transpose_backend,
-                                                     max_neighbors = max_neighbors_)
+                                                     max_neighbors = max_neighbors_,
+                                                     sort_neighbor_lists = nhs.sort_neighbor_lists)
 end
 
 @inline function freeze_neighborhood_search(search::PrecomputedNeighborhoodSearch)
diff --git a/src/vector_of_vectors.jl b/src/vector_of_vectors.jl
index c05cddbd..d00f2e9d 100644
--- a/src/vector_of_vectors.jl
+++ b/src/vector_of_vectors.jl
@@ -73,13 +73,13 @@ end
     # Outer bounds check
     @boundscheck checkbounds(vov, i)
 
-    lengths[i] += 1
+    @inbounds lengths[i] += 1
 
     # Inner bounds check
     @boundscheck check_list_bounds(vov, i)
 
     # Activate new entry in column `i`
-    backend[lengths[i], i] = value
+    @inbounds backend[lengths[i], i] = value
 
     return vov
 end
@@ -175,6 +175,16 @@ end
 
 # Sort each inner vector
 @inline function sorteach!(vov::DynamicVectorOfVectors)
+    @threaded default_backend(vov.backend) for i in axes(vov.backend, 2)
+        # QuickSort is ~1.6x faster here than the default
+        sort!(view(vov.backend, 1:vov.lengths[i], i), alg=QuickSort)
+    end
+
+    return vov
+end
+
+# GPU version
+@inline function sorteach!(vov::DynamicVectorOfVectors{<:Any, <:Any, <:AbstractGPUArray})
     # TODO remove this check when Metal supports sorting
     if nameof(typeof(default_backend(vov.backend))) == :MetalBackend
         @warn "sorting neighbor lists is not supported on Metal. Skipping sort."

From e19d132091598e4ffb83da19357f6e8ed6e7e32a Mon Sep 17 00:00:00 2001
From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com>
Date: Fri, 22 May 2026 14:16:21 +0200
Subject: [PATCH 2/2] Fix sorting length

---
 src/nhs_precomputed.jl   | 2 +-
 src/vector_of_vectors.jl | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/nhs_precomputed.jl b/src/nhs_precomputed.jl
index cf23a23e..97b11288 100644
--- a/src/nhs_precomputed.jl
+++ b/src/nhs_precomputed.jl
@@ -192,7 +192,7 @@ function initialize_neighbor_lists!(neighbor_lists::DynamicVectorOfVectors,
     # Fill neighbor lists
     foreach_point_neighbor(x, y, neighborhood_search;
                            parallelization_backend) do point, neighbor, _, _
-        @inbounds pushat!(neighbor_lists, point, neighbor)
+        pushat!(neighbor_lists, point, neighbor)
     end
 
     if sort_neighbor_lists
diff --git a/src/vector_of_vectors.jl b/src/vector_of_vectors.jl
index d00f2e9d..72632dae 100644
--- a/src/vector_of_vectors.jl
+++ b/src/vector_of_vectors.jl
@@ -175,9 +175,9 @@ end
 
 # Sort each inner vector
 @inline function sorteach!(vov::DynamicVectorOfVectors)
-    @threaded default_backend(vov.backend) for i in axes(vov.backend, 2)
+    @threaded default_backend(vov.backend) for i in eachindex(vov)
         # QuickSort is ~1.6x faster here than the default
-        sort!(view(vov.backend, 1:vov.lengths[i], i), alg=QuickSort)
+        sort!(view(vov.backend, 1:vov.lengths[i], i), alg = QuickSort)
     end
 
     return vov