From 3c7e0f3fdabea0a08a003475075d75e24bedd1b0 Mon Sep 17 00:00:00 2001 From: Teddy Tennant Date: Wed, 8 Jul 2026 13:15:56 -0400 Subject: [PATCH] Fix stack overflow in quickselect selection routines on large arrays The quickselect-based selection routines in `src/sort.rs` recursed once per partition step. On pathological inputs such as a large constant (all-equal) array, Hoare's partition peels off only a single element at each step, so the recursion depth grows linearly with the array length. For large arrays this overflows the thread stack and aborts the process (e.g. taking the median of `Array1::::ones(100_000)`). Rewrite `_get_many_from_sorted_mut_unchecked` and `Sort1dExt::get_from_sorted_mut` to recurse into the sub-problem with fewer array elements and iterate on the larger one, the standard technique that bounds the recursion depth to O(log n). The computed results are unchanged; only the stack usage is bounded. Closes #86 --- src/sort.rs | 159 ++++++++++++++++++++++++++-------------------- tests/quantile.rs | 8 +++ 2 files changed, 97 insertions(+), 70 deletions(-) diff --git a/src/sort.rs b/src/sort.rs index e38e205b..e0c20024 100644 --- a/src/sort.rs +++ b/src/sort.rs @@ -96,25 +96,31 @@ pub trait Sort1dExt { } impl Sort1dExt for ArrayRef { - fn get_from_sorted_mut(&mut self, i: usize) -> A + fn get_from_sorted_mut(&mut self, mut i: usize) -> A where A: Ord + Clone, { - let n = self.len(); - if n == 1 { - self[0].clone() - } else { + // We narrow the view in place instead of recursing, so that the stack + // depth stays constant regardless of the number of partition steps. + // This matters for pathological inputs (e.g. large constant arrays), + // where each partition step only peels off a single element and the + // recursive formulation would use O(n) stack frames. + let mut array = self.view_mut(); + loop { + let n = array.len(); + if n == 1 { + return array[0].clone(); + } let mut rng = thread_rng(); let pivot_index = rng.gen_range(0..n); - let partition_index = self.partition_mut(pivot_index); + let partition_index = array.partition_mut(pivot_index); if i < partition_index { - self.slice_axis_mut(Axis(0), Slice::from(..partition_index)) - .get_from_sorted_mut(i) + array.slice_axis_inplace(Axis(0), Slice::from(..partition_index)); } else if i == partition_index { - self[i].clone() + return array[i].clone(); } else { - self.slice_axis_mut(Axis(0), Slice::from(partition_index + 1..)) - .get_from_sorted_mut(i - (partition_index + 1)) + array.slice_axis_inplace(Axis(0), Slice::from(partition_index + 1..)); + i -= partition_index + 1; } } } @@ -211,73 +217,86 @@ where /// initial element values are ignored. fn _get_many_from_sorted_mut_unchecked( mut array: ArrayViewMut1<'_, A>, - indexes: &mut [usize], - values: &mut [A], + mut indexes: &mut [usize], + mut values: &mut [A], ) where A: Ord + Clone, { - let n = array.len(); - debug_assert!(n >= indexes.len()); // because indexes must be unique and in-bounds - debug_assert_eq!(indexes.len(), values.len()); + // After each partition step we recurse into the sub-problem with *fewer* + // array elements and iterate (in this `loop`) on the larger one. This is + // the standard technique for bounding quickselect's stack depth to + // O(log n): without it, pathological inputs such as large constant arrays + // (where each partition step only peels off a single element) would use + // O(n) stack frames and overflow the stack. + loop { + let n = array.len(); + debug_assert!(n >= indexes.len()); // because indexes must be unique and in-bounds + debug_assert_eq!(indexes.len(), values.len()); - if indexes.is_empty() { - // Nothing to do in this case. - return; - } + if indexes.is_empty() { + // Nothing to do in this case. + return; + } - // At this point, `n >= 1` since `indexes.len() >= 1`. - if n == 1 { - // We can only reach this point if `indexes.len() == 1`, so we only - // need to assign the single value, and then we're done. - debug_assert_eq!(indexes.len(), 1); - values[0] = array[0].clone(); - return; - } + // At this point, `n >= 1` since `indexes.len() >= 1`. + if n == 1 { + // We can only reach this point if `indexes.len() == 1`, so we only + // need to assign the single value, and then we're done. + debug_assert_eq!(indexes.len(), 1); + values[0] = array[0].clone(); + return; + } + + // We pick a random pivot index: the corresponding element is the pivot value + let mut rng = thread_rng(); + let pivot_index = rng.gen_range(0..n); - // We pick a random pivot index: the corresponding element is the pivot value - let mut rng = thread_rng(); - let pivot_index = rng.gen_range(0..n); + // We partition the array with respect to the pivot value. + // The pivot value moves to `array_partition_index`. + // Elements strictly smaller than the pivot value have indexes < `array_partition_index`. + // Elements greater or equal to the pivot value have indexes > `array_partition_index`. + let array_partition_index = array.partition_mut(pivot_index); - // We partition the array with respect to the pivot value. - // The pivot value moves to `array_partition_index`. - // Elements strictly smaller than the pivot value have indexes < `array_partition_index`. - // Elements greater or equal to the pivot value have indexes > `array_partition_index`. - let array_partition_index = array.partition_mut(pivot_index); + // We use a divide-and-conquer strategy, splitting the indexes we are + // searching for (`indexes`) and the corresponding portions of the output + // slice (`values`) into pieces with respect to `array_partition_index`. + let (found_exact, index_split) = match indexes.binary_search(&array_partition_index) { + Ok(index) => (true, index), + Err(index) => (false, index), + }; + let (smaller_indexes, other_indexes) = indexes.split_at_mut(index_split); + let (smaller_values, other_values) = values.split_at_mut(index_split); + let (bigger_indexes, bigger_values) = if found_exact { + other_values[0] = array[array_partition_index].clone(); // Write exactly found value. + (&mut other_indexes[1..], &mut other_values[1..]) + } else { + (other_indexes, other_values) + }; - // We use a divide-and-conquer strategy, splitting the indexes we are - // searching for (`indexes`) and the corresponding portions of the output - // slice (`values`) into pieces with respect to `array_partition_index`. - let (found_exact, index_split) = match indexes.binary_search(&array_partition_index) { - Ok(index) => (true, index), - Err(index) => (false, index), - }; - let (smaller_indexes, other_indexes) = indexes.split_at_mut(index_split); - let (smaller_values, other_values) = values.split_at_mut(index_split); - let (bigger_indexes, bigger_values) = if found_exact { - other_values[0] = array[array_partition_index].clone(); // Write exactly found value. - (&mut other_indexes[1..], &mut other_values[1..]) - } else { - (other_indexes, other_values) - }; + // The indexes to the right of the pivot need to be shifted by the length + // of the removed portion (the smaller part plus the pivot itself). + bigger_indexes + .iter_mut() + .for_each(|x| *x -= array_partition_index + 1); - // We search recursively for the values corresponding to strictly smaller - // indexes to the left of `partition_index`. - _get_many_from_sorted_mut_unchecked( - array.slice_axis_mut(Axis(0), Slice::from(..array_partition_index)), - smaller_indexes, - smaller_values, - ); + // We split the array into the two disjoint sub-views, dropping the pivot + // (which is already in its final sorted position and, if requested, has + // already been written to `bigger_values`/`other_values` above). + let (smaller_array, rest) = array.split_at(Axis(0), array_partition_index); + let bigger_array = rest.slice_axis_move(Axis(0), Slice::from(1..)); - // We search recursively for the values corresponding to strictly bigger - // indexes to the right of `partition_index`. Since only the right portion - // of the array is passed in, the indexes need to be shifted by length of - // the removed portion. - bigger_indexes - .iter_mut() - .for_each(|x| *x -= array_partition_index + 1); - _get_many_from_sorted_mut_unchecked( - array.slice_axis_mut(Axis(0), Slice::from(array_partition_index + 1..)), - bigger_indexes, - bigger_values, - ); + // Recurse into the sub-problem with fewer array elements and continue + // iterating on the larger one, so the recursion depth stays O(log n). + if smaller_array.len() <= bigger_array.len() { + _get_many_from_sorted_mut_unchecked(smaller_array, smaller_indexes, smaller_values); + array = bigger_array; + indexes = bigger_indexes; + values = bigger_values; + } else { + _get_many_from_sorted_mut_unchecked(bigger_array, bigger_indexes, bigger_values); + array = smaller_array; + indexes = smaller_indexes; + values = smaller_values; + } + } } diff --git a/tests/quantile.rs b/tests/quantile.rs index 9d58071f..a103df1d 100644 --- a/tests/quantile.rs +++ b/tests/quantile.rs @@ -277,6 +277,14 @@ fn test_midpoint_overflow() { assert_eq!(median, expected_median); } +#[test] +fn quantile_mut_on_large_constant_array_does_not_overflow_stack() { + // Regression test for https://github.com/rust-ndarray/ndarray-stats/issues/86 + let mut array: Array1 = Array1::ones(100_000); + let median = array.quantile_mut(n64(0.5), &Linear).unwrap(); + assert_eq!(median, n64(1.)); +} + #[quickcheck] fn test_quantiles_mut(xs: Vec) -> bool { let v = Array::from(xs.clone());