Skip to content

Make map! correctly handle differently-sized containers. #50647

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 8 commits into from
Closed

Make map! correctly handle differently-sized containers. #50647

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b9b48cc
Make map! correctly handle differently-sized containers.
maleadt Jul 23, 2023
3e363ac
Handle OffsetArrays.
maleadt Jul 24, 2023
98d3bf1
Remove some dead code and unnecessary specializations.
maleadt Jul 24, 2023
5b4ac94
Fix destination boundscheck.
maleadt Jul 25, 2023
25dc0c1
Avoid bootstrap failure.
maleadt Jul 25, 2023
54a5d02
Make mapfirst use inbounds map to avoid regression OffsetArrays alloc…
maleadt Jul 25, 2023
84acf71
Optimizations.
maleadt Jul 26, 2023
e9fc1c2
Fix error type.
maleadt Jul 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 9 additions & 29 deletions base/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3261,14 +3261,6 @@ map(f) = f()

## 1 argument

function map!(f::F, dest::AbstractArray, A::AbstractArray) where F
for (i,j) in zip(eachindex(dest),eachindex(A))
val = f(@inbounds A[j])
@inbounds dest[i] = val
end
return dest
end

# map on collections
map(f, A::AbstractArray) = collect_similar(A, Generator(f,A))

Expand Down Expand Up @@ -3303,30 +3295,17 @@ map(f, A) = collect(Generator(f,A)) # default to returning an Array for `map` on
map(f, ::AbstractDict) = error("map is not defined on dictionaries")
map(f, ::AbstractSet) = error("map is not defined on sets")

## 2 argument
function map!(f::F, dest::AbstractArray, A::AbstractArray, B::AbstractArray) where F
for (i, j, k) in zip(eachindex(dest), eachindex(A), eachindex(B))
@inbounds a, b = A[j], B[k]
val = f(a, b)
@inbounds dest[i] = val
end
return dest
end

## N argument

@inline ith_all(i, ::Tuple{}) = ()
function ith_all(i, as)
@_propagate_inbounds_meta
return (as[1][i], ith_all(i, tail(as))...)
end
@noinline throw_map_mismatch(ndest, nvals) =
throw(DimensionMismatch("map! over $nvals values, but destination only has length $ndest"))

function map_n!(f::F, dest::AbstractArray, As) where F
idxs1 = LinearIndices(As[1])
@boundscheck LinearIndices(dest) == idxs1 && all(x -> LinearIndices(x) == idxs1, As)
for i = idxs1
@inbounds I = ith_all(i, As)
val = f(I...)
@inline function map_n!(f::F, dest::AbstractArray, As) where F
Is = zip(map(eachindex, As)...)
@boundscheck length(Is) <= length(dest) || throw_map_mismatch(length(dest), length(Is))
for (i, js...) in zip(eachindex(dest), map(eachindex, As)...)
J = ntuple(d -> @inbounds(As[d][js[d]]), Val(length(As)))
val = f(J...)
@inbounds dest[i] = val
end
return dest
Expand Down Expand Up @@ -3362,6 +3341,7 @@ julia> map!(+, zeros(Int, 5), 100:999, 1:3)
```
"""
function map!(f::F, dest::AbstractArray, As::AbstractArray...) where {F}
@_propagate_inbounds_meta
isempty(As) && throw(ArgumentError(
"""map! requires at least one "source" argument"""))
map_n!(f, dest, As)
Expand Down
2 changes: 1 addition & 1 deletion base/reducedim.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,7 @@ copyfirst!(R::AbstractArray, A::AbstractArray) = mapfirst!(identity, R, A)
function mapfirst!(f::F, R::AbstractArray, A::AbstractArray{<:Any,N}) where {N, F}
lsiz = check_reducedims(R, A)
t = _firstreducedslice(axes(R), axes(A))
map!(f, R, view(A, t...))
@inbounds map!(f, R, view(A, t...))
end
# We know that the axes of R and A are compatible, but R might have a different number of
# dimensions than A, which is trickier than it seems due to offset arrays and type stability
Expand Down
6 changes: 6 additions & 0 deletions test/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -844,6 +844,12 @@ test_ind2sub(TestAbstractArray)
include("generic_map_tests.jl")
generic_map_tests(map, map!)
@test_throws ArgumentError map!(-, [1])
# Issue #30624
@test map!(+, [0,0,0], [1,2], [10,20,30], [100]) == [111,0,0]
## destination container should be large enough
@test_throws DimensionMismatch map!(+, [0], [1,2])
@test_throws DimensionMismatch map!(+, [0], [1,2], [1,2])
@test_throws DimensionMismatch map!(+, [0], [1,2], [1,2], [1,2])

test_UInt_indexing(TestAbstractArray)
test_13315(TestAbstractArray)
Expand Down