fixed general case cat (fixes #116)

j/abstractarray.j

@@ -628,20 +628,40 @@ end
 ## cat: general case
 
 function cat(catdim::Int, X...)
- typeC = promote_type(map(typeof, X)...)
  nargs = length(X)
- local dimsC::(Size...)
- if catdim == 1
- dimsC = (nargs,)
- elseif catdim == 2
- dimsC = (1, nargs)
- else
- dimsC = tuple(ntuple(catdim-1, i->1)..., nargs)
+ dimsA = map((a->isa(a,AbstractArray) ? size(a) : (1,)), X)
+ ndimsA = map((a->isa(a,AbstractArray) ? ndims(a) : 1), X)
+ d_max = max(ndimsA)
+
+ cat_ranges = ntuple(nargs, i->(catdim <= ndimsA[i] ? dimsA[i][catdim] : 1))
+
+ function compute_dims(d)
+ if d == catdim
+ if catdim <= d_max
+ return sum(cat_ranges)
+ else
+ return nargs
+ end
+ else
+ if d <= ndimsA[1]
+ return dimsA[1][d]
+ else
+ return 1
+ end
+ end
  end
- C = similar(C, typeC, dimsC)
 
- for i=1:nargs
- C[i] = X[i]
+ ndimsC = max(catdim, d_max)
+ dimsC = ntuple(ndimsC, compute_dims)::(Size...)
+ typeC = promote_type(ntuple(nargs, i->isa(X[i],AbstractArray) ? typeof(X[i]).parameters[1] : typeof(X[i]))...)
+ C = similar(isa(X[1],AbstractArray) ? X[1] : [X[1]], typeC, dimsC)
+
+ cat_ranges = cumsum(1, cat_ranges...)::(Size...)
+ for k=1:nargs
+ cat_one = ntuple(ndimsC, i->(i != catdim ?
+ Range1(1,dimsC[i]) :
+ Range1(cat_ranges[k],cat_ranges[k+1]-1) ))
+ C[cat_one...] = X[k]
  end
  return C
 end
@@ -668,7 +688,7 @@ function cat(catdim::Int, A::AbstractArray...)
  return nargs
  end
  else
- if d <= d_max
+ if d <= ndimsA[1]
  return dimsA[1][d]
  else
  return 1
@@ -1328,19 +1348,18 @@ type SubArray{T,N,A<:AbstractArray,I<:(RangeIndex...,)} <: AbstractArray{T,N}
  dims::Dims
  strides::Array{Index,1}
  first_index::Index
+ indexed_dims::Array{Index,1}
 
  #linear indexing constructor
  if N == 1 && length(I) == 1 && A <: Array
- function SubArray(p::A, i::I)
- t = new(p, i, (length(i[1]),))
- if isa(i[1], Index)
- t.strides = [1]
- t.first_index = i[1]
- else
- t.strides = [isa(i[1], Range1) ? 1 : i[1].step]
- t.first_index = i[1].start
- end
- t
+ function SubArray(p::A, i::(Index,))
+ new(p, i, (length(i[1]),), [1], i[1])
+ end
+ function SubArray(p::A, i::(Range1{Index},))
+ new(p, i, (length(i[1]),), [1], i[1].start)
+ end
+ function SubArray(p::A, i::(Range{Index},))
+ new(p, i, (length(i[1]),), i[1].step, i[1].start) 
  end
  else
  function SubArray(p::A, i::I)
@@ -1410,7 +1429,7 @@ function slice(A::SubArray, i::RangeIndex...)
 end
 
 ### rename the old slice function ###
-##slice all dimensions of length 1
+##squeeze all dimensions of length 1
 #slice{T,N}(a::AbstractArray{T,N}) = sub(a, map(i-> i == 1 ? 1 : (1:i), size(a)))
 #slice{T,N}(s::SubArray{T,N}) =
 # sub(s.parent, map(i->!isa(i, Index) && length(i)==1 ?i[1] : i, s.indexes))
@@ -1476,6 +1495,35 @@ function ref(s::SubArray, is::Int...)
  s.parent[index]
 end
 
+function ref(s::SubArray, I::Indices...)
+ j = 1 #the jth dimension in subarray
+ n = ndims(s.parent)
+ newindexes = Array(Indices, n)
+ for i = 1:n
+ t = s.indexes[i]
+ #TODO: don't generate the dense vector indexes if they can be ranges
+ newindexes[i] = isa(t, Index) ? t : t[I[j]]
+ j += 1
+ end
+
+ reshape(ref(s.parent, newindexes...), s.dims)
+end
+#TODO: don't generate the dense vector indexes if they can be ranges
+#function ref{T}(s::SubArray{T,1}, I::Range{Index})
+
+#end
+#function ref{T}(s::SubArray{T,1}, I::Range1{Index})
+
+#end
+#TODO: see if using strides and first index is faster for 1d
+function ref2{T}(s::SubArray{T,1}, I::AbstractArray{Int})
+ t = Array(Index, length(I))
+ for i = 1:length(I)
+ t[i] = s.first_index + (I[i]-1)*s.strides
+ end
+ ref(s.parent, t)
+end
+
 assign(s::SubArray, v::AbstractArray, i::Int) =
  invoke(assign, (SubArray, Any, Int), s, v, i)
 assign(s::SubArray, v, i::Int) = assign(s, v, ind2sub(size(s), i)...)