diff --git a/base/views.jl b/base/views.jl index 10fb3b0dc5e1b..d205b1b38dee4 100644 --- a/base/views.jl +++ b/base/views.jl @@ -25,6 +25,81 @@ function replace_ref_begin_end_!(__module__::Module, ex, withex, in_quote_contex end return ex end + function handle_refexpr!(__module__::Module, ref_ex::Expr, main_ex::Expr, withex, in_quote_context, escs::Int) + @assert !in_quote_context + local used_withex + ref_ex.args[1], used_withex = replace_ref_begin_end_!(__module__, ref_ex.args[1], withex, in_quote_context, escs) + 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 + assignments = [] + used_S = false # whether we actually need S + # new :ref, so redefine withex + nargs = length(ref_ex.args) - 1 + if nargs == 0 + return main_ex, used_withex + elseif nargs == 1 + # replace with lastindex(S) + ref_ex.args[2], used_S = replace_ref_begin_end_!(__module__, ref_ex.args[2], (:($firstindex($S)),:($lastindex($S))), in_quote_context, escs) + else + ni = 1 + nx = 0 + J = nargs + 1 + need_temps = false # whether any arg needs temporaries + + # First pass: determine if any argument will needs temporaries + for j = 2:J + exj = ref_ex.args[j] + if isexpr(exj, :...) + need_temps = true + break + end + end + + # Second pass: if any need temps, create temps for all args + temp_vars = Tuple{Int,Symbol}[] + for j = 2:J + n = nx === 0 ? ni : :($nx + $ni) + exj, used = replace_ref_begin_end_!(__module__, ref_ex.args[j], (:($firstindex($S,$n)),:($lastindex($S,$n))), in_quote_context, escs) + used_S |= used + ref_ex.args[j] = exj + ni += 1 + if need_temps + isva = isexpr(exj, :...) # implied need_temps + if isva + exj = exj.args[1] + end + if isa_ast_node(exj) # create temp to preserve evaluation order and count in case `used` gets set later + exj = gensym(:arg) + push!(temp_vars, (j, exj)) + end + if isva + ni -= 1 + nx = nx === 0 ? :(length($exj)) : :($nx + length($exj)) + end + end + end + + # Third pass: if `used`, need to actually make those temp assignments now + if used_S + for (j, temp_var) in temp_vars + exj = ref_ex.args[j] + isva = isexpr(exj, :...) # implied need_temps + if isva + exj = exj.args[1] + end + push!(assignments, :(local $temp_var = $exj)) + ref_ex.args[j] = isva ? Expr(:..., temp_var) : temp_var + end + end + end + + if used_S + S0 = ref_ex.args[1] + S = escapes(S, escs) + ref_ex.args[1] = S + main_ex = :(local $S = $S0; $(assignments...); $main_ex) + end + return main_ex, used_withex + end if ex isa Expr && ex.head === :macrocall # Blithly modifying the arguments to another macro is unwise, so call # macroexpand first on it. @@ -46,77 +121,7 @@ function replace_ref_begin_end_!(__module__::Module, ex, withex, in_quote_contex end elseif isa(ex,Expr) 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 - ex.args[1], used_withex = replace_ref_begin_end_!(__module__, ex.args[1], withex, in_quote_context, escs) - 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 - assignments = [] - used_S = false # whether we actually need S - # new :ref, so redefine withex - nargs = length(ex.args)-1 - if nargs == 0 - return ex, used_withex - elseif nargs == 1 - # replace with lastindex(S) - ex.args[2], used_S = replace_ref_begin_end_!(__module__, ex.args[2], (:($firstindex($S)),:($lastindex($S))), in_quote_context, escs) - else - ni = 1 - nx = 0 - J = lastindex(ex.args) - need_temps = false # whether any arg needs temporaries - - # First pass: determine if any argument will needs temporaries - for j = 2:J - exj = ex.args[j] - if isexpr(exj, :...) - need_temps = true - break - end - end - - # Second pass: if any need temps, create temps for all args - temp_vars = Tuple{Int,Symbol}[] - for j = 2:J - n = nx === 0 ? ni : :($nx + $ni) - exj, used = replace_ref_begin_end_!(__module__, ex.args[j], (:($firstindex($S,$n)),:($lastindex($S,$n))), in_quote_context, escs) - used_S |= used - ex.args[j] = exj - ni += 1 - if need_temps - isva = isexpr(exj, :...) # implied need_temps - if isva - exj = exj.args[1] - end - if isa_ast_node(exj) # create temp to preserve evaluation order and count in case `used` gets set later - exj = gensym(:arg) - push!(temp_vars, (j, exj)) - end - if isva - ni -= 1 - nx = nx === 0 ? :(length($exj)) : :($nx + length($exj)) - end - end - end - - # Third pass: if `used`, need to actually make those temp assignments now - if used_S - for (j, temp_var) in temp_vars - exj = ex.args[j] - isva = isexpr(exj, :...) # implied need_temps - if isva - exj = exj.args[1] - end - push!(assignments, :(local $temp_var = $exj)) - ex.args[j] = isva ? Expr(:..., temp_var) : temp_var - end - end - end - - if used_S - S0 = ex.args[1] - S = escapes(S, escs) - ex.args[1] = S - ex = :(local $S = $S0; $(assignments...); $ex) - end - return ex, used_withex + return handle_refexpr!(__module__, ex, ex, withex, in_quote_context, escs) elseif ex.head === :$ # no longer an executable expression (handle all equivalent forms of :inert, :quote, and QuoteNode the same way) in_quote_context = false @@ -132,6 +137,18 @@ function replace_ref_begin_end_!(__module__::Module, ex, withex, in_quote_contex escs -= 1 elseif ex.head === :meta || ex.head === :inert return ex, used_withex + elseif !in_quote_context && last(string(ex.head)) == '=' && Meta.isexpr(ex.args[1], :ref) + for i = eachindex(ex.args) + if i == 1 + # we'll deal with the ref expression later + continue + end + ex.args[i], used = replace_ref_begin_end_!(__module__, ex.args[i], withex, in_quote_context, escs) + used_withex |= used + end + ex, used = handle_refexpr!(__module__, ex.args[1]::Expr, ex, withex, in_quote_context, escs) + used_withex |= used + return ex, used_withex end # recursive search for i = eachindex(ex.args) diff --git a/test/subarray.jl b/test/subarray.jl index 45ad4851b3b56..977a06442d7c2 100644 --- a/test/subarray.jl +++ b/test/subarray.jl @@ -1101,3 +1101,55 @@ end @test @views quote var"begin" + var"end" end isa Expr + +@testset "@views handling of assignment" begin + @test @macroexpand(@views x[a:b] = c) == :(x[a:b] = c) + # Assignments should still work + let array = [1, 2, 3, 4, 5, 6, 7, 8] + @views array[begin:2] = [-1, -2] + @test array == [-1, -2, 3, 4, 5, 6, 7, 8] + @views array[7:end] = [-7, -8] + @test array == [-1, -2, 3, 4, 5, 6, -7, -8] + @views array[begin + 2:end - 4] = [-3, -4] + @test array == [-1, -2, -3, -4, 5, 6, -7, -8] + @views identity(array)[begin + 4:end - 2] = [-5, -6] + @test array == [-1, -2, -3, -4, -5, -6, -7, -8] + + @views array[begin:2] .= 100 + @test array == [100, 100, -3, -4, -5, -6, -7, -8] + @views array[7:end] .= 200 + @test array == [100, 100, -3, -4, -5, -6, 200, 200] + @views array[begin + 2:end - 4] .= 300 + @test array == [100, 100, 300, 300, -5, -6, 200, 200] + @views identity(array)[begin + 4:end - 2] .= 400 + @test array == [100, 100, 300, 300, 400, 400, 200, 200] + + @views identity(array)[begin:end] .-= 1 + @test array == [99, 99, 299, 299, 399, 399, 199, 199] + + @views identity(array)[begin:end] += [1, 2, 3, 4, 5, 6, 7, 8] + @test array == [100, 101, 302, 303, 404, 405, 206, 207] + end + # Nested getindex in assignment should be transformed + let array = [1, 2, 3, 4, 5, 6, 7, 8], array2 = [1, 2, 3, 4, 5, 6, 7, 8] + array[begin + 1:end - 2][2] = -1 + array[begin + 1:end - 2][end] = -2 + @test array == [1, 2, 3, 4, 5, 6, 7, 8] + + @views array[begin + 1:end - 2][2] = -1 + @views array[begin + 1:end - 2][end] = -2 + @test array == [1, 2, -1, 4, 5, -2, 7, 8] + + function swap_ele(ary, i, v) + res = ary[i] + ary[i] = v + return res + end + array2[swap_ele(array[begin:end], 1, -3):swap_ele(array[begin:end], 7, -4)] = [-1, -2, -3, -4, -5, -6, -7] + @test array == [1, 2, -1, 4, 5, -2, 7, 8] + @test array2 == [-1, -2, -3, -4, -5, -6, -7, 8] + @views array2[swap_ele(array[begin:end], 1, -3):swap_ele(array[begin:end], 7, -4)] = [-10, 2, -30, 4, -50, 6, -70] + @test array == [-3, 2, -1, 4, 5, -2, -4, 8] + @test array2 == [-10, 2, -30, 4, -50, 6, -70, 8] + end +end