Preserve blocked unit range in more broadcasting operations (#464)

Fixes #334.

Co-authored-by: Sheehan Olver <[email protected]>

Author: mtfishman

src/blockbroadcast.jl

@@ -276,3 +276,13 @@ for op in (:*, :\)
         broadcasted(::AbstractBlockStyle, ::typeof($op), a::AbstractArray{T}, b::Ones{V}) where {T,V} = LinearAlgebra.copy_oftype(a, Base.promote_op(*, T, V))
     end
 end
+
+
+
+###
+# Ranges
+###
+
+broadcasted(::DefaultArrayStyle{1}, ::typeof(+), r::AbstractBlockedUnitRange, x::Integer) = _BlockedUnitRange(first(r) + x, blocklasts(r) .+ x)
+broadcasted(::DefaultArrayStyle{1}, ::typeof(+), x::Integer, r::AbstractBlockedUnitRange) = _BlockedUnitRange(x + first(r), x .+ blocklasts(r))
+broadcasted(::DefaultArrayStyle{1}, ::typeof(-), r::AbstractBlockedUnitRange, x::Integer) = _BlockedUnitRange(first(r) - x, blocklasts(r) .- x)

test/test_blockbroadcast.jl

@@ -193,6 +193,24 @@ using StaticArrays
         @test blockisequal(axes(A .* Ones(axes(A))), axes(Ones(axes(A)) .* A), axes(A .* ones(6)))
     end
 
+    @testset "Blocked unit range broadcast" begin
+        r = blockedrange([2,3])
+        @test r .+ 3 == blockedrange(4, [2,3])
+        @test blocklengths(r .+ 3) == [2,3]
+        @test 3 .+ r == blockedrange(4, [2,3])
+        @test blocklengths(3 .+ r .+ 3) == [2,3]
+        @test r .- 3 == blockedrange(-2, [2,3])
+        @test blocklengths(r .- 3) == [2,3]
+
+        r = blockedrange(2, [2,3])
+        @test r .+ 3 == blockedrange(5, [2,3])
+        @test blocklengths(r .+ 3) == [2,3]
+        @test 3 .+ r == blockedrange(5, [2,3])
+        @test blocklengths(3 .+ r .+ 3) == [2,3]
+        @test r .- 3 == blockedrange(-1, [2,3])
+        @test blocklengths(r .- 3) == [2,3]
+    end
+
     @testset "type inference" begin
         u = BlockArray(randn(5), [2,3]);
         A = zeros(size(u))