Skip to content

Commit 16832c1

Browse files
adienesKristofferC
authored andcommitted
Switch from segfault to zip behavior for mismatched indices in map! (#56673)
(cherry picked from commit 0947114)
1 parent fcf0174 commit 16832c1

File tree

2 files changed

+33
-6
lines changed

2 files changed

+33
-6
lines changed

base/abstractarray.jl

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3420,12 +3420,19 @@ function ith_all(i, as)
34203420
end
34213421

34223422
function map_n!(f::F, dest::AbstractArray, As) where F
3423-
idxs1 = LinearIndices(As[1])
3424-
@boundscheck LinearIndices(dest) == idxs1 && all(x -> LinearIndices(x) == idxs1, As)
3425-
for i = idxs1
3426-
@inbounds I = ith_all(i, As)
3427-
val = f(I...)
3428-
@inbounds dest[i] = val
3423+
idxs = LinearIndices(dest)
3424+
if all(x -> LinearIndices(x) == idxs, As)
3425+
for i in idxs
3426+
@inbounds as = ith_all(i, As)
3427+
val = f(as...)
3428+
@inbounds dest[i] = val
3429+
end
3430+
else
3431+
for (i, Is...) in zip(eachindex(dest), map(eachindex, As)...)
3432+
as = ntuple(j->getindex(As[j], Is[j]), length(As))
3433+
val = f(as...)
3434+
dest[i] = val
3435+
end
34293436
end
34303437
return dest
34313438
end

test/abstractarray.jl

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -889,6 +889,26 @@ include("generic_map_tests.jl")
889889
generic_map_tests(map, map!)
890890
@test_throws ArgumentError map!(-, [1])
891891

892+
@testset "#30624" begin
893+
### unstructured
894+
@test map!(+, ones(3), ones(3), ones(3), [1]) == [3, 1, 1]
895+
@test map!(+, ones(3), [1], ones(3), ones(3)) == [3, 1, 1]
896+
@test map!(+, [1], [1], [], []) == [1]
897+
@test map!(+, [[1]], [1], [], []) == [[1]]
898+
899+
# TODO: decide if input axes & lengths should be validated
900+
# @test_throws BoundsError map!(+, ones(1), ones(2))
901+
# @test_throws BoundsError map!(+, ones(1), ones(2, 2))
902+
903+
@test map!(+, ones(3), view(ones(2, 3), 1:2, 2:3), ones(3)) == [2, 2, 2]
904+
@test map!(+, ones(3), ones(2, 2), ones(3)) == [2, 2, 2]
905+
906+
### structured (all mapped arguments are <:AbstractArray equal ndims > 1)
907+
@test map!(+, ones(4), ones(2, 2), ones(2, 2)) == [2, 2, 2, 2]
908+
@test map!(+, ones(4), ones(2, 2), ones(1, 2)) == [2, 2, 1, 1]
909+
# @test_throws BoundsError map!(+, ones(3), ones(2, 2), ones(2, 2))
910+
end
911+
892912
test_UInt_indexing(TestAbstractArray)
893913
test_13315(TestAbstractArray)
894914
test_checksquare()

0 commit comments

Comments
 (0)