Skip to content

Commit

Permalink
Fix #2644 in release-0.1
Browse files Browse the repository at this point in the history
  • Loading branch information
timholy committed Mar 22, 2013
1 parent a703335 commit 2730b5a
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 31 deletions.
78 changes: 47 additions & 31 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,8 @@ function sub(A::SubArray, i::RangeIndex...)
j += 1
end
end
sub(A.parent, tuple(newindexes...))
ni = tuple(newindexes...)
SubArray{eltype(A),L,typeof(A.parent),typeof(ni)}(A.parent, ni)
end

function slice{T,N}(A::AbstractArray{T,N}, i::NTuple{N,RangeIndex})
Expand Down Expand Up @@ -188,33 +189,63 @@ function ref{T,S<:Integer}(s::SubArray{T,1}, I::AbstractVector{S})
ref(s.parent, t)
end

function translate_indexes(s::SubArray, I::Union(Real,AbstractArray)...)
I = indices(I)
nds = ndims(s)
n = length(I)
if n > nds
throw(BoundsError())
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 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
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)
idx = Array(Int, length(I))
ssztail = size(s)[n:]
psztail = size(s.parent)[n:]
pdims = parentdims(s)
psztail = size(s.parent)[pdims[n:]]
for j=1:length(I)
su = ind2sub(ssztail,I[j])
idx[j] = sub2ind(psztail, [ s.indexes[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) ]...)
end
idx
end

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

function ref(s::SubArray, I::Union(Real,AbstractArray)...)
newindexes = translate_indexes(s, I...)

rs = ref_shape(I...)
result = ref(s.parent, newindexes...)
Expand Down Expand Up @@ -270,22 +301,7 @@ function assign{T,S<:Integer}(s::SubArray{T,1}, v, I::AbstractVector{S})
end

function assign(s::SubArray, v, I::Union(Real,AbstractArray)...)
I = indices(I)
j = 1 #the jth dimension in subarray
ndp = ndims(s.parent)
n = length(I)
newindexes = cell(n)
for i = 1:n
t = s.indexes[i]
#TODO: don't generate the dense vector indexes if they can be ranges
if i==n && n < ndp
newindexes[i] = translate_linear_indexes(s, i, I[i])
else
newindexes[i] = isa(t, Int) ? t : t[I[j]]
end
j += 1
end

newindexes = translate_indexes(s, I...)
assign(s.parent, v, newindexes...)
end

Expand Down
22 changes: 22 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,28 @@ b = [4, 6, 2, -7, 1]
ind = findin(a, b)
@test ind == [3,4]

# sub
A = reshape(1:120, 3, 5, 8)
sA = sub(A, 2, 1:5, 1:8)
@test size(sA) == (1, 5, 8)
@test_fails sA[2, 1:8]
@test sA[1, 2, 1:8][:] == 5:15:120
sA[2:5:end] = -1
@test all(sA[2:5:end] .== -1)
@test all(A[5:15:120] .== -1)

# slice
A = reshape(1:120, 3, 5, 8)
sA = slice(A, 2, 1:5, 1:8)
@test size(sA) == (5, 8)
@test sA[2, 1:8][:] == 5:15:120
@test sA[:,1] == 2:3:14
@test sA[2:5:end] == 5:15:120
sA[2:5:end] = -1
@test all(sA[2:5:end] .== -1)
@test all(A[5:15:120] .== -1)


# get
let
A = reshape(1:24, 3, 8)
Expand Down

0 comments on commit 2730b5a

Please sign in to comment.