Skip to content

Commit

Permalink
Fix #3755
Browse files Browse the repository at this point in the history
Significantly simpler and shorter than what we had previously. Should
be more maintainable.
  • Loading branch information
timholy committed Jul 18, 2013
1 parent c9a537a commit 7eb374e
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 26 deletions.
44 changes: 18 additions & 26 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -205,53 +205,45 @@ function getindex{T,S<:Integer}(s::SubArray{T,1}, I::AbstractVector{S})
end

function translate_indexes(s::SubArray, I::Union(Real,AbstractArray)...)
I = indices(I)
nds = ndims(s)
n = length(I)
if n > nds
throw(BoundsError())
newindexes = Any[s.indexes...]
pdims = parentdims(s)
havelinear = n < ndims(s)
for i = 1:n-havelinear
newindexes[pdims[i]] = s.indexes[pdims[i]][I[i]]
end
ndp = ndims(s.parent) - (nds-n)
newindexes = Array(Any, ndp)
sp = strides(s.parent)
j = 1
for i = 1:ndp
t = s.indexes[i]
if j <= nds && s.strides[j] == sp[i]
#TODO: don't generate the dense vector indexes if they can be ranges
if j==n && n < nds
newindexes[i] = translate_linear_indexes(s, j, I[j])
else
newindexes[i] = isa(t, Int) ? t : t[I[j]]
end
j += 1
else
newindexes[i] = t
end
lastdim = pdims[n]
if havelinear
newindexes = newindexes[1:lastdim]
newindexes[pdims[n]] = translate_linear_indexes(s, n, I[end], pdims)
end
newindexes
end

# translate a linear index vector I for dim n to a linear index vector for
# the parent array
function translate_linear_indexes(s, n, I)
function translate_linear_indexes(s, n, I, pdims)
idx = Array(Int, length(I))
ssztail = size(s)[n:]
pdims = parentdims(s)
psztail = size(s.parent)[pdims[n:]]
taildimsoffset = 0
for i = pdims[end]+1:ndims(s.parent)
taildimsoffset += (s.indexes[i]-1)*stride(s.parent, i)
end
for j=1:length(I)
su = ind2sub(ssztail,I[j])
idx[j] = sub2ind(psztail, [ s.indexes[pdims[n+k-1]][su[k]] for k=1:length(su) ]...)
idx[j] = sub2ind(psztail, [ s.indexes[pdims[n+k-1]][su[k]] for k=1:length(su) ]...) + taildimsoffset
end
idx
end

function parentdims(s::SubArray)
dimindex = Array(Int, ndims(s))
nd = ndims(s)
dimindex = Array(Int, nd)
sp = strides(s.parent)
j = 1
for i = 1:ndims(s.parent)
if sp[i] == s.strides[j]
if j <= nd && sp[i] == s.strides[j]
dimindex[j] = i
j += 1
end
Expand Down
2 changes: 2 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,8 @@ sA[2:5:end] = -1
sA = sub(A, 1:3, 1:5, 5)
sA[1:3,1:5] = -2
@test all(A[:,:,5] .== -2)
sA[:] = -3
@test all(A[:,:,5] .== -3)

# slice
A = reshape(1:120, 3, 5, 8)
Expand Down

1 comment on commit 7eb374e

@JeffBezanson
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Much appreciated! Index gymnastics routinely make by brain freeze. The new code does look much nicer too.

Please sign in to comment.