Attempt to improve type stability

brenhinkeller · Dec 28, 2024 · c4d4be3 · c4d4be3
1 parent 9aef19a
commit c4d4be3
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/src/Sorting/quicksort.jl b/src/Sorting/quicksort.jl
@@ -110,7 +110,7 @@ argsortnans!(I::AbstractArray, A::AbstractArray{<:Integer}, iₗ::Int=firstindex
 end
 
 # Sort `A`, assuming no NaNs
-function quicksort!(A, iₗ=firstindex(A), iᵤ=lastindex(A))
+function quicksort!(A::TA, iₗ=firstindex(A), iᵤ=lastindex(A)) where {TA<:AbstractArray}
     if issortedrange(A, iₗ, iᵤ)
         # If already sorted, we're done here
         return A
@@ -169,13 +169,14 @@ function quicksort!(A, iₗ=firstindex(A), iᵤ=lastindex(A))
         iₚ = iₗ + Nₗ - 1
         A[iₗ], A[iₚ] = A[iₚ], A[iₗ]
         # Recurse: sort both upper and lower partitions
-        quicksort!(A, iₗ, iₚ)
-        quicksort!(A, iₚ+1, iᵤ)
+        quicksort!(A, iₗ, iₚ)::TA
+        quicksort!(A, iₚ+1, iᵤ)::TA
+        return A
     end
 end
 
 # Argsort: sort A and permute I to match `A`, assuming no NaNs
-function argsort!(I::AbstractArray, A::AbstractArray, iₗ::Int=firstindex(A), iᵤ::Int=lastindex(A))
+function argsort!(I::TI, A::TA, iₗ::Int=firstindex(A), iᵤ::Int=lastindex(A)) where {TI<:AbstractArray, TA<:AbstractArray}
     if issortedrange(A, iₗ, iᵤ)
         # If already sorted, we're done here
         return I, A
@@ -242,8 +243,9 @@ function argsort!(I::AbstractArray, A::AbstractArray, iₗ::Int=firstindex(A), i
         A[iₗ], A[iₚ] = A[iₚ], A[iₗ]
         I[iₗ], I[iₚ] = I[iₚ], I[iₗ]
         # Recurse: sort both upper and lower partitions
-        argsort!(I, A, iₗ, iₚ)
-        argsort!(I, A, iₚ+1, iᵤ)
+        argsort!(I, A, iₗ, iₚ)::Tuple{TI, TA}
+        argsort!(I, A, iₚ+1, iᵤ)::Tuple{TI, TA}
+        return I, A
     end
 end