Skip to content

Commit 70e4a6a

Browse files
yuyichaoyuyichao
committed
Fix at-views for setindex expression
Fix #59383
1 parent b24014a commit 70e4a6a

File tree

2 files changed

+140
-71
lines changed

2 files changed

+140
-71
lines changed

‎base/views.jl

Lines changed: 88 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,81 @@ function replace_ref_begin_end_!(__module__::Module, ex, withex, in_quote_contex
2525
end
2626
return ex
2727
end
28+
function handle_refexpr!(__module__::Module, ref_ex::Expr, main_ex::Expr, withex, in_quote_context, escs::Int)
29+
@assert !in_quote_context
30+
local used_withex
31+
ref_ex.args[1], used_withex = replace_ref_begin_end_!(__module__, ref_ex.args[1], withex, in_quote_context, escs)
32+
S = gensym(:S) # temp var to cache ex.args[1] if needed. if S is a global or expression, then it has side effects to use
33+
assignments = []
34+
used_S = false # whether we actually need S
35+
# new :ref, so redefine withex
36+
nargs = length(ref_ex.args) - 1
37+
if nargs == 0
38+
return main_ex, used_withex
39+
elseif nargs == 1
40+
# replace with lastindex(S)
41+
ref_ex.args[2], used_S = replace_ref_begin_end_!(__module__, ref_ex.args[2], (:($firstindex($S)),:($lastindex($S))), in_quote_context, escs)
42+
else
43+
ni = 1
44+
nx = 0
45+
J = nargs + 1
46+
need_temps = false # whether any arg needs temporaries
47+
48+
# First pass: determine if any argument will needs temporaries
49+
for j = 2:J
50+
exj = ref_ex.args[j]
51+
if isexpr(exj, :...)
52+
need_temps = true
53+
break
54+
end
55+
end
56+
57+
# Second pass: if any need temps, create temps for all args
58+
temp_vars = Tuple{Int,Symbol}[]
59+
for j = 2:J
60+
n = nx === 0 ? ni : :($nx + $ni)
61+
exj, used = replace_ref_begin_end_!(__module__, ref_ex.args[j], (:($firstindex($S,$n)),:($lastindex($S,$n))), in_quote_context, escs)
62+
used_S |= used
63+
ref_ex.args[j] = exj
64+
ni += 1
65+
if need_temps
66+
isva = isexpr(exj, :...) # implied need_temps
67+
if isva
68+
exj = exj.args[1]
69+
end
70+
if isa_ast_node(exj) # create temp to preserve evaluation order and count in case `used` gets set later
71+
exj = gensym(:arg)
72+
push!(temp_vars, (j, exj))
73+
end
74+
if isva
75+
ni -= 1
76+
nx = nx === 0 ? :(length($exj)) : :($nx + length($exj))
77+
end
78+
end
79+
end
80+
81+
# Third pass: if `used`, need to actually make those temp assignments now
82+
if used_S
83+
for (j, temp_var) in temp_vars
84+
exj = ref_ex.args[j]
85+
isva = isexpr(exj, :...) # implied need_temps
86+
if isva
87+
exj = exj.args[1]
88+
end
89+
push!(assignments, :(local $temp_var = $exj))
90+
ref_ex.args[j] = isva ? Expr(:..., temp_var) : temp_var
91+
end
92+
end
93+
end
94+
95+
if used_S
96+
S0 = ref_ex.args[1]
97+
S = escapes(S, escs)
98+
ref_ex.args[1] = S
99+
main_ex = :(local $S = $S0; $(assignments...); $main_ex)
100+
end
101+
return main_ex, used_withex
102+
end
28103
if ex isa Expr && ex.head === :macrocall
29104
# Blithly modifying the arguments to another macro is unwise, so call
30105
# macroexpand first on it.
@@ -46,77 +121,7 @@ function replace_ref_begin_end_!(__module__::Module, ex, withex, in_quote_contex
46121
end
47122
elseif isa(ex,Expr)
48123
if !in_quote_context && ex.head === :ref # n.b. macroexpand.scm design is incapable of tracking :begin and :end scope, so emulate that here too and ignore escs
49-
ex.args[1], used_withex = replace_ref_begin_end_!(__module__, ex.args[1], withex, in_quote_context, escs)
50-
S = gensym(:S) # temp var to cache ex.args[1] if needed. if S is a global or expression, then it has side effects to use
51-
assignments = []
52-
used_S = false # whether we actually need S
53-
# new :ref, so redefine withex
54-
nargs = length(ex.args)-1
55-
if nargs == 0
56-
return ex, used_withex
57-
elseif nargs == 1
58-
# replace with lastindex(S)
59-
ex.args[2], used_S = replace_ref_begin_end_!(__module__, ex.args[2], (:($firstindex($S)),:($lastindex($S))), in_quote_context, escs)
60-
else
61-
ni = 1
62-
nx = 0
63-
J = lastindex(ex.args)
64-
need_temps = false # whether any arg needs temporaries
65-
66-
# First pass: determine if any argument will needs temporaries
67-
for j = 2:J
68-
exj = ex.args[j]
69-
if isexpr(exj, :...)
70-
need_temps = true
71-
break
72-
end
73-
end
74-
75-
# Second pass: if any need temps, create temps for all args
76-
temp_vars = Tuple{Int,Symbol}[]
77-
for j = 2:J
78-
n = nx === 0 ? ni : :($nx + $ni)
79-
exj, used = replace_ref_begin_end_!(__module__, ex.args[j], (:($firstindex($S,$n)),:($lastindex($S,$n))), in_quote_context, escs)
80-
used_S |= used
81-
ex.args[j] = exj
82-
ni += 1
83-
if need_temps
84-
isva = isexpr(exj, :...) # implied need_temps
85-
if isva
86-
exj = exj.args[1]
87-
end
88-
if isa_ast_node(exj) # create temp to preserve evaluation order and count in case `used` gets set later
89-
exj = gensym(:arg)
90-
push!(temp_vars, (j, exj))
91-
end
92-
if isva
93-
ni -= 1
94-
nx = nx === 0 ? :(length($exj)) : :($nx + length($exj))
95-
end
96-
end
97-
end
98-
99-
# Third pass: if `used`, need to actually make those temp assignments now
100-
if used_S
101-
for (j, temp_var) in temp_vars
102-
exj = ex.args[j]
103-
isva = isexpr(exj, :...) # implied need_temps
104-
if isva
105-
exj = exj.args[1]
106-
end
107-
push!(assignments, :(local $temp_var = $exj))
108-
ex.args[j] = isva ? Expr(:..., temp_var) : temp_var
109-
end
110-
end
111-
end
112-
113-
if used_S
114-
S0 = ex.args[1]
115-
S = escapes(S, escs)
116-
ex.args[1] = S
117-
ex = :(local $S = $S0; $(assignments...); $ex)
118-
end
119-
return ex, used_withex
124+
return handle_refexpr!(__module__, ex, ex, withex, in_quote_context, escs)
120125
elseif ex.head === :$
121126
# no longer an executable expression (handle all equivalent forms of :inert, :quote, and QuoteNode the same way)
122127
in_quote_context = false
@@ -132,6 +137,18 @@ function replace_ref_begin_end_!(__module__::Module, ex, withex, in_quote_contex
132137
escs -= 1
133138
elseif ex.head === :meta || ex.head === :inert
134139
return ex, used_withex
140+
elseif !in_quote_context && last(string(ex.head)) == '=' && Meta.isexpr(ex.args[1], :ref)
141+
for i = eachindex(ex.args)
142+
if i == 1
143+
# we'll deal with the ref expression later
144+
continue
145+
end
146+
ex.args[i], used = replace_ref_begin_end_!(__module__, ex.args[i], withex, in_quote_context, escs)
147+
used_withex |= used
148+
end
149+
ex, used = handle_refexpr!(__module__, ex.args[1]::Expr, ex, withex, in_quote_context, escs)
150+
used_withex |= used
151+
return ex, used_withex
135152
end
136153
# recursive search
137154
for i = eachindex(ex.args)

‎test/subarray.jl

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1101,3 +1101,55 @@ end
11011101

11021102

11031103
@test @views quote var"begin" + var"end" end isa Expr
1104+
1105+
@testset "@views handling of assignment" begin
1106+
@test @macroexpand(@views x[a:b] = c) == :(x[a:b] = c)
1107+
# Assignments should still work
1108+
let array = [1, 2, 3, 4, 5, 6, 7, 8]
1109+
@views array[begin:2] = [-1, -2]
1110+
@test array == [-1, -2, 3, 4, 5, 6, 7, 8]
1111+
@views array[7:end] = [-7, -8]
1112+
@test array == [-1, -2, 3, 4, 5, 6, -7, -8]
1113+
@views array[begin + 2:end - 4] = [-3, -4]
1114+
@test array == [-1, -2, -3, -4, 5, 6, -7, -8]
1115+
@views identity(array)[begin + 4:end - 2] = [-5, -6]
1116+
@test array == [-1, -2, -3, -4, -5, -6, -7, -8]
1117+
1118+
@views array[begin:2] .= 100
1119+
@test array == [100, 100, -3, -4, -5, -6, -7, -8]
1120+
@views array[7:end] .= 200
1121+
@test array == [100, 100, -3, -4, -5, -6, 200, 200]
1122+
@views array[begin + 2:end - 4] .= 300
1123+
@test array == [100, 100, 300, 300, -5, -6, 200, 200]
1124+
@views identity(array)[begin + 4:end - 2] .= 400
1125+
@test array == [100, 100, 300, 300, 400, 400, 200, 200]
1126+
1127+
@views identity(array)[begin:end] .-= 1
1128+
@test array == [99, 99, 299, 299, 399, 399, 199, 199]
1129+
1130+
@views identity(array)[begin:end] += [1, 2, 3, 4, 5, 6, 7, 8]
1131+
@test array == [100, 101, 302, 303, 404, 405, 206, 207]
1132+
end
1133+
# Nested getindex in assignment should be transformed
1134+
let array = [1, 2, 3, 4, 5, 6, 7, 8], array2 = [1, 2, 3, 4, 5, 6, 7, 8]
1135+
array[begin + 1:end - 2][2] = -1
1136+
array[begin + 1:end - 2][end] = -2
1137+
@test array == [1, 2, 3, 4, 5, 6, 7, 8]
1138+
1139+
@views array[begin + 1:end - 2][2] = -1
1140+
@views array[begin + 1:end - 2][end] = -2
1141+
@test array == [1, 2, -1, 4, 5, -2, 7, 8]
1142+
1143+
function swap_ele(ary, i, v)
1144+
res = ary[i]
1145+
ary[i] = v
1146+
return res
1147+
end
1148+
array2[swap_ele(array[begin:end], 1, -3):swap_ele(array[begin:end], 7, -4)] = [-1, -2, -3, -4, -5, -6, -7]
1149+
@test array == [1, 2, -1, 4, 5, -2, 7, 8]
1150+
@test array2 == [-1, -2, -3, -4, -5, -6, -7, 8]
1151+
@views array2[swap_ele(array[begin:end], 1, -3):swap_ele(array[begin:end], 7, -4)] = [-10, 2, -30, 4, -50, 6, -70]
1152+
@test array == [-3, 2, -1, 4, 5, -2, -4, 8]
1153+
@test array2 == [-10, 2, -30, 4, -50, 6, -70, 8]
1154+
end
1155+
end

0 commit comments

Comments
 (0)