Fix infinite recursion, overflow, and off-by-one error in triu/tril (#1418)

* Fixes infinite recursion and off-by-one error

* Avoids overflow using saturating arithmetic

* Removes unused import

* Fixes bug for isize::MAX for triu

* Fix formatting

* Uses broadcast indices to remove D::Smaller: Copy trait bound

Author: akern40

src/tri.rs

@@ -6,18 +6,25 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use core::cmp::{max, min};
+use core::cmp::min;
 
 use num_traits::Zero;
 
-use crate::{dimension::is_layout_f, Array, ArrayBase, Axis, Data, Dimension, IntoDimension, Zip};
+use crate::{
+    dimension::{is_layout_c, is_layout_f},
+    Array,
+    ArrayBase,
+    Axis,
+    Data,
+    Dimension,
+    Zip,
+};
 
 impl<S, A, D> ArrayBase<S, D>
 where
     S: Data<Elem = A>,
     D: Dimension,
     A: Clone + Zero,
-    D::Smaller: Copy,
 {
     /// Upper triangular of an array.
     ///
@@ -30,38 +37,56 @@ where
     /// ```
     /// use ndarray::array;
     ///
-    /// let arr = array![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
-    /// let res = arr.triu(0);
-    /// assert_eq!(res, array![[1, 2, 3], [0, 5, 6], [0, 0, 9]]);
+    /// let arr = array![
+    ///     [1, 2, 3],
+    ///     [4, 5, 6],
+    ///     [7, 8, 9]
+    /// ];
+    /// assert_eq!(
+    ///     arr.triu(0),
+    ///     array![
+    ///         [1, 2, 3],
+    ///         [0, 5, 6],
+    ///         [0, 0, 9]
+    ///     ]
+    /// );
     /// ```
     pub fn triu(&self, k: isize) -> Array<A, D>
     {
         if self.ndim() <= 1 {
             return self.to_owned();
         }
-        match is_layout_f(&self.dim, &self.strides) {
-            true => {
-                let n = self.ndim();
-                let mut x = self.view();
-                x.swap_axes(n - 2, n - 1);
-                let mut tril = x.tril(-k);
-                tril.swap_axes(n - 2, n - 1);
-
-                tril
-            }
-            false => {
-                let mut res = Array::zeros(self.raw_dim());
-                Zip::indexed(self.rows())
-                    .and(res.rows_mut())
-                    .for_each(|i, src, mut dst| {
-                        let row_num = i.into_dimension().last_elem();
-                        let lower = max(row_num as isize + k, 0);
-                        dst.slice_mut(s![lower..]).assign(&src.slice(s![lower..]));
-                    });
-
-                res
-            }
+
+        // Performance optimization for F-order arrays.
+        // C-order array check prevents infinite recursion in edge cases like [[1]].
+        // k-size check prevents underflow when k == isize::MIN
+        let n = self.ndim();
+        if is_layout_f(&self.dim, &self.strides) && !is_layout_c(&self.dim, &self.strides) && k > isize::MIN {
+            let mut x = self.view();
+            x.swap_axes(n - 2, n - 1);
+            let mut tril = x.tril(-k);
+            tril.swap_axes(n - 2, n - 1);
+
+            return tril;
         }
+
+        let mut res = Array::zeros(self.raw_dim());
+        let ncols = self.len_of(Axis(n - 1));
+        let nrows = self.len_of(Axis(n - 2));
+        let indices = Array::from_iter(0..nrows);
+        Zip::from(self.rows())
+            .and(res.rows_mut())
+            .and_broadcast(&indices)
+            .for_each(|src, mut dst, row_num| {
+                let mut lower = match k >= 0 {
+                    true => row_num.saturating_add(k as usize),        // Avoid overflow
+                    false => row_num.saturating_sub(k.unsigned_abs()), // Avoid underflow, go to 0
+                };
+                lower = min(lower, ncols);
+                dst.slice_mut(s![lower..]).assign(&src.slice(s![lower..]));
+            });
+
+        res
     }
 
     /// Lower triangular of an array.
@@ -75,45 +100,65 @@ where
     /// ```
     /// use ndarray::array;
     ///
-    /// let arr = array![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
-    /// let res = arr.tril(0);
-    /// assert_eq!(res, array![[1, 0, 0], [4, 5, 0], [7, 8, 9]]);
+    /// let arr = array![
+    ///     [1, 2, 3],
+    ///     [4, 5, 6],
+    ///     [7, 8, 9]
+    /// ];
+    /// assert_eq!(
+    ///     arr.tril(0),
+    ///     array![
+    ///         [1, 0, 0],
+    ///         [4, 5, 0],
+    ///         [7, 8, 9]
+    ///     ]
+    /// );
     /// ```
     pub fn tril(&self, k: isize) -> Array<A, D>
     {
         if self.ndim() <= 1 {
             return self.to_owned();
         }
-        match is_layout_f(&self.dim, &self.strides) {
-            true => {
-                let n = self.ndim();
-                let mut x = self.view();
-                x.swap_axes(n - 2, n - 1);
-                let mut tril = x.triu(-k);
-                tril.swap_axes(n - 2, n - 1);
-
-                tril
-            }
-            false => {
-                let mut res = Array::zeros(self.raw_dim());
-                let ncols = self.len_of(Axis(self.ndim() - 1)) as isize;
-                Zip::indexed(self.rows())
-                    .and(res.rows_mut())
-                    .for_each(|i, src, mut dst| {
-                        let row_num = i.into_dimension().last_elem();
-                        let upper = min(row_num as isize + k, ncols) + 1;
-                        dst.slice_mut(s![..upper]).assign(&src.slice(s![..upper]));
-                    });
-
-                res
-            }
+
+        // Performance optimization for F-order arrays.
+        // C-order array check prevents infinite recursion in edge cases like [[1]].
+        // k-size check prevents underflow when k == isize::MIN
+        let n = self.ndim();
+        if is_layout_f(&self.dim, &self.strides) && !is_layout_c(&self.dim, &self.strides) && k > isize::MIN {
+            let mut x = self.view();
+            x.swap_axes(n - 2, n - 1);
+            let mut tril = x.triu(-k);
+            tril.swap_axes(n - 2, n - 1);
+
+            return tril;
         }
+
+        let mut res = Array::zeros(self.raw_dim());
+        let ncols = self.len_of(Axis(n - 1));
+        let nrows = self.len_of(Axis(n - 2));
+        let indices = Array::from_iter(0..nrows);
+        Zip::from(self.rows())
+            .and(res.rows_mut())
+            .and_broadcast(&indices)
+            .for_each(|src, mut dst, row_num| {
+                // let row_num = i.into_dimension().last_elem();
+                let mut upper = match k >= 0 {
+                    true => row_num.saturating_add(k as usize).saturating_add(1), // Avoid overflow
+                    false => row_num.saturating_sub((k + 1).unsigned_abs()),      // Avoid underflow
+                };
+                upper = min(upper, ncols);
+                dst.slice_mut(s![..upper]).assign(&src.slice(s![..upper]));
+            });
+
+        res
     }
 }
 
 #[cfg(test)]
 mod tests
 {
+    use core::isize;
+
     use crate::{array, dimension, Array0, Array1, Array2, Array3, ShapeBuilder};
     use alloc::vec;
 
@@ -188,6 +233,19 @@ mod tests
         assert_eq!(res, array![[1, 0, 0], [4, 5, 0], [7, 8, 9]]);
     }
 
+    #[test]
+    fn test_2d_single()
+    {
+        let x = array![[1]];
+
+        assert_eq!(x.triu(0), array![[1]]);
+        assert_eq!(x.tril(0), array![[1]]);
+        assert_eq!(x.triu(1), array![[0]]);
+        assert_eq!(x.tril(1), array![[1]]);
+        assert_eq!(x.triu(-1), array![[1]]);
+        assert_eq!(x.tril(-1), array![[0]]);
+    }
+
     #[test]
     fn test_3d()
     {
@@ -285,8 +343,25 @@ mod tests
         let res = x.triu(0);
         assert_eq!(res, array![[1, 2, 3], [0, 5, 6]]);
 
+        let res = x.tril(0);
+        assert_eq!(res, array![[1, 0, 0], [4, 5, 0]]);
+
         let x = array![[1, 2], [3, 4], [5, 6]];
         let res = x.triu(0);
         assert_eq!(res, array![[1, 2], [0, 4], [0, 0]]);
+
+        let res = x.tril(0);
+        assert_eq!(res, array![[1, 0], [3, 4], [5, 6]]);
+    }
+
+    #[test]
+    fn test_odd_k()
+    {
+        let x = array![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
+        let z = Array2::zeros([3, 3]);
+        assert_eq!(x.triu(isize::MIN), x);
+        assert_eq!(x.tril(isize::MIN), z);
+        assert_eq!(x.triu(isize::MAX), z);
+        assert_eq!(x.tril(isize::MAX), x);
     }
 }