Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Create @view macro for creating SubArrays via indexing. #16564

Merged
merged 1 commit into from
Jun 27, 2016
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions base/exports.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1412,6 +1412,7 @@ export
@enum,
@label,
@goto,
@view,

# SparseArrays module re-exports
SparseArrays,
Expand Down
73 changes: 73 additions & 0 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -318,3 +318,76 @@ function parentdims(s::SubArray)
end
dimindex
end

"""
replace_ref_end!(ex)

Recursively replace occurences of the symbol :end in a "ref" expression (i.e. A[...]) `ex`
with the appropriate function calls (`endof`, `size` or `trailingsize`). Replacement uses
the closest enclosing ref, so

A[B[end]]

should transform to

A[B[endof(B)]]

"""
function replace_ref_end!(ex,withex=nothing)
if isa(ex,Symbol) && ex == :end
withex == nothing && error("Invalid use of end")
return withex
elseif isa(ex,Expr)
if ex.head == :ref
S = ex.args[1] = replace_ref_end!(ex.args[1],withex)
# new :ref, so redefine withex
nargs = length(ex.args)-1
if nargs == 0
return ex
elseif nargs == 1
# replace with endof(S)
ex.args[2] = replace_ref_end!(ex.args[2],:(Base.endof($S)))
else
n = 1
J = endof(ex.args)
for j = 2:J-1
exj = ex.args[j] = replace_ref_end!(ex.args[j],:(Base.size($S,$n)))
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
ex.args[J] = replace_ref_end!(ex.args[J],:(Base.trailingsize($S,$n)))
end
else
# recursive search
for i = eachindex(ex.args)
ex.args[i] = replace_ref_end!(ex.args[i],withex)
end
end
end
ex
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]) = ...`).
"""
macro view(ex)
if isa(ex, Expr) && ex.head == :ref
ex = replace_ref_end!(ex)
Expr(:&&, true, esc(Expr(:call,:(Base.view),ex.args...)))
else
throw(ArgumentError("Invalid use of @view macro: argument must be a reference expression A[...]."))
end
end
6 changes: 6 additions & 0 deletions doc/stdlib/arrays.rst
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,12 @@ Indexing, Assignment, and Concatenation

Like :func:`getindex`\ , but returns a view into the parent array ``A`` with the given indices instead of making a copy. Calling :func:`getindex` or :func:`setindex!` on the returned :obj:`SubArray` computes the indices to the parent array on the fly without checking bounds.

.. function:: @view A[inds...]

.. Docstring generated from Julia source

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]) = ...``\ ).

.. function:: parent(A)

.. Docstring generated from Julia source
Expand Down
26 changes: 26 additions & 0 deletions test/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -465,3 +465,29 @@ end
# the following segfaults with LLVM 3.8 on Windows, ref #15417
@test collect(view(view(reshape(1:13^3, 13, 13, 13), 3:7, 6:6, :), 1:2:5, :, 1:2:5)) ==
cat(3,[68,70,72],[406,408,410],[744,746,748])



# tests @view (and replace_ref_end!)
X = reshape(1:24,2,3,4)
Y = 4:-1:1

@test isa(@view(X[1:3]), SubArray)


@test X[1:end] == @view X[1:end]
@test X[1:end-3] == @view X[1:end-3]
@test X[1:end,2,2] == @view X[1:end,2,2]
@test X[1,1:end-2] == @view X[1,1:end-2]
@test X[1,2,1:end-2] == @view X[1,2,1:end-2]
@test X[1,2,Y[2:end]] == @view X[1,2,Y[2:end]]
@test X[1:end,2,Y[2:end]] == @view X[1:end,2,Y[2:end]]

u = (1,2:3)
@test X[u...,2:end] == @view X[u...,2:end]
@test X[(1,)...,(2,)...,2:end] == @view X[(1,)...,(2,)...,2:end]

# test macro hygiene
let size=(x,y)-> error("should not happen")
@test X[1:end,2,2] == @view X[1:end,2,2]
end