Skip to content

Commit 368516b

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

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
@@ -3335,12 +3335,19 @@ function ith_all(i, as)
33353335
end
33363336

33373337
function map_n!(f::F, dest::AbstractArray, As) where F
3338-
idxs1 = LinearIndices(As[1])
3339-
@boundscheck LinearIndices(dest) == idxs1 && all(x -> LinearIndices(x) == idxs1, As)
3340-
for i = idxs1
3341-
@inbounds I = ith_all(i, As)
3342-
val = f(I...)
3343-
@inbounds dest[i] = val
3338+
idxs = LinearIndices(dest)
3339+
if all(x -> LinearIndices(x) == idxs, As)
3340+
for i in idxs
3341+
@inbounds as = ith_all(i, As)
3342+
val = f(as...)
3343+
@inbounds dest[i] = val
3344+
end
3345+
else
3346+
for (i, Is...) in zip(eachindex(dest), map(eachindex, As)...)
3347+
as = ntuple(j->getindex(As[j], Is[j]), length(As))
3348+
val = f(as...)
3349+
dest[i] = val
3350+
end
33443351
end
33453352
return dest
33463353
end

test/abstractarray.jl

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

838+
@testset "#30624" begin
839+
### unstructured
840+
@test map!(+, ones(3), ones(3), ones(3), [1]) == [3, 1, 1]
841+
@test map!(+, ones(3), [1], ones(3), ones(3)) == [3, 1, 1]
842+
@test map!(+, [1], [1], [], []) == [1]
843+
@test map!(+, [[1]], [1], [], []) == [[1]]
844+
845+
# TODO: decide if input axes & lengths should be validated
846+
# @test_throws BoundsError map!(+, ones(1), ones(2))
847+
# @test_throws BoundsError map!(+, ones(1), ones(2, 2))
848+
849+
@test map!(+, ones(3), view(ones(2, 3), 1:2, 2:3), ones(3)) == [2, 2, 2]
850+
@test map!(+, ones(3), ones(2, 2), ones(3)) == [2, 2, 2]
851+
852+
### structured (all mapped arguments are <:AbstractArray equal ndims > 1)
853+
@test map!(+, ones(4), ones(2, 2), ones(2, 2)) == [2, 2, 2, 2]
854+
@test map!(+, ones(4), ones(2, 2), ones(1, 2)) == [2, 2, 1, 1]
855+
# @test_throws BoundsError map!(+, ones(3), ones(2, 2), ones(2, 2))
856+
end
857+
838858
test_UInt_indexing(TestAbstractArray)
839859
test_13315(TestAbstractArray)
840860
test_checksquare()

0 commit comments

Comments
 (0)