"""
replace_ref_end!(ex)
Recursively replace occurrences of the symbol :end in a "ref" expression (i.e. A[...]) `ex`
with the appropriate function calls (`lastindex` or `size`). Replacement uses
the closest enclosing ref, so
A[B[end]]
should transform to
A[B[lastindex(B)]]
"""
replace_ref_end!(ex) = replace_ref_end_!(ex, nothing)[1]
# replace_ref_end_!(ex,withex) returns (new ex, whether withex was used)
function replace_ref_end_!(ex, withex)
used_withex = false
if isa(ex,Symbol) && ex == :end
withex === nothing && error("Invalid use of end")
return withex, true
elseif isa(ex,Expr)
if ex.head == :ref
ex.args[1], used_withex = replace_ref_end_!(ex.args[1],withex)
S = isa(ex.args[1],Symbol) ? ex.args[1]::Symbol : gensym(:S) # temp var to cache ex.args[1] if needed
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_end_!(ex.args[2],:($lastindex($S)))
else
n = 1
J = lastindex(ex.args)
for j = 2:J
exj, used = replace_ref_end_!(ex.args[j],:($lastindex($S,$n)))
used_S |= used
ex.args[j] = exj
if isa(exj,Expr) && exj.head == :...
# splatted object
exjs = exj.args[1]
n = :($n + length($exjs))
elseif isa(n, Expr)
# previous expression splatted
n = :($n + 1)
else
# an integer
n += 1
end
end
end
if used_S && S !== ex.args[1]
S0 = ex.args[1]
ex.args[1] = S
ex = Expr(:let, :($S = $S0), ex)
end
else
# recursive search
for i = eachindex(ex.args)
ex.args[i], used = replace_ref_end_!(ex.args[i],withex)
used_withex |= used
end
end
end
ex, used_withex
end
"""
@view A[inds...]
Creates a `SubArray` from an indexing expression. This can only be applied directly to a
reference expression (e.g. `@view A[1,2:end]`), and should *not* be used as the target of
an assignment (e.g. `@view(A[1,2:end]) = ...`). See also [`@views`](@ref)
to switch an entire block of code to use views for slicing.
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
1 2
3 4
julia> b = @view A[:, 1]
2-element view(::Array{Int64,2}, :, 1) with eltype Int64:
1
3
julia> fill!(b, 0)
2-element view(::Array{Int64,2}, :, 1) with eltype Int64:
0
0
julia> A
2×2 Array{Int64,2}:
0 2
0 4
```
"""
macro view(ex)
if Meta.isexpr(ex, :ref)
ex = replace_ref_end!(ex)
if Meta.isexpr(ex, :ref)
ex = Expr(:call, view, ex.args...)
else # ex replaced by let ...; foo[...]; end
@assert Meta.isexpr(ex, :let) && Meta.isexpr(ex.args[2], :ref)
ex.args[2] = Expr(:call, view, ex.args[2].args...)
end
Expr(:&&, true, esc(ex))
else
throw(ArgumentError("Invalid use of @view macro: argument must be a reference expression A[...]."))
end
end
############################################################################
# @views macro code:
# maybeview is like getindex, but returns a view for slicing operations
# (while remaining equivalent to getindex for scalar indices and non-array types)
@propagate_inbounds maybeview(A, args...) = getindex(A, args...)
@propagate_inbounds maybeview(A::AbstractArray, args...) = view(A, args...)
@propagate_inbounds maybeview(A::AbstractArray, args::Number...) = getindex(A, args...)
@propagate_inbounds maybeview(A) = getindex(A)
@propagate_inbounds maybeview(A::AbstractArray) = getindex(A)
# _views implements the transformation for the @views macro.
# @views calls esc(_views(...)) to work around #20241,
# so any function calls we insert (to maybeview, or to
# lastindex in replace_ref_end!) must be interpolated
# as values rather than as symbols to ensure that they are called
# from Base rather than from the caller's scope.
_views(x) = x
function _views(ex::Expr)
if ex.head in (:(=), :(.=))
# don't use view for ref on the lhs of an assignment,
# but still use views for the args of the ref:
lhs = ex.args[1]
Expr(ex.head, Meta.isexpr(lhs, :ref) ?
Expr(:ref, _views.(lhs.args)...) : _views(lhs),
_views(ex.args[2]))
elseif ex.head == :ref
Expr(:call, maybeview, _views.(ex.args)...)
else
h = string(ex.head)
# don't use view on the lhs of an op-assignment a[i...] += ...
if last(h) == '=' && Meta.isexpr(ex.args[1], :ref)
lhs = ex.args[1]
# temp vars to avoid recomputing a and i,
# which will be assigned in a let block:
a = gensym(:a)
i = [gensym(:i) for k = 1:length(lhs.args)-1]
# for splatted indices like a[i, j...], we need to
# splat the corresponding temp var.
I = similar(i, Any)
for k = 1:length(i)
if Meta.isexpr(lhs.args[k+1], :...)
I[k] = Expr(:..., i[k])
lhs.args[k+1] = lhs.args[k+1].args[1] # unsplat
else
I[k] = i[k]
end
end
Expr(:let,
Expr(:block,
:($a = $(_views(lhs.args[1]))),
[:($(i[k]) = $(_views(lhs.args[k+1]))) for k=1:length(i)]...),
Expr(first(h) == '.' ? :(.=) : :(=), :($a[$(I...)]),
Expr(:call, Symbol(h[1:end-1]),
:($maybeview($a, $(I...))),
_views.(ex.args[2:end])...)))
else
Expr(ex.head, _views.(ex.args)...)
end
end
end
"""
@views expression
Convert every array-slicing operation in the given expression
(which may be a `begin`/`end` block, loop, function, etc.)
to return a view. Scalar indices, non-array types, and
explicit `getindex` calls (as opposed to `array[...]`) are
unaffected.
!!! note
The `@views` macro only affects `array[...]` expressions
that appear explicitly in the given `expression`, not array slicing that
occurs in functions called by that code.
# Examples
```jldoctest
julia> A = zeros(3, 3);
julia> @views for row in 1:3
b = A[row, :]
b[:] = row
end
julia> A
3×3 Array{Float64,2}:
1.0 1.0 1.0
2.0 2.0 2.0
3.0 3.0 3.0
```
"""
macro views(x)
esc(_views(replace_ref_end!(x)))
end