Skip to content

Commit

Permalink
Force specialization on number of args passed to anonymous function c…
Browse files Browse the repository at this point in the history
…onstructed by `oror` and `andand` (JuliaLang#54152)

These functions should only accept one argument signature so we should fully specialize.
  • Loading branch information
LilithHafner authored Apr 19, 2024
1 parent 2f90962 commit 3364aa5
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -196,14 +196,18 @@ const andand = AndAnd()
broadcasted(::AndAnd, a, b) = broadcasted((a, b) -> a && b, a, b)
function broadcasted(::AndAnd, a, bc::Broadcasted)
bcf = flatten(bc)
broadcasted((a, args...) -> a && bcf.f(args...), a, bcf.args...)
# Vararg type signature to specialize on args count. This is necessary for performance
# and innexpensive because this should only ever get called with 1+N = length(bc.args)
broadcasted(((a, args::Vararg{Any, N}) where {N}) -> a && bcf.f(args...), a, bcf.args...)
end
struct OrOr end
const oror = OrOr()
broadcasted(::OrOr, a, b) = broadcasted((a, b) -> a || b, a, b)
function broadcasted(::OrOr, a, bc::Broadcasted)
bcf = flatten(bc)
broadcasted((a, args...) -> a || bcf.f(args...), a, bcf.args...)
# Vararg type signature to specialize on args count. This is necessary for performance
# and innexpensive because this should only ever get called with 1+N = length(bc.args)
broadcasted(((a, args::Vararg{Any, N}) where {N}) -> a || bcf.f(args...), a, bcf.args...)
end

Base.convert(::Type{Broadcasted{NewStyle}}, bc::Broadcasted{<:Any,Axes,F,Args}) where {NewStyle,Axes,F,Args} =
Expand Down Expand Up @@ -345,6 +349,7 @@ function flatten(bc::Broadcasted)
# makeargs[3] = ((w, x, y, z)) -> z
makeargs = make_makeargs(bc.args)
f = Base.maybeconstructor(bc.f)
# TODO: consider specializing on args... if performance problems emerge:
newf = (args...) -> (@inline; f(prepare_args(makeargs, args)...))
return Broadcasted(bc.style, newf, args, bc.axes)
end
Expand Down

0 comments on commit 3364aa5

Please sign in to comment.