Add map_with_index! to Pointer, Array and StaticArray (#4456)

crystal-lang · May 26, 2017 · a6a4f22 · a6a4f22
1 parent d65ea1a
commit a6a4f22
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 0 deletions.
diff --git a/spec/std/array_spec.cr b/spec/std/array_spec.cr
@@ -1319,6 +1319,15 @@ describe "Array" do
     ary2.should eq([1, 2, 4, 5])
   end
 
+  it "does map_with_index!" do
+    ary = [0, 1, 2]
+    ary2 = ary.map_with_index! { |e, i| i * 2 }
+    ary[0].should eq(0)
+    ary[1].should eq(2)
+    ary[2].should eq(4)
+    ary2.should be(ary)
+  end
+
   it "does + with different types (#568)" do
     a = [1, 2, 3]
     a += ["hello"]

diff --git a/spec/std/pointer_spec.cr b/spec/std/pointer_spec.cr
@@ -196,6 +196,14 @@ describe "Pointer" do
     a[2].should eq(4)
   end
 
+  it "maps_with_index!" do
+    a = Pointer(Int32).malloc(3) { |i| i + 1 }
+    a.map_with_index!(3) { |e, i| e + i }
+    a[0].should eq(1)
+    a[1].should eq(3)
+    a[2].should eq(5)
+  end
+
   it "raises if mallocs negative size" do
     expect_raises(ArgumentError) { Pointer.malloc(-1, 0) }
   end

diff --git a/spec/std/static_array_spec.cr b/spec/std/static_array_spec.cr
@@ -117,6 +117,15 @@ describe "StaticArray" do
     a[2].should eq(4)
   end
 
+  it "map_with_index!" do
+    a = StaticArray(Int32, 3).new { |i| i + 1 }
+    a.map_with_index! { |e, i| i * 2 }
+    a[0].should eq(0)
+    a[1].should eq(2)
+    a[2].should eq(4)
+    a.should be_a(StaticArray(Int32, 3))
+  end
+
   it "updates value" do
     a = StaticArray(Int32, 3).new { |i| i + 1 }
     a.update(1) { |x| x * 2 }

diff --git a/src/array.cr b/src/array.cr
@@ -936,6 +936,12 @@ class Array(T)
     Array(U).new(size) { |i| yield @buffer[i], i }
   end
 
+  # Like `map_with_index`, but mutates `self` instead of allocating a new object.
+  def map_with_index!(&block : (T, Int32) -> T)
+    to_unsafe.map_with_index!(size) { |e, i| yield e, i }
+    self
+  end
+
   # Returns an `Array` with all possible permutations of *size*.
   #
   # ```

diff --git a/src/pointer.cr b/src/pointer.cr
@@ -377,6 +377,14 @@ struct Pointer(T)
     end
   end
 
+  # Like `map!`, but yield 2 arugments, the element and it's index
+  def map_with_index!(count : Int, &block)
+    count.times do |i|
+      self[i] = yield self[i], i
+    end
+    self
+  end
+
   # Returns a pointer whose memory address is zero. This doesn't allocate memory.
   #
   # When calling a C function you can also pass `nil` instead of constructing a

diff --git a/src/static_array.cr b/src/static_array.cr
@@ -165,6 +165,12 @@ struct StaticArray(T, N)
     self
   end
 
+  # Like `map`, but the block gets passed both the element and its index.
+  def map_with_index!(&block : (T, Int32) -> T)
+    to_unsafe.map_with_index!(size) { |e, i| yield e, i }
+    self
+  end
+
   # Reverses the elements of this array in-place, then returns `self`.
   #
   # ```