From b35f3f35c0d9f62520f97fa1fe108f1a0444f097 Mon Sep 17 00:00:00 2001
From: Leandro Santiago <leandrosansilva@gmail.com>
Date: Thu, 20 Feb 2025 22:27:06 +0100
Subject: [PATCH] Add Matrix::apply_with_location public function

It works like apply, but passes row and column to the callback function
---
 src/base/matrix.rs   | 19 +++++++++++++++++++
 tests/core/matrix.rs | 11 +++++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/base/matrix.rs b/src/base/matrix.rs
index 77f6c1c88..2e0122eb8 100644
--- a/src/base/matrix.rs
+++ b/src/base/matrix.rs
@@ -994,6 +994,25 @@ impl<T, R: Dim, C: Dim, S: RawStorage<T, R, C>> Matrix<T, R, C, S> {
         }
     }
 
+    /// Applies a closure `f` to modify each component of `self`. Unlike `apply`,
+    /// `f` also gets passed the row and column index, i.e. `f(row, col, value)`.
+    #[inline]
+    pub fn apply_with_location<F: FnMut(usize, usize, &mut T)>(&mut self, mut f: F)
+    where
+        S: RawStorageMut<T, R, C>,
+    {
+        let (nrows, ncols) = self.shape();
+
+        for j in 0..ncols {
+            for i in 0..nrows {
+                unsafe {
+                    let e = self.data.get_unchecked_mut(i, j);
+                    f(i, j, e)
+                }
+            }
+        }
+    }
+
     /// Replaces each component of `self` by the result of a closure `f` applied on its components
     /// joined with the components from `rhs`.
     #[inline]
diff --git a/tests/core/matrix.rs b/tests/core/matrix.rs
index 50c0e426a..2091c1260 100644
--- a/tests/core/matrix.rs
+++ b/tests/core/matrix.rs
@@ -498,6 +498,17 @@ fn apply() {
     assert_eq!(a, expected);
 }
 
+#[test]
+fn apply_with_location() {
+    let mut a = Matrix3::new(1.1f32, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9);
+
+    let expected = Matrix3::new(1.1, 3.2, 5.3, 5.4, 7.5, 9.6, 9.7, 11.8, 13.9);
+
+    a.apply_with_location(|i, j, e| *e += i as f32 + j as f32);
+
+    assert_eq!(a, expected);
+}
+
 #[test]
 fn map() {
     let a = Matrix4::new(