forked from JuliaLang/julia
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsubarray.jl
539 lines (476 loc) · 21.6 KB
/
subarray.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
# This file is a part of Julia. License is MIT: http://julialang.org/license
abstract type AbstractCartesianIndex{N} end # This is a hacky forward declaration for CartesianIndex
typealias ViewIndex Union{Real, AbstractArray}
typealias ScalarIndex Real
# L is true if the view itself supports fast linear indexing
struct SubArray{T,N,P,I,L} <: AbstractArray{T,N}
parent::P
indexes::I
offset1::Int # for linear indexing and pointer, only valid when L==true
stride1::Int # used only for linear indexing
function SubArray{T,N,P,I,L}(parent, indexes, offset1, stride1) where {T,N,P,I,L}
@_inline_meta
check_parent_index_match(parent, indexes)
new(parent, indexes, offset1, stride1)
end
end
# Compute the linear indexability of the indices, and combine it with the linear indexing of the parent
function SubArray(parent::AbstractArray, indexes::Tuple)
@_inline_meta
SubArray(linearindexing(viewindexing(indexes), linearindexing(parent)), parent, ensure_indexable(indexes), index_dimsum(indexes...))
end
function SubArray(::LinearSlow, parent::P, indexes::I, ::NTuple{N,Any}) where {P,I,N}
@_inline_meta
SubArray{eltype(P), N, P, I, false}(parent, indexes, 0, 0)
end
function SubArray(::LinearFast, parent::P, indexes::I, ::NTuple{N,Any}) where {P,I,N}
@_inline_meta
# Compute the stride and offset
stride1 = compute_stride1(parent, indexes)
SubArray{eltype(P), N, P, I, true}(parent, indexes, compute_offset1(parent, stride1, indexes), stride1)
end
check_parent_index_match(parent, indexes) = check_parent_index_match(parent, index_ndims(indexes...))
check_parent_index_match{T,N}(parent::AbstractArray{T,N}, ::NTuple{N, Bool}) = nothing
check_parent_index_match{N}(parent, ::NTuple{N, Bool}) =
throw(ArgumentError("number of indices ($N) must match the parent dimensionality ($(ndims(parent)))"))
# This computes the linear indexing compatability for a given tuple of indices
viewindexing() = LinearFast()
# Leading scalar indexes simply increase the stride
viewindexing(I::Tuple{ScalarIndex, Vararg{Any}}) = (@_inline_meta; viewindexing(tail(I)))
# Slices may begin a section which may be followed by any number of Slices
viewindexing(I::Tuple{Slice, Slice, Vararg{Any}}) = (@_inline_meta; viewindexing(tail(I)))
# A UnitRange can follow Slices, but only if all other indices are scalar
viewindexing(I::Tuple{Slice, UnitRange, Vararg{ScalarIndex}}) = LinearFast()
# In general, ranges are only fast if all other indices are scalar
viewindexing(I::Tuple{Union{Range, Slice}, Vararg{ScalarIndex}}) = LinearFast()
# All other index combinations are slow
viewindexing(I::Tuple{Vararg{Any}}) = LinearSlow()
# Of course, all other array types are slow
viewindexing(I::Tuple{AbstractArray, Vararg{Any}}) = LinearSlow()
# Simple utilities
size(V::SubArray) = (@_inline_meta; map(n->Int(unsafe_length(n)), indices(V)))
similar(V::SubArray, T::Type, dims::Dims) = similar(V.parent, T, dims)
parent(V::SubArray) = V.parent
parentindexes(V::SubArray) = V.indexes
parent(a::AbstractArray) = a
"""
parentindexes(A)
From an array view `A`, returns the corresponding indexes in the parent.
"""
parentindexes(a::AbstractArray) = ntuple(i->OneTo(size(a,i)), ndims(a))
## SubArray creation
# We always assume that the dimensionality of the parent matches the number of
# indices that end up getting passed to it, so we store the parent as a
# ReshapedArray view if necessary. The trouble is that arrays of `CartesianIndex`
# can make the number of effective indices not equal to length(I).
_maybe_reshape_parent(A::AbstractArray, ::NTuple{1, Bool}) = reshape(A, Val{1})
_maybe_reshape_parent{_}(A::AbstractArray{_,1}, ::NTuple{1, Bool}) = reshape(A, Val{1})
_maybe_reshape_parent{_,N}(A::AbstractArray{_,N}, ::NTuple{N, Bool}) = A
_maybe_reshape_parent{N}(A::AbstractArray, ::NTuple{N, Bool}) = reshape(A, Val{N}) # TODO: DEPRECATE FOR #14770
"""
view(A, inds...)
Like [`getindex`](@ref), but returns a view into the parent array `A` with the
given indices instead of making a copy. Calling [`getindex`](@ref) or
[`setindex!`](@ref) on the returned `SubArray` computes the
indices to the parent array on the fly without checking bounds.
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
1 2
3 4
julia> b = view(A, :, 1)
2-element SubArray{Int64,1,Array{Int64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}:
1
3
julia> fill!(b, 0)
2-element SubArray{Int64,1,Array{Int64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}:
0
0
julia> A # Note A has changed even though we modified b
2×2 Array{Int64,2}:
0 2
0 4
```
"""
function view(A::AbstractArray, I...)
@_inline_meta
J = to_indices(A, I)
@boundscheck checkbounds(A, J...)
unsafe_view(_maybe_reshape_parent(A, index_ndims(J...)), J...)
end
function unsafe_view(A::AbstractArray, I::ViewIndex...)
@_inline_meta
SubArray(A, I)
end
# When we take the view of a view, it's often possible to "reindex" the parent
# view's indices such that we can "pop" the parent view and keep just one layer
# of indirection. But we can't always do this because arrays of `CartesianIndex`
# might span multiple parent indices, making the reindex calculation very hard.
# So we use _maybe_reindex to figure out if there are any arrays of
# `CartesianIndex`, and if so, we punt and keep two layers of indirection.
unsafe_view(V::SubArray, I::ViewIndex...) = (@_inline_meta; _maybe_reindex(V, I))
_maybe_reindex(V, I) = (@_inline_meta; _maybe_reindex(V, I, I))
_maybe_reindex(V, I, ::Tuple{AbstractArray{<:AbstractCartesianIndex}, Vararg{Any}}) =
(@_inline_meta; SubArray(V, I))
# But allow arrays of CartesianIndex{1}; they behave just like arrays of Ints
_maybe_reindex(V, I, A::Tuple{AbstractArray{<:AbstractCartesianIndex{1}}, Vararg{Any}}) =
(@_inline_meta; _maybe_reindex(V, I, tail(A)))
_maybe_reindex(V, I, A::Tuple{Any, Vararg{Any}}) = (@_inline_meta; _maybe_reindex(V, I, tail(A)))
function _maybe_reindex(V, I, ::Tuple{})
@_inline_meta
@inbounds idxs = to_indices(V.parent, reindex(V, V.indexes, I))
SubArray(V.parent, idxs)
end
## Re-indexing is the heart of a view, transforming A[i, j][x, y] to A[i[x], j[y]]
#
# Recursively look through the heads of the parent- and sub-indexes, considering
# the following cases:
# * Parent index is array -> re-index that with one or more sub-indexes (one per dimension)
# * Parent index is Colon -> just use the sub-index as provided
# * Parent index is scalar -> that dimension was dropped, so skip the sub-index and use the index as is
typealias AbstractZeroDimArray{T} AbstractArray{T, 0}
reindex(V, ::Tuple{}, ::Tuple{}) = ()
# Skip dropped scalars, so simply peel them off the parent indices and continue
reindex(V, idxs::Tuple{ScalarIndex, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) =
(@_propagate_inbounds_meta; (idxs[1], reindex(V, tail(idxs), subidxs)...))
# Slices simply pass their subindexes straight through
reindex(V, idxs::Tuple{Slice, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) =
(@_propagate_inbounds_meta; (subidxs[1], reindex(V, tail(idxs), tail(subidxs))...))
# Re-index into parent vectors with one subindex
reindex(V, idxs::Tuple{AbstractVector, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) =
(@_propagate_inbounds_meta; (idxs[1][subidxs[1]], reindex(V, tail(idxs), tail(subidxs))...))
# Parent matrices are re-indexed with two sub-indices
reindex(V, idxs::Tuple{AbstractMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any, Vararg{Any}}) =
(@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], reindex(V, tail(idxs), tail(tail(subidxs)))...))
# In general, we index N-dimensional parent arrays with N indices
@generated function reindex{T,N}(V, idxs::Tuple{AbstractArray{T,N}, Vararg{Any}}, subidxs::Tuple{Vararg{Any}})
if length(subidxs.parameters) >= N
subs = [:(subidxs[$d]) for d in 1:N]
tail = [:(subidxs[$d]) for d in N+1:length(subidxs.parameters)]
:(@_propagate_inbounds_meta; (idxs[1][$(subs...)], reindex(V, tail(idxs), ($(tail...),))...))
else
:(throw(ArgumentError("cannot re-index $(ndims(V)) dimensional SubArray with fewer than $(ndims(V)) indices\nThis should not occur; please submit a bug report.")))
end
end
# In general, we simply re-index the parent indices by the provided ones
typealias SlowSubArray{T,N,P,I} SubArray{T,N,P,I,false}
function getindex{T,N}(V::SlowSubArray{T,N}, I::Vararg{Int,N})
@_inline_meta
@boundscheck checkbounds(V, I...)
@inbounds r = V.parent[reindex(V, V.indexes, I)...]
r
end
typealias FastSubArray{T,N,P,I} SubArray{T,N,P,I,true}
function getindex(V::FastSubArray, i::Int)
@_inline_meta
@boundscheck checkbounds(V, i)
@inbounds r = V.parent[V.offset1 + V.stride1*i]
r
end
# We can avoid a multiplication if the first parent index is a Colon or UnitRange
typealias FastContiguousSubArray{T,N,P,I<:Tuple{Union{Slice, UnitRange}, Vararg{Any}}} SubArray{T,N,P,I,true}
function getindex(V::FastContiguousSubArray, i::Int)
@_inline_meta
@boundscheck checkbounds(V, i)
@inbounds r = V.parent[V.offset1 + i]
r
end
function setindex!{T,N}(V::SlowSubArray{T,N}, x, I::Vararg{Int,N})
@_inline_meta
@boundscheck checkbounds(V, I...)
@inbounds V.parent[reindex(V, V.indexes, I)...] = x
V
end
function setindex!(V::FastSubArray, x, i::Int)
@_inline_meta
@boundscheck checkbounds(V, i)
@inbounds V.parent[V.offset1 + V.stride1*i] = x
V
end
function setindex!(V::FastContiguousSubArray, x, i::Int)
@_inline_meta
@boundscheck checkbounds(V, i)
@inbounds V.parent[V.offset1 + i] = x
V
end
linearindexing(::Type{<:FastSubArray}) = LinearFast()
linearindexing(::Type{<:SubArray}) = LinearSlow()
# Strides are the distance between adjacent elements in a given dimension,
# so they are well-defined even for non-linear memory layouts
strides(V::SubArray) = substrides(V.parent, V.indexes)
substrides(parent, I::Tuple) = substrides(1, parent, 1, I)
substrides(s, parent, dim, ::Tuple{}) = ()
substrides(s, parent, dim, I::Tuple{ScalarIndex, Vararg{Any}}) = (substrides(s*size(parent, dim), parent, dim+1, tail(I))...)
substrides(s, parent, dim, I::Tuple{Slice, Vararg{Any}}) = (s, substrides(s*size(parent, dim), parent, dim+1, tail(I))...)
substrides(s, parent, dim, I::Tuple{Range, Vararg{Any}}) = (s*step(I[1]), substrides(s*size(parent, dim), parent, dim+1, tail(I))...)
substrides(s, parent, dim, I::Tuple{Any, Vararg{Any}}) = throw(ArgumentError("strides is invalid for SubArrays with indices of type $(typeof(I[1]))"))
stride(V::SubArray, d::Integer) = d <= ndims(V) ? strides(V)[d] : strides(V)[end] * size(V)[end]
compute_stride1{N}(parent::AbstractArray, I::NTuple{N,Any}) =
(@_inline_meta; compute_stride1(1, fill_to_length(indices(parent), OneTo(1), Val{N}), I))
compute_stride1(s, inds, I::Tuple{}) = s
compute_stride1(s, inds, I::Tuple{ScalarIndex, Vararg{Any}}) =
(@_inline_meta; compute_stride1(s*unsafe_length(inds[1]), tail(inds), tail(I)))
compute_stride1(s, inds, I::Tuple{Range, Vararg{Any}}) = s*step(I[1])
compute_stride1(s, inds, I::Tuple{Slice, Vararg{Any}}) = s
compute_stride1(s, inds, I::Tuple{Any, Vararg{Any}}) = throw(ArgumentError("invalid strided index type $(typeof(I[1]))"))
iscontiguous(A::SubArray) = iscontiguous(typeof(A))
iscontiguous(::Type{<:SubArray}) = false
iscontiguous(::Type{<:FastContiguousSubArray}) = true
first_index(V::FastSubArray) = V.offset1 + V.stride1 # cached for fast linear SubArrays
function first_index(V::SubArray)
P, I = parent(V), V.indexes
s1 = compute_stride1(P, I)
s1 + compute_offset1(P, s1, I)
end
# Computing the first index simply steps through the indices, accumulating the
# sum of index each multiplied by the parent's stride.
# The running sum is `f`; the cumulative stride product is `s`.
# If the result is one-dimensional and it's a Colon, then linear
# indexing uses the indices along the given dimension. Otherwise
# linear indexing always starts with 1.
compute_offset1(parent, stride1::Integer, I::Tuple) =
(@_inline_meta; compute_offset1(parent, stride1, find_extended_dims(I)..., I))
compute_offset1(parent, stride1::Integer, dims::Tuple{Int}, inds::Tuple{Slice}, I::Tuple) =
(@_inline_meta; compute_linindex(parent, I) - stride1*first(indices(parent, dims[1]))) # index-preserving case
compute_offset1(parent, stride1::Integer, dims, inds, I::Tuple) =
(@_inline_meta; compute_linindex(parent, I) - stride1) # linear indexing starts with 1
function compute_linindex{N}(parent, I::NTuple{N,Any})
@_inline_meta
IP = fill_to_length(indices(parent), OneTo(1), Val{N})
compute_linindex(1, 1, IP, I)
end
function compute_linindex(f, s, IP::Tuple, I::Tuple{ScalarIndex, Vararg{Any}})
@_inline_meta
Δi = I[1]-first(IP[1])
compute_linindex(f + Δi*s, s*unsafe_length(IP[1]), tail(IP), tail(I))
end
function compute_linindex(f, s, IP::Tuple, I::Tuple{Any, Vararg{Any}})
@_inline_meta
Δi = first(I[1])-first(IP[1])
compute_linindex(f + Δi*s, s*unsafe_length(IP[1]), tail(IP), tail(I))
end
compute_linindex(f, s, IP::Tuple, I::Tuple{}) = f
find_extended_dims(I) = (@_inline_meta; _find_extended_dims((), (), 1, I...))
_find_extended_dims(dims, inds, dim) = dims, inds
_find_extended_dims(dims, inds, dim, ::ScalarIndex, I...) =
(@_inline_meta; _find_extended_dims(dims, inds, dim+1, I...))
_find_extended_dims(dims, inds, dim, i1, I...) =
(@_inline_meta; _find_extended_dims((dims..., dim), (inds..., i1), dim+1, I...))
unsafe_convert{T,N,P}(::Type{Ptr{T}}, V::SubArray{T,N,P,<:Tuple{Vararg{RangeIndex}}}) =
unsafe_convert(Ptr{T}, V.parent) + (first_index(V)-1)*sizeof(T)
pointer(V::FastSubArray, i::Int) = pointer(V.parent, V.offset1 + V.stride1*i)
pointer(V::FastContiguousSubArray, i::Int) = pointer(V.parent, V.offset1 + i)
pointer(V::SubArray, i::Int) = _pointer(V, i)
_pointer(V::SubArray{<:Any,1}, i::Int) = pointer(V, (i,))
_pointer(V::SubArray, i::Int) = pointer(V, ind2sub(indices(V), i))
function pointer{T,N}(V::SubArray{T,N,<:Array,<:Tuple{Vararg{RangeIndex}}}, is::Tuple{Vararg{Int}})
index = first_index(V)
strds = strides(V)
for d = 1:length(is)
index += (is[d]-1)*strds[d]
end
return pointer(V.parent, index)
end
# indices are taken from the range/vector
# Since bounds-checking is performance-critical and uses
# indices, it's worth optimizing these implementations thoroughly
indices(S::SubArray) = (@_inline_meta; _indices_sub(S, S.indexes...))
_indices_sub(S::SubArray) = ()
_indices_sub(S::SubArray, ::Real, I...) = (@_inline_meta; _indices_sub(S, I...))
function _indices_sub(S::SubArray, i1::AbstractArray, I...)
@_inline_meta
(unsafe_indices(i1)..., _indices_sub(S, I...)...)
end
## Compatability
# deprecate?
function parentdims(s::SubArray)
nd = ndims(s)
dimindex = Array{Int}(nd)
sp = strides(s.parent)
sv = strides(s)
j = 1
for i = 1:ndims(s.parent)
r = s.indexes[i]
if j <= nd && (isa(r,Union{Slice,Range}) ? sp[i]*step(r) : sp[i]) == sv[j]
dimindex[j] = i
j += 1
end
end
dimindex
end
"""
replace_ref_end!(ex)
Recursively replace occurrences 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)]]
"""
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 endof(S)
ex.args[2], used_S = replace_ref_end_!(ex.args[2],:($endof($S)))
else
n = 1
J = endof(ex.args)
for j = 2:J-1
exj, used = replace_ref_end_!(ex.args[j],:($size($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
ex.args[J], used = replace_ref_end_!(ex.args[J],:($trailingsize($S,$n)))
used_S |= used
end
if used_S && S !== ex.args[1]
S0 = ex.args[1]
ex.args[1] = S
ex = Expr(:let, ex, :($S = $S0))
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 SubArray{Int64,1,Array{Int64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}:
1
3
julia> fill!(b, 0)
2-element SubArray{Int64,1,Array{Int64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}:
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[1], :ref))
ex.args[1] = Expr(:call, view, ex.args[1].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
# size and endof 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(first(h) == '.' ? :(.=) : :(=), :($a[$(I...)]),
Expr(:call, Symbol(h[1:end-1]),
:($maybeview($a, $(I...))),
_views.(ex.args[2:end])...)),
:($a = $(_views(lhs.args[1]))),
[:($(i[k]) = $(_views(lhs.args[k+1]))) for k=1:length(i)]...)
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 that 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.
"""
macro views(x)
esc(_views(replace_ref_end!(x)))
end