fast path for common_chunks where only one disk array (#224)

* fast path for common_chunks where only one disk array

* bugfix tt

* fix if block nesting

* add chunk equality comparisons

* fix == for small RegularChunks

Author: rafaqz

src/broadcast.jl

@@ -52,19 +52,24 @@ function eachchunk(a::BroadcastDiskArray)
 end
 function common_chunks(s, args...)
     N = length(s)
-    chunkedars = filter(i -> haschunks(i) === Chunked(), collect(args))
-    all(ar -> isa(eachchunk(ar), GridChunks), chunkedars) ||
+    chunkedarrays = reduce(args; init=()) do acc, x
+        haschunks(x) === Chunked() ? (acc..., x) : acc
+    end
+    all(ar -> isa(eachchunk(ar), GridChunks), chunkedarrays) ||
         error("Currently only chunks of type GridChunks can be merged by broadcast")
-    if isempty(chunkedars)
+    if isempty(chunkedarrays)
         totalsize = sum(sizeof ∘ eltype, args)
         return estimate_chunksize(s, totalsize)
+    elseif length(chunkedarrays) == 1
+        return eachchunk(only(chunkedarrays))
     else
-        allcs = eachchunk.(chunkedars)
+        allchunks = collect(map(eachchunk, chunkedarrays))
         tt = ntuple(N) do n
-            csnow = filter(allcs) do cs
-                ndims(cs) >= n && first(first(cs.chunks[n])) < last(last(cs.chunks[n]))
+            csnow = filter(allchunks) do cs
+            ndims(cs) >= n && first(first(cs.chunks[n])) < last(last(cs.chunks[n]))
             end
             isempty(csnow) && return RegularChunks(1, 0, s[n])
+            
             cs = first(csnow).chunks[n]
             if all(s -> s.chunks[n] == cs, csnow)
                 return cs

src/chunks.jl

@@ -38,6 +38,28 @@ function Base.getindex(r::RegularChunks, i::Int)
 end
 Base.size(r::RegularChunks, _) = div(r.s + r.offset - 1, r.cs) + 1
 Base.size(r::RegularChunks) = (size(r, 1),)
+function Base.:(==)(r1::RegularChunks, r2::RegularChunks)
+    # The axis sizes must always match
+    r1.s == r2.s || return false
+    # The number of chunks must also match
+    nchunks = length(r1)
+    nchunks == length(r2) || return false
+    # But after that we need to take the number of chunks into account
+    if nchunks > 2 
+        # For longer RegularChunks the offsets and chunk sizes 
+        # must match for the chunks to be the same. 
+        # So we compare them directly rather than iterating all of the ranges
+        return r1.cs == r2.cs && r1.offset == r2.offset
+    elseif nchunks == 2
+        # Smaller RegularChunks can match with different chunk sizes and offsets
+        # So we compare the ranges
+        return first(r1) == first(r2) && last(r1) == last(r2)
+    elseif nchunks == 1
+        return first(r1) == first(r2)
+    else
+        return true
+    end
+end
 
 # DiskArrays interface
 
@@ -135,6 +157,9 @@ function Base.getindex(r::IrregularChunks, i::Int)
     return (r.offsets[i] + 1):r.offsets[i + 1]
 end
 Base.size(r::IrregularChunks) = (length(r.offsets) - 1,)
+Base.:(==)(r1::IrregularChunks, r2::IrregularChunks) =
+    r1 === r2 || r1.offsets == r2.offsets
+
 function subsetchunks(r::IrregularChunks, subs::UnitRange)
     c1 = findchunk(r, first(subs))
     c2 = findchunk(r, last(subs))