Visitar URL original
fix: prevent stack overflow in F-order triu/tril (#1616) · rust-ndarray/ndarray@b301b7a · GitHub
Skip to content

Commit b301b7a

Browse files
cesterciancestercian
authored
fix: prevent stack overflow in F-order triu/tril (#1616)
The F-order fast paths in triu and tril swap the last two axes and call each other. That is only safe for 2D arrays, where the swap yields a C-contiguous view. Higher-D F-and-not-C views (for example a C-order transpose followed by insert_axis) can remain F-and-not-C after the swap, so the two functions recurse until they overflow the stack. Restrict the shortcut to ndim == 2 so normal 2D .f() arrays keep the optimization, while the issue #1615 repro takes the non-recursive path.
1 parent ace739e commit b301b7a

1 file changed

Lines changed: 36 additions & 7 deletions

File tree

‎src/tri.rs‎

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,11 +55,18 @@ where
5555
return self.to_owned();
5656
}
5757

58-
// Performance optimization for F-order arrays.
59-
// C-order array check prevents infinite recursion in edge cases like [[1]].
58+
// Performance optimization for 2D F-order arrays: swap axes so the callee
59+
// hits the C-contiguous path, then swap back to preserve F-order.
60+
// Restricting to ndim == 2 avoids infinite recursion on higher-D views that
61+
// remain F-and-not-C after swapping the last two axes (see #1615).
62+
// The C-order check still guards the both-C-and-F `[[1]]` edge case.
6063
// k-size check prevents underflow when k == isize::MIN
6164
let n = self.ndim();
62-
if is_layout_f(self._dim(), self._strides()) && !is_layout_c(self._dim(), self._strides()) && k > isize::MIN {
65+
if n == 2
66+
&& is_layout_f(self._dim(), self._strides())
67+
&& !is_layout_c(self._dim(), self._strides())
68+
&& k > isize::MIN
69+
{
6370
let mut x = self.view();
6471
x.swap_axes(n - 2, n - 1);
6572
let mut tril = x.tril(-k);
@@ -120,11 +127,18 @@ where
120127
return self.to_owned();
121128
}
122129

123-
// Performance optimization for F-order arrays.
124-
// C-order array check prevents infinite recursion in edge cases like [[1]].
130+
// Performance optimization for 2D F-order arrays: swap axes so the callee
131+
// hits the C-contiguous path, then swap back to preserve F-order.
132+
// Restricting to ndim == 2 avoids infinite recursion on higher-D views that
133+
// remain F-and-not-C after swapping the last two axes (see #1615).
134+
// The C-order check still guards the both-C-and-F `[[1]]` edge case.
125135
// k-size check prevents underflow when k == isize::MIN
126136
let n = self.ndim();
127-
if is_layout_f(self._dim(), self._strides()) && !is_layout_c(self._dim(), self._strides()) && k > isize::MIN {
137+
if n == 2
138+
&& is_layout_f(self._dim(), self._strides())
139+
&& !is_layout_c(self._dim(), self._strides())
140+
&& k > isize::MIN
141+
{
128142
let mut x = self.view();
129143
x.swap_axes(n - 2, n - 1);
130144
let mut triu = x.triu(-k);
@@ -159,7 +173,7 @@ where
159173
#[cfg(test)]
160174
mod tests
161175
{
162-
use crate::{array, dimension, Array0, Array1, Array2, Array3, ShapeBuilder};
176+
use crate::{array, dimension, Array0, Array1, Array2, Array3, Axis, ShapeBuilder};
163177
use alloc::vec;
164178

165179
#[test]
@@ -364,4 +378,19 @@ mod tests
364378
assert_eq!(x.triu(isize::MAX), z);
365379
assert_eq!(x.tril(isize::MAX), x);
366380
}
381+
382+
#[test]
383+
fn test_f_view_after_insert_axis_does_not_recurse()
384+
{
385+
// Regression for https://github.com/rust-ndarray/ndarray/issues/1615:
386+
// transpose of a C-order (3, 4) array, then insert_axis(Axis(2)), yields a
387+
// (4, 3, 1) view that is F-and-not-C both before and after swapping the last
388+
// two axes. The F-order shortcut must not recurse into triu/tril.
389+
let a: Array2<f64> = Array2::from_shape_fn((3, 4), |(i, j)| (i * 4 + j) as f64);
390+
let widened = a.t().insert_axis(Axis(2));
391+
let standard = widened.as_standard_layout();
392+
393+
assert_eq!(widened.triu(0), standard.triu(0));
394+
assert_eq!(widened.tril(0), standard.tril(0));
395+
}
367396
}

0 commit comments

Comments
 (0)