Fix dimensions for row view of a vcat (#388)

* Clearer sublayout for concat

* special case for row view of Vcat

* Add a test

* Case where k is not in args range

* v2.9.2

---------

Co-authored-by: Daniel VandenHeuvel <[email protected]>

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "2.9.1"
+version = "2.9.2"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/lazyconcat.jl

@@ -791,10 +791,16 @@ end
 # subarrays
 ###
 
-sublayout(::ApplyLayout{typeof(vcat)}, ::Type{<:Tuple{Vararg{Union{AbstractRange{Int},Int}}}}) = ApplyLayout{typeof(vcat)}()
-sublayout(::ApplyLayout{typeof(hcat)}, ::Type{<:Tuple{Vararg{Union{AbstractRange{Int},Int}}}}) = ApplyLayout{typeof(hcat)}()
+sublayout(::ApplyLayout{typeof(vcat)}, ::Type{<:Tuple{AbstractRange{Int}}}) = ApplyLayout{typeof(vcat)}()
+sublayout(::ApplyLayout{typeof(vcat)}, ::Type{<:Tuple{AbstractRange{Int},Int}}) = ApplyLayout{typeof(vcat)}()
+sublayout(::ApplyLayout{typeof(vcat)}, ::Type{<:Tuple{AbstractRange{Int},AbstractRange{Int}}}) = ApplyLayout{typeof(vcat)}()
+sublayout(::ApplyLayout{typeof(vcat)}, ::Type{<:Tuple{Int,AbstractRange{Int}}}) = ApplyLayout{typeof(vcat)}() # a Vcat but with only 1 argument
+
+
 # a row-slice of an Hcat is equivalent to a Vcat
 sublayout(::ApplyLayout{typeof(hcat)}, ::Type{<:Tuple{Int,AbstractRange{Int}}}) = ApplyLayout{typeof(vcat)}()
+sublayout(::ApplyLayout{typeof(hcat)}, ::Type{<:Tuple{AbstractRange{Int},Int}}) = ApplyLayout{typeof(hcat)}()
+sublayout(::ApplyLayout{typeof(hcat)}, ::Type{<:Tuple{AbstractRange{Int},AbstractRange{Int}}}) = ApplyLayout{typeof(hcat)}()
 
 _vcat_lastinds(sz) = _vcat_cumsum(sz...)
 _vcat_firstinds(sz) = (1, (1 .+ Base.front(_vcat_lastinds(sz)))...)
@@ -816,13 +822,25 @@ function _vcat_sub_arguments(lay::ApplyLayout{typeof(vcat)}, A, V, kr)
     _reverse_if_neg_step(map(_view_vcat, arguments(lay, A), skr2), kr)
 end
 
-function _vcat_sub_arguments(::ApplyLayout{typeof(vcat)}, A, V, kr, jr)
-    sz = size.(arguments(A),1)
+function _vcat_sub_arguments(lay::ApplyLayout{typeof(vcat)}, A, V, kr, jr)
+    sz = size.(arguments(lay, A),1)
     skr = intersect.(_argsindices(sz), Ref(kr))
     skr2 = broadcast((a,b) -> a .- b .+ 1, skr, _vcat_firstinds(sz))
     _reverse_if_neg_step(_view_vcat.(arguments(A), skr2, Ref(jr)), kr)
 end
 
+function _vcat_sub_arguments(lay::ApplyLayout{typeof(vcat)}, A, V, k::Int, jr)
+    args = arguments(lay, A)
+    sz = size.(args,1)
+    for (a,inds, firstind) in zip(args, _argsindices(sz), _vcat_firstinds(sz))
+        if k in inds
+            return (_view_vcat(a, k - firstind + 1, jr),)
+        end
+    end
+
+    () # no args
+end
+
 _vcat_sub_arguments(LAY::ApplyLayout{typeof(vcat)}, A, V) = _vcat_sub_arguments(LAY, A, V, parentindices(V)...)
 
 

test/concattests.jl

@@ -698,6 +698,18 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, materialize!, call, paddeddat
         @test copy(view(V,2:-1:1,1:-1:1)) == [2 ; 1 ;;]
         @test copy(view(H,1:-1:1,2:-1:1)) == [2 1]
     end
+
+    @testset "Vcat arguments with integer-range view" begin
+        A = Vcat([1, 2, 3, 4, 5]', [6, 7, 8, 9, 10]')
+        V = view(A, 1, 1:3)
+        args = arguments(V)
+        @test length(args) == 1 
+        @test args[1] == [1, 2, 3] 
+        @test args[1] isa SubArray{Int, 1, Vector{Int}}
+
+        args = LazyArrays._vcat_sub_arguments(MemoryLayout(V), V, (), 0, 1:3)
+        @test args == ()
+    end
 end
 
 end # module