From 548d5ce665430fc4476455eb6e0b9ce827e5a895 Mon Sep 17 00:00:00 2001
From: timholy <tim.holy@gmail.com>
Date: Sun, 21 Jul 2013 09:12:56 -0500
Subject: [PATCH] Merging range indexes for faster subarray linear indexing

This only helps when the indexes are commensurate, but when that's true it's a huge benefit
---
 base/subarray.jl | 79 +++++++++++++++++++++++++++++++++++++++++++++++-
 test/arrayops.jl | 19 ++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/base/subarray.jl b/base/subarray.jl
index 9cc546b84f8a4..5595c665b3654 100644
--- a/base/subarray.jl
+++ b/base/subarray.jl
@@ -215,7 +215,13 @@ function translate_indexes(s::SubArray, I::Union(Real,AbstractArray)...)
     lastdim = pdims[n]
     if havelinear
         newindexes = newindexes[1:lastdim]
-        newindexes[pdims[n]] = translate_linear_indexes(s, n, I[end], pdims)
+        pstrides = strides(s.parent)[lastdim:end]
+        pindexes = s.indexes[lastdim:end]
+        if ismergeable(pstrides, pindexes)
+            newindexes[lastdim] = mergeindexes(pstrides, pindexes)[I[end]]
+        else
+            newindexes[lastdim] = translate_linear_indexes(s, n, I[end], pdims)
+        end
     end
     newindexes
 end
@@ -272,6 +278,77 @@ function parentdims(s::SubArray)
     dimindex
 end
 
+# True if indexing by a list of RangeIndexes is equivalent to indexing with a single Range
+function ismergeable(strides, indexes::(RangeIndex...))
+    m = true
+    n = length(indexes)
+    i = 1
+    local I
+    # Find the first non-singleton index
+    while i <= n
+        I = indexes[i]
+        i += 1
+        if length(I) > 1
+            break
+        end
+    end
+    if i > n
+        return true
+    end
+    # Check that strides in all later ranges are commensurate
+    pold = step(I)*strides[i-1]*length(I)
+    while i <= n
+        I = indexes[i]
+        i += 1
+        lI = length(I)
+        if lI == 1
+            continue
+        end
+        p = step(I)*strides[i-1]
+        m = pold == p
+        if !m
+            break
+        end
+        pold = p*lI
+    end
+    m
+end
+ismergeable(strides, indexes::RangeIndex...) = ismergeable(strides, indexes)
+
+function mergeindexes(strides::Union(Dims,AbstractVector), indexes::(RangeIndex...))
+    N = 1
+    f = 1
+    n = length(indexes)
+    i = 1
+    local I
+    # Find the first non-singleton index
+    while i <= n
+        I = indexes[i]
+        f += (first(I)-1)*strides[i]
+        i += 1
+        if length(I) > 1
+            break
+        end
+    end
+    if i > length(strides)
+        return f:f
+    end
+    N = length(I)
+    stp = step(I)*strides[i-1]/strides[1]
+    while i <= n
+        I = indexes[i]
+        f += (first(I)-1)*strides[i]
+        lI = length(I)
+        i += 1
+        if lI == 1
+            continue
+        end
+        N *= lI
+    end
+    return stp == 1 ? Range1{Int}(f, N) : Range{Int}(f, stp, N)   # is this a good idea or should it always be Range?
+end
+mergeindexes(strides::Union(Dims,AbstractVector), indexes::RangeIndex...) = mergeindexes(strides, indexes)
+
 function getindex(s::SubArray, I::Union(Real,AbstractVector)...)
     newindexes = translate_indexes(s, I...)
 
diff --git a/test/arrayops.jl b/test/arrayops.jl
index cb2be430ca34c..f22d0cc39fc8c 100644
--- a/test/arrayops.jl
+++ b/test/arrayops.jl
@@ -90,6 +90,25 @@ b = [4, 6, 2, -7, 1]
 ind = findin(a, b)
 @test ind == [3,4]
 
+# merging multidimensional range indexes into a single index
+strds = (1,8,40)
+@test Base.ismergeable(strds, 1:8, 1:5, 1:3) == true           # all elements
+@test Base.ismergeable(strds, 8:-1:1, 5:-1:1, 3:-1:1) == true  # all elements in reverse order
+@test Base.ismergeable(strds, 8:-1:1, 1:5, 1:3) == false
+@test Base.ismergeable(strds, 1, 1, 1) == true
+@test Base.ismergeable(strds, 1:8, 2:4, 3) == true
+@test Base.ismergeable(strds, 2:2:8, 2:4, 3) == true
+@test Base.ismergeable(strds, 1:8, 2:5, 1:2) == false
+@test Base.ismergeable(strds, 2, 3:5) == true
+
+@test Base.mergeindexes(strds, 1:8, 1:5, 1:3) == 1:120
+@test Base.mergeindexes(strds, 8:-1:1, 5:-1:1, 3:-1:1) == 120:-1:1
+@test Base.mergeindexes(strds, 1, 1, 1) == 1:1
+@test Base.mergeindexes(strds, 2, 3, 1) == 18:18
+@test Base.mergeindexes(strds, 1:8, 2:4, 3) == 89:112
+@test Base.mergeindexes(strds, 2:2:8, 2:4, 3) == 90:2:112
+@test Base.mergeindexes(strds, 2, 3:5) == 18:8:40
+
 # sub
 A = reshape(1:120, 3, 5, 8)
 sA = sub(A, 2, 1:5, 1:8)