Specialize blockedrange BroadcastStyle for LazyArrayStyle (#384)

* Specialize blockedrange BroadcastStyle for LazyArrayStyle

* Add compat for LazyArrays

Author: jishnub

Project.toml

@@ -5,14 +5,22 @@ version = "1.0.0-dev"
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
+[weakdeps]
+LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
+
+[extensions]
+BlockArraysLazyArraysExt = "LazyArrays"
+
 [compat]
 Aqua = "0.8"
 ArrayLayouts = "1.0.8"
 Documenter = "1"
 FillArrays = "1"
 InfiniteArrays = "0.13"
+LazyArrays = "1"
 LinearAlgebra = "1.6"
 OffsetArrays = "1"
 Random = "1.6"

ext/BlockArraysLazyArraysExt.jl

@@ -0,0 +1,8 @@
+module BlockArraysLazyArraysExt
+
+import BlockArrays
+import LazyArrays
+
+BlockArrays._broadcaststyle(S::LazyArrays.LazyArrayStyle{1}) = S
+
+end

src/BlockArrays.jl

@@ -74,6 +74,8 @@ include("blockarrayinterface.jl")
 @deprecate getblock!(X, A::AbstractBlockArray{T,N}, I::Vararg{Integer, N}) where {T,N} copyto!(X, view(A, Block(I)))
 @deprecate setblock!(A::AbstractBlockArray{T,N}, v, I::Vararg{Integer, N}) where {T,N} (A[Block(I...)] = v)
 
-
+if !isdefined(Base, :get_extension)
+    include("../ext/BlockArraysLazyArraysExt.jl")
+end
 
 end # module

src/blockaxis.jl

@@ -535,7 +535,8 @@ blockaxes(S::Base.Slice) = blockaxes(S.indices)
 
 
 # This supports broadcasting with infinite block arrays
-Base.BroadcastStyle(::Type{<:AbstractBlockedUnitRange{<:Any,R}}) where R = Base.BroadcastStyle(R)
+_broadcaststyle(_) = Broadcast.DefaultArrayStyle{1}()
+Base.BroadcastStyle(::Type{<:AbstractBlockedUnitRange{<:Any,R}}) where R = _broadcaststyle(Base.BroadcastStyle(R))
 
 
 ###

test/test_blockbroadcast.jl

@@ -1,6 +1,7 @@
 using BlockArrays, FillArrays, Test
 import BlockArrays: SubBlockIterator, BlockIndexRange, Diagonal
 import InfiniteArrays
+using StaticArrays
 
 @testset "broadcast" begin
     @testset "BlockArray" begin
@@ -129,6 +130,9 @@ import InfiniteArrays
         for i in 1:10
             @test b2[Block(i)] == b[Block(i)] + b[Block(i)]
         end
+
+        b = blockedrange(SVector{2}([1,2]))
+        @test b .+ b == 2:2:6
     end
 
     @testset "Special broadcast" begin