Skip to content

Commit 602d6ee

Browse files
serenity4serenity4
authored
InteractiveUtils: Fully support broadcasting expressions for code introspection macros (#58349)
This PR includes a fix and a feature for code introspection macros (`@code_typed`, `@code_llvm` and friends *but not* `@which`, `@edit`, etc): - Fixes a bug for expressions of the form `f.(x; y = 3)`, for which keyword arguments were not properly handled and led to an internal error. - Adds support for broadcasting assignments of the form `x .= f(y)`, `x .<<= f.(y, z)`, etc. The way this was (and still is) implemented is by constructing a temporary function, `f(x1, x2, x3, ...) = <body>` and feeding that to code introspection functions. This trick doesn't apply to `@which` and `@edit`, which need to map to a single function call (we could arguably choose to target `materialize`/`materialize!`, but this behavior could be a bit surprising and difficult to support). The switch differentiating the families of macro `@code_typed`/`@code_llvm` and `@which`/`@edit` etc was further exposed as an additional argument to `gen_call_with_extracted_types` and `gen_call_with_extracted_types_and_kwargs`, which default to the previous behavior (differentiating them based on whether their name starts with `code_`). The intent is to allow other macros such as `Cthulhu.@descend` to register themselves as code introspection macros. Quick tests indicate that it works as intended, e.g. with this PR Cthulhu supports `@descend [1, 2] .+= [2, 3]` (or equivalently, as added in #57909, `@descend ::Vector{Int} .+= ::Vector{Int}`). I originally just went for the fix, and after some refactoring I realized the feature was very straightforward to implement.
1 parent b04b104 commit 602d6ee

File tree

5 files changed

+137
-67
lines changed

5 files changed

+137
-67
lines changed

‎Compiler/test/irutils.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ macro code_typed1(ex0...)
1414
end
1515
get_code(args...; kwargs...) = code_typed1(args...; kwargs...).code
1616
macro get_code(ex0...)
17-
return gen_call_with_extracted_types_and_kwargs(__module__, :get_code, ex0)
17+
return gen_call_with_extracted_types_and_kwargs(__module__, :get_code, ex0; is_source_reflection = false)
1818
end
1919

2020
# check if `x` is a statement with a given `head`
@@ -58,7 +58,7 @@ function fully_eliminated(code::Vector{Any}; retval=(@__FILE__), kwargs...)
5858
return retval′ == retval
5959
end
6060
macro fully_eliminated(ex0...)
61-
return gen_call_with_extracted_types_and_kwargs(__module__, :fully_eliminated, ex0)
61+
return gen_call_with_extracted_types_and_kwargs(__module__, :fully_eliminated, ex0; is_source_reflection = false)
6262
end
6363

6464
let m = Meta.@lower 1 + 1

‎NEWS.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ Standard library changes
7272
#### InteractiveUtils
7373

7474
* Introspection utilities such as `@code_typed`, `@which` and `@edit` now accept type annotations as substitutes for values, recognizing forms such as `f(1, ::Float64, 3)` or even `sum(::Vector{T}; init = ::T) where {T<:Real}`. Type-annotated variables as in `f(val::Int; kw::Float64)` are not evaluated if the type annotation provides the necessary information, making this syntax compatible with signatures found in stacktraces ([#57909], [#58222]).
75+
* Code introspection macros such as `@code_lowered` and `@code_typed` now have a much better support for broadcasting expressions, including broadcasting assignments of the form `x .+= f(y)` ([#58349]).
7576

7677
External dependencies
7778
---------------------

‎stdlib/InteractiveUtils/src/macros.jl

Lines changed: 110 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -32,30 +32,65 @@ function get_typeof(@nospecialize ex)
3232
return :(Core.Typeof($(esc(ex))))
3333
end
3434

35+
function is_broadcasting_call(ex)
36+
isa(ex, Expr) || return false
37+
# Standard broadcasting: f.(x)
38+
isexpr(ex, :.) && length(ex.args) ≥ 2 && isexpr(ex.args[2], :tuple) && return true
39+
# Infix broadcasting: x .+ y, x .<< y, etc.
40+
if isexpr(ex, :call)
41+
f = ex.args[1]
42+
f == :.. && return false
43+
string(f)[1] == '.' && return true
44+
end
45+
return false
46+
end
47+
is_broadcasting_expr(ex) = is_broadcasting_call(ex) || is_broadcasting_assignment(ex)
48+
function is_broadcasting_assignment(ex)
49+
isa(ex, Expr) || return false
50+
isexpr(ex, :.) && return false
51+
head = string(ex.head)
52+
# x .= y, x .+= y, x .<<= y, etc.
53+
head[begin] == '.' && head[end] == '=' && return true
54+
return false
55+
end
56+
3557
"""
3658
Transform a dot expression into one where each argument has been replaced by a
3759
variable "xj" (with j an integer from 1 to the returned i).
3860
The list `args` contains the original arguments that have been replaced.
3961
"""
4062
function recursive_dotcalls!(ex, args, i=1)
41-
if !(ex isa Expr) || ((ex.head !== :. || !(ex.args[2] isa Expr)) &&
42-
(ex.head !== :call || string(ex.args[1])[1] != '.'))
43-
newarg = Symbol('x', i)
44-
if isexpr(ex, :...)
45-
push!(args, only(ex.args))
46-
return Expr(:..., newarg), i+1
63+
if is_broadcasting_expr(ex)
64+
if is_broadcasting_assignment(ex)
65+
(start, branches) = (1, ex.args)
66+
elseif isexpr(ex, :.)
67+
(start, branches) = (1, ex.args[2].args)
4768
else
48-
push!(args, ex)
49-
return newarg, i+1
69+
(start, branches) = (2, ex.args)
5070
end
71+
for j in start:length(branches)::Int
72+
branch, i = recursive_dotcalls!(branches[j], args, i)
73+
branches[j] = branch
74+
end
75+
return ex, i
76+
elseif isexpr(ex, :parameters)
77+
for j in eachindex(ex.args)
78+
param, i = recursive_dotcalls!(ex.args[j], args, i)
79+
ex.args[j] = param
80+
end
81+
return ex, i
5182
end
52-
(start, branches) = ex.head === :. ? (1, ex.args[2].args) : (2, ex.args)
53-
length_branches = length(branches)::Int
54-
for j in start:length_branches
55-
branch, i = recursive_dotcalls!(branches[j], args, i)
56-
branches[j] = branch
83+
newarg = Symbol('x', i)
84+
if isexpr(ex, :...)
85+
newarg = Expr(:..., newarg)
86+
push!(args, only(ex.args))
87+
elseif isexpr(ex, :kw)
88+
newarg = Expr(:kw, ex.args[1], newarg)
89+
push!(args, ex.args[end])
90+
else
91+
push!(args, ex)
5792
end
58-
return ex, i
93+
return newarg, i+1
5994
end
6095

6196
function extract_farg(@nospecialize arg)
@@ -158,20 +193,23 @@ function merge_namedtuple_types(nt::Type{<:NamedTuple}, nts::Type{<:NamedTuple}.
158193
NamedTuple{Tuple(names), Tuple{types...}}
159194
end
160195

161-
function gen_call_with_extracted_types(__module__, fcn, ex0, kws=Expr[])
196+
is_code_macro(fcn) = startswith(string(fcn), "code_")
197+
198+
function gen_call_with_extracted_types(__module__, fcn, ex0, kws = Expr[]; is_source_reflection = !is_code_macro(fcn), supports_binding_reflection = false)
162199
if isexpr(ex0, :ref)
163200
ex0 = replace_ref_begin_end!(ex0)
164201
end
165202
# assignments get bypassed: @edit a = f(x) <=> @edit f(x)
166203
if isa(ex0, Expr) && ex0.head == :(=) && isa(ex0.args[1], Symbol) && isempty(kws)
167-
return gen_call_with_extracted_types(__module__, fcn, ex0.args[2])
204+
return gen_call_with_extracted_types(__module__, fcn, ex0.args[2], kws; is_source_reflection, supports_binding_reflection)
168205
end
169206
where_params = nothing
170207
if isa(ex0, Expr)
171208
ex0, where_params = extract_where_parameters(ex0)
172209
end
173210
if isa(ex0, Expr)
174211
if ex0.head === :do && isexpr(get(ex0.args, 1, nothing), :call)
212+
# Normalize `f(args...) do ... end` calls to `f(do_anonymous_function, args...)`
175213
if length(ex0.args) != 2
176214
return Expr(:call, :error, "ill-formed do call")
177215
end
@@ -180,53 +218,60 @@ function gen_call_with_extracted_types(__module__, fcn, ex0, kws=Expr[])
180218
insert!(args, (isnothing(i) ? 2 : 1+i::Int), ex0.args[2])
181219
ex0 = Expr(:call, args...)
182220
end
183-
if ex0.head === :. || (ex0.head === :call && ex0.args[1] !== :.. && string(ex0.args[1])[1] == '.')
184-
codemacro = startswith(string(fcn), "code_")
185-
if codemacro && (ex0.head === :call || ex0.args[2] isa Expr)
186-
# Manually wrap a dot call in a function
187-
args = Any[]
188-
ex, i = recursive_dotcalls!(copy(ex0), args)
189-
xargs = [Symbol('x', j) for j in 1:i-1]
190-
dotfuncname = gensym("dotfunction")
191-
dotfuncdef = Expr(:local, Expr(:(=), Expr(:call, dotfuncname, xargs...), ex))
192-
return quote
193-
$(esc(dotfuncdef))
194-
$(fcn)($(esc(dotfuncname)), $(typesof_expr(args, where_params)); $(kws...))
195-
end
196-
elseif !codemacro
197-
fully_qualified_symbol = true # of the form A.B.C.D
198-
ex1 = ex0
199-
while ex1 isa Expr && ex1.head === :.
200-
fully_qualified_symbol = (length(ex1.args) == 2 &&
201-
ex1.args[2] isa QuoteNode &&
202-
ex1.args[2].value isa Symbol)
203-
fully_qualified_symbol || break
204-
ex1 = ex1.args[1]
221+
if is_broadcasting_expr(ex0) && !is_source_reflection
222+
# Manually wrap top-level broadcasts in a function.
223+
# We don't do that if `fcn` reflects into the source,
224+
# because that destroys provenance information.
225+
args = Any[]
226+
ex, i = recursive_dotcalls!(copy(ex0), args)
227+
xargs = [Symbol('x', j) for j in 1:i-1]
228+
dotfuncname = gensym("dotfunction")
229+
dotfuncdef = :(local $dotfuncname($(xargs...)) = $ex)
230+
return quote
231+
$(esc(dotfuncdef))
232+
$(gen_call_with_extracted_types(__module__, fcn, :($dotfuncname($(args...))), kws; is_source_reflection, supports_binding_reflection))
233+
end
234+
elseif isexpr(ex0, :.) && is_source_reflection
235+
# If `ex0` has the form A.B (or some chain A.B.C.D) and `fcn` reflects into the source,
236+
# `A` (or `A.B.C`) may be a module, in which case `fcn` is probably more interested in
237+
# the binding rather than the `getproperty` call.
238+
# If binding reflection is not supported, we generate an error; `getproperty(::Module, field)`
239+
# is not going to be interesting to reflect into, so best to allow future non-breaking support
240+
# for binding reflection in case the macro may eventually support that.
241+
fully_qualified_symbol = true
242+
ex1 = ex0
243+
while ex1 isa Expr && ex1.head === :.
244+
fully_qualified_symbol = (length(ex1.args) == 2 &&
245+
ex1.args[2] isa QuoteNode &&
246+
ex1.args[2].value isa Symbol)
247+
fully_qualified_symbol || break
248+
ex1 = ex1.args[1]
249+
end
250+
fully_qualified_symbol &= ex1 isa Symbol
251+
if fully_qualified_symbol || isexpr(ex1, :(::), 1)
252+
call_reflection = :($(fcn)(Base.getproperty, $(typesof_expr(ex0.args, where_params))))
253+
isexpr(ex0.args[1], :(::), 1) && return call_reflection
254+
if supports_binding_reflection
255+
binding_reflection = :($fcn(arg1, $(ex0.args[2])))
256+
else
257+
binding_reflection = :(error("expression is not a function call"))
205258
end
206-
fully_qualified_symbol &= ex1 isa Symbol
207-
if fully_qualified_symbol || isexpr(ex1, :(::), 1)
208-
getproperty_ex = :($(fcn)(Base.getproperty, $(typesof_expr(ex0.args, where_params))))
209-
isexpr(ex0.args[1], :(::), 1) && return getproperty_ex
210-
return quote
211-
local arg1 = $(esc(ex0.args[1]))
212-
if isa(arg1, Module)
213-
$(if string(fcn) == "which"
214-
:(which(arg1, $(ex0.args[2])))
215-
else
216-
:(error("expression is not a function call"))
217-
end)
218-
else
219-
$getproperty_ex
220-
end
259+
return quote
260+
local arg1 = $(esc(ex0.args[1]))
261+
if isa(arg1, Module)
262+
$binding_reflection
263+
else
264+
$call_reflection
221265
end
222-
else
223-
return Expr(:call, :error, "dot expressions are not lowered to "
224-
* "a single function call, so @$fcn cannot analyze "
225-
* "them. You may want to use Meta.@lower to identify "
226-
* "which function call to target.")
227266
end
228267
end
229268
end
269+
if is_broadcasting_expr(ex0)
270+
return Expr(:call, :error, "dot expressions are not lowered to "
271+
* "a single function call, so @$fcn cannot analyze "
272+
* "them. You may want to use Meta.@lower to identify "
273+
* "which function call to target.")
274+
end
230275
if any(@nospecialize(a)->(isexpr(a, :kw) || isexpr(a, :parameters)), ex0.args)
231276
args, kwargs = separate_kwargs(ex0.args)
232277
are_kwargs_valid(kwargs) || return quote
@@ -318,7 +363,7 @@ Same behaviour as `gen_call_with_extracted_types` except that keyword arguments
318363
of the form "foo=bar" are passed on to the called function as well.
319364
The keyword arguments must be given before the mandatory argument.
320365
"""
321-
function gen_call_with_extracted_types_and_kwargs(__module__, fcn, ex0)
366+
function gen_call_with_extracted_types_and_kwargs(__module__, fcn, ex0; is_source_reflection = !is_code_macro(fcn), supports_binding_reflection = false)
322367
kws = Expr[]
323368
arg = ex0[end] # Mandatory argument
324369
for i in 1:length(ex0)-1
@@ -332,13 +377,15 @@ function gen_call_with_extracted_types_and_kwargs(__module__, fcn, ex0)
332377
return Expr(:call, :error, "@$fcn expects only one non-keyword argument")
333378
end
334379
end
335-
return gen_call_with_extracted_types(__module__, fcn, arg, kws)
380+
return gen_call_with_extracted_types(__module__, fcn, arg, kws; is_source_reflection, supports_binding_reflection)
336381
end
337382

338383
for fname in [:which, :less, :edit, :functionloc]
339384
@eval begin
340385
macro ($fname)(ex0)
341-
gen_call_with_extracted_types(__module__, $(Expr(:quote, fname)), ex0)
386+
gen_call_with_extracted_types(__module__, $(Expr(:quote, fname)), ex0, Expr[];
387+
is_source_reflection = true,
388+
supports_binding_reflection = $(fname === :which))
342389
end
343390
end
344391
end
@@ -351,13 +398,13 @@ end
351398
for fname in [:code_warntype, :code_llvm, :code_native,
352399
:infer_return_type, :infer_effects, :infer_exception_type]
353400
@eval macro ($fname)(ex0...)
354-
gen_call_with_extracted_types_and_kwargs(__module__, $(QuoteNode(fname)), ex0)
401+
gen_call_with_extracted_types_and_kwargs(__module__, $(QuoteNode(fname)), ex0; is_source_reflection = false)
355402
end
356403
end
357404

358405
for fname in [:code_typed, :code_lowered, :code_ircode]
359406
@eval macro ($fname)(ex0...)
360-
thecall = gen_call_with_extracted_types_and_kwargs(__module__, $(QuoteNode(fname)), ex0)
407+
thecall = gen_call_with_extracted_types_and_kwargs(__module__, $(QuoteNode(fname)), ex0; is_source_reflection = false)
361408
quote
362409
local results = $thecall
363410
length(results) == 1 ? results[1] : results

‎stdlib/InteractiveUtils/test/runtests.jl

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,13 @@ end
373373
@test (@which round(1.2; digits = ::Float64, kwargs_1...)).name === :round
374374
@test (@which round(1.2; sigdigits = ::Int, kwargs_1...)).name === :round
375375
@test (@which round(1.2; kwargs_1..., kwargs_2..., base)).name === :round
376+
@test (@code_typed optimize=false round.([1.0, 2.0]; digits = ::Int64))[2] == Vector{Float64}
377+
@test (@code_typed optimize=false round.(::Vector{Float64}, base = 2; digits = ::Int64))[2] == Vector{Float64}
378+
@test (@code_typed optimize=false round.(base = ::Int64, ::Vector{Float64}; digits = ::Int64))[2] == Vector{Float64}
379+
@test (@code_typed optimize=false [1, 2] .= ::Int)[2] == Vector{Int}
380+
@test (@code_typed optimize=false ::Vector{Int} .= ::Int)[2] == Vector{Int}
381+
@test (@code_typed optimize=false ::Vector{Float64} .= 1 .+ ::Vector{Int})[2] == Vector{Float64}
382+
@test (@code_typed optimize=false ::Vector{Float64} .= 1 .+ round.(base = ::Int, ::Vector{Int}; digits = 3))[2] == Vector{Float64}
376383
end
377384

378385
module MacroTest
@@ -505,7 +512,22 @@ a14637 = A14637(0)
505512
@test (@code_typed optimize=true max.([1,7], UInt.([4])))[2] == Vector{UInt}
506513
@test (@code_typed Ref.([1,2])[1].x)[2] == Int
507514
@test (@code_typed max.(Ref(true).x))[2] == Bool
515+
@test (@code_typed optimize=false round.([1.0, 2.0]; digits = 3))[2] == Vector{Float64}
516+
@test (@code_typed optimize=false round.([1.0, 2.0], base = 2; digits = 3))[2] == Vector{Float64}
517+
@test (@code_typed optimize=false round.(base = 2, [1.0, 2.0], digits = 3))[2] == Vector{Float64}
518+
@test (@code_typed optimize=false [1, 2] .= 2)[2] == Vector{Int}
519+
@test (@code_typed optimize=false [1, 2] .<<= 2)[2] == Vector{Int}
520+
@test (@code_typed optimize=false [1, 2.0] .= 1 .+ [2, 3])[2] == Vector{Float64}
521+
@test (@code_typed optimize=false [1, 2.0] .= 1 .+ round.(base = 1, [1, 3]; digits = 3))[2] == Vector{Float64}
522+
@test (@code_typed optimize=false [1] .+ [2])[2] == Vector{Int}
508523
@test !isempty(@code_typed optimize=false max.(Ref.([5, 6])...))
524+
expansion = string(@macroexpand @code_typed optimize=false max.(Ref.([5, 6])...))
525+
@test contains(expansion, "(x1) =") # presence of wrapper function
526+
# Make sure broadcasts in nested arguments are not processed.
527+
v = Any[1]
528+
expansion = string(@macroexpand @code_typed v[1] = rand.(Ref(1)))
529+
@test contains(expansion, "Typeof(rand.(Ref(1)))")
530+
@test !contains(expansion, "(x1) =")
509531

510532
# Issue # 45889
511533
@test !isempty(@code_typed 3 .+ 6)
@@ -619,7 +641,7 @@ end
619641
@test_throws err @code_lowered 1
620642
@test_throws err @code_lowered 1.0
621643

622-
@test_throws "is too complex" @code_lowered a .= 1 + 2
644+
@test_throws "dot expressions are not lowered to a single function call" @which a .= 1 + 2
623645
@test_throws "invalid keyword argument syntax" @eval @which round(1; digits(3))
624646
end
625647

‎stdlib/Test/src/Test.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2134,7 +2134,7 @@ function _inferred(ex, mod, allow = :(Union{}))
21342134
kwargs = gensym()
21352135
quote
21362136
$(esc(args)), $(esc(kwargs)), result = $(esc(Expr(:call, _args_and_call, ex.args[2:end]..., ex.args[1])))
2137-
inftype = $(gen_call_with_extracted_types(mod, Base.infer_return_type, :($(ex.args[1])($(args)...; $(kwargs)...))))
2137+
inftype = $(gen_call_with_extracted_types(mod, Base.infer_return_type, :($(ex.args[1])($(args)...; $(kwargs)...)); is_source_reflection = false))
21382138
end
21392139
else
21402140
# No keywords

0 commit comments

Comments
 (0)