Only reuse memory in mapslices when it's safe (fixes #18524) (#18570)

Author: timholy

base/abstractarray.jl

@@ -1734,6 +1734,7 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
 
     Aslice = A[idx...]
     r1 = f(Aslice)
+    safe_for_reuse = isa(r1, Number) || (isa(r1, AbstractArray) && eltype(r1) <: Number)
 
     # determine result size and allocate
     Rsize = copy(dimsA)
@@ -1758,15 +1759,31 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
 
     isfirst = true
     nidx = length(otherdims)
-    for I in CartesianRange(itershape)
-        if isfirst
-            isfirst = false
-        else
-            for i in 1:nidx
-                idx[otherdims[i]] = ridx[otherdims[i]] = I.I[i]
+    if safe_for_reuse
+        # when f returns an array, R[ridx...] = f(Aslice) line copies elements,
+        # so we can reuse Aslice
+        for I in CartesianRange(itershape)
+            if isfirst
+                isfirst = false  # skip the first element, we already handled it
+            else
+                for i in 1:nidx
+                    idx[otherdims[i]] = ridx[otherdims[i]] = I.I[i]
+                end
+                _unsafe_getindex!(Aslice, A, idx...)
+                R[ridx...] = f(Aslice)
+            end
+        end
+    else
+        # we can't guarantee safety (#18524), so allocate new storage for each slice
+        for I in CartesianRange(itershape)
+            if isfirst
+                isfirst = false
+            else
+                for i in 1:nidx
+                    idx[otherdims[i]] = ridx[otherdims[i]] = I.I[i]
+                end
+                R[ridx...] = f(A[idx...])
             end
-            _unsafe_getindex!(Aslice, A, idx...)
-            R[ridx...] = f(Aslice)
         end
     end
 

test/arrayops.jl

@@ -953,6 +953,17 @@ end
         @test lexless(asc[:,1],asc[:,2])
         @test lexless(asc[:,2],asc[:,3])
 
+        # mutating functions
+        o = ones(3, 4)
+        m = mapslices(x->fill!(x, 0), o, 2)
+        @test m == zeros(3, 4)
+        @test o == ones(3, 4)
+
+        # issue #18524
+        m = mapslices(x->tuple(x), [1 2; 3 4], 1)
+        @test m[1,1] == ([1,3],)
+        @test m[1,2] == ([2,4],)
+
         asr = sortrows(a, rev=true)
         @test lexless(asr[2,:],asr[1,:])
         @test lexless(asr[3,:],asr[2,:])