From e4ddcba1189b67f4519814d55737668565b5b222 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9my=20Marronnier?= <rmarronnier@gmail.com>
Date: Sat, 21 Sep 2019 04:26:30 +0200
Subject: [PATCH] Adds #map_with_index to Matrix

---
 spec/apatite/matrix_spec.cr | 10 ++++++++++
 src/apatite/matrix.cr       | 14 +++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 spec/apatite/matrix_spec.cr

diff --git a/spec/apatite/matrix_spec.cr b/spec/apatite/matrix_spec.cr
new file mode 100644
index 0000000..3ed1bfb
--- /dev/null
+++ b/spec/apatite/matrix_spec.cr
@@ -0,0 +1,10 @@
+require "../spec_helper"
+
+describe "Apatite::Matrix" do
+  describe ".map_with_index" do
+    it "returns a new matrix with each element processed according to their indexes" do
+      matrix = Apatite::Matrix[[1, 2], [3, 4]]
+      matrix.map_with_index { |e, i, j| e*i + j }.should eq Apatite::Matrix[[0, 1], [3, 5]]
+    end
+  end
+end
diff --git a/src/apatite/matrix.cr b/src/apatite/matrix.cr
index a1368df..d3268e9 100644
--- a/src/apatite/matrix.cr
+++ b/src/apatite/matrix.cr
@@ -384,6 +384,18 @@ module Apatite
       Matrix.new(rows, column_count)
     end
 
+    # Same as #map, but yields the row index and column index in addition to the element
+    #
+    # ```
+    # matrix = Matrix[[1, 2], [3, 4]].map_with_index do |e, row, col|
+    #   e*row + col
+    # end
+    # matrix == Matrix[[0, 1], [3, 5]] # => true
+    # ```
+    def map_with_index(&block : T, Int32, Int32 -> U) forall U
+      Matrix(U).build(@rows.size, @column_count) { |i, j| yield @rows[i][j], i, j }
+    end
+
     # Yields all elements of the matrix, starting with those of the first row,
     # or returns an Enumerator if no block given.
     # Elements can be restricted by passing an argument:
@@ -446,7 +458,7 @@ module Apatite
       end
     end
 
-    # Same as #each, but the row index and column index in addition to the element
+    # Same as #each, but yields the row index and column index in addition to the element
     #
     # ```
     # Matrix[[1, 2], [3, 4]].each_with_index do |e, row, col|