Skip to content

Commit 5bf8e41

Browse files
committed
Create @sub macro for creating SubArrays via indexing.
1 parent 01e3c8a commit 5bf8e41

File tree

4 files changed

+106
-0
lines changed

4 files changed

+106
-0
lines changed

base/exports.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1412,6 +1412,7 @@ export
14121412
@enum,
14131413
@label,
14141414
@goto,
1415+
@view,
14151416

14161417
# SparseArrays module re-exports
14171418
SparseArrays,

base/subarray.jl

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,3 +318,76 @@ function parentdims(s::SubArray)
318318
end
319319
dimindex
320320
end
321+
322+
"""
323+
replace_ref_end!(ex)
324+
325+
Recursively replace occurences of the symbol :end in a "ref" expression (i.e. A[...]) `ex`
326+
with the appropriate function calls (`endof`, `size` or `trailingsize`). Replacement uses
327+
the closest enclosing ref, so
328+
329+
A[B[end]]
330+
331+
should transform to
332+
333+
A[B[endof(B)]]
334+
335+
"""
336+
function replace_ref_end!(ex,withex=nothing)
337+
if isa(ex,Symbol) && ex == :end
338+
withex == nothing && error("Invalid use of end")
339+
return withex
340+
elseif isa(ex,Expr)
341+
if ex.head == :ref
342+
S = ex.args[1] = replace_ref_end!(ex.args[1],withex)
343+
# new :ref, so redefine withex
344+
nargs = length(ex.args)-1
345+
if nargs == 0
346+
return ex
347+
elseif nargs == 1
348+
# replace with endof(S)
349+
ex.args[2] = replace_ref_end!(ex.args[2],:(Base.endof($S)))
350+
else
351+
n = 1
352+
J = endof(ex.args)
353+
for j = 2:J-1
354+
exj = ex.args[j] = replace_ref_end!(ex.args[j],:(Base.size($S,$n)))
355+
if isa(exj,Expr) && exj.head == :...
356+
# splatted object
357+
exjs = exj.args[1]
358+
n = :($n + length($exjs))
359+
elseif isa(n, Expr)
360+
# previous expression splatted
361+
n = :($n + 1)
362+
else
363+
# an integer
364+
n += 1
365+
end
366+
end
367+
ex.args[J] = replace_ref_end!(ex.args[J],:(Base.trailingsize($S,$n)))
368+
end
369+
else
370+
# recursive search
371+
for i = eachindex(ex.args)
372+
ex.args[i] = replace_ref_end!(ex.args[i],withex)
373+
end
374+
end
375+
end
376+
ex
377+
end
378+
379+
"""
380+
@view A[inds...]
381+
382+
Creates a `SubArray` from an indexing expression. This can only be applied directly to a
383+
reference expression (e.g. `@view A[1,2:end]`), and should *not* be used as the target of
384+
an assignment (e.g. `@view(A[1,2:end]) = ...`).
385+
"""
386+
macro view(ex)
387+
if isa(ex, Expr) && ex.head == :ref
388+
ex = replace_ref_end!(ex)
389+
Expr(:&&, true, esc(Expr(:call,:(Base.view),ex.args...)))
390+
else
391+
throw(ArgumentError("Invalid use of @view macro: argument must be a reference expression A[...]."))
392+
end
393+
end

doc/stdlib/arrays.rst

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,12 @@ Indexing, Assignment, and Concatenation
383383
384384
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.
385385

386+
.. function:: @view A[inds...]
387+
388+
.. Docstring generated from Julia source
389+
390+
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]) = ...``\ ).
391+
386392
.. function:: parent(A)
387393

388394
.. Docstring generated from Julia source

test/subarray.jl

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -465,3 +465,29 @@ end
465465
# the following segfaults with LLVM 3.8 on Windows, ref #15417
466466
@test collect(view(view(reshape(1:13^3, 13, 13, 13), 3:7, 6:6, :), 1:2:5, :, 1:2:5)) ==
467467
cat(3,[68,70,72],[406,408,410],[744,746,748])
468+
469+
470+
471+
# tests @view (and replace_ref_end!)
472+
X = reshape(1:24,2,3,4)
473+
Y = 4:-1:1
474+
475+
@test isa(@view(X[1:3]), SubArray)
476+
477+
478+
@test X[1:end] == @view X[1:end]
479+
@test X[1:end-3] == @view X[1:end-3]
480+
@test X[1:end,2,2] == @view X[1:end,2,2]
481+
@test X[1,1:end-2] == @view X[1,1:end-2]
482+
@test X[1,2,1:end-2] == @view X[1,2,1:end-2]
483+
@test X[1,2,Y[2:end]] == @view X[1,2,Y[2:end]]
484+
@test X[1:end,2,Y[2:end]] == @view X[1:end,2,Y[2:end]]
485+
486+
u = (1,2:3)
487+
@test X[u...,2:end] == @view X[u...,2:end]
488+
@test X[(1,)...,(2,)...,2:end] == @view X[(1,)...,(2,)...,2:end]
489+
490+
# test macro hygiene
491+
let size=(x,y)-> error("should not happen")
492+
@test X[1:end,2,2] == @view X[1:end,2,2]
493+
end

0 commit comments

Comments
 (0)