Restore inferrable return type for block tensor slicing

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockTensorKit"
 uuid = "5f87ffc2-9cf1-4a46-8172-465d160bd8cd"
-version = "0.3.16"
+version = "0.3.17"
 authors = ["Lukas Devos <ldevos98@gmail.com> and contributors"]
 
 [deps]

src/tensors/abstractblocktensor/abstractarray.jl

@@ -100,20 +100,33 @@ end
         t::AbstractBlockTensorMap, indices::Vararg{Any, M}
     ) where {M}
     M == numind(t) || return _slice_getindex_single(t, indices...)
-    inds = Base.to_indices(t, indices)
-    inds isa NTuple{M, Int} && return parent(t)[inds...]
+    return _slice_getindex_full(t, Base.to_indices(t, indices))
+end
+
+# single block vs. sliced block tensor is chosen by dispatch, so the return type is inferrable
+@propagate_inbounds _slice_getindex_full(
+    t::AbstractBlockTensorMap, inds::NTuple{M, Int}
+) where {M} = parent(t)[inds...]
+@propagate_inbounds function _slice_getindex_full(t::AbstractBlockTensorMap, inds::Tuple)
     tdst = similar(t, space(eachspace(t)[inds...]))
     length(tdst) == 0 && return tdst
     return _copyslice!(tdst, t, inds)
 end
 
+# a mask may consume several dimensions, which `to_indices` cannot infer through the fillers below
+@propagate_inbounds _slice_getindex_single(
+    t::AbstractBlockTensorMap, index::AbstractArray{Bool}
+) = _slice_getindex_single(t, Base.to_index(t, index))
+
 # a single index is only supported when it selects a single nontrivial dimension
 @noinline function _slice_getindex_single(
         t::AbstractBlockTensorMap, indices::Vararg{Any, M}
     ) where {M}
     space(eachspace(t)[indices...]) # errors as before if unsupported
     d = something(findfirst(>(1), size(t)), 1)
-    return _slice_getindex(t, ntuple(i -> i == d ? only(indices) : 1, numind(t))...)
+    # `1:1` rather than `1` keeps `Int` out of the tuple, so the dispatch above stays static
+    inds = ntuple(i -> i == d ? only(indices) : (1:1), numind(t))
+    return _slice_getindex_full(t, Base.to_indices(t, inds))
 end
 
 # TODO: check if this fallback is fair

test/abstracttensor/indexing.jl

@@ -92,6 +92,45 @@ for (label, t1) in (("dense", rand(V ← one(V))), ("sparse", sprand(V ← one(V
     end
 end
 
+# single-index slicing of a multi-index block tensor takes a different path than the case above
+W = SumSpace(ℂ^2)
+for (label, ts) in (("dense", rand(W ⊗ V ← W)), ("sparse", sprand(W ⊗ V ← W, 0.8)))
+    @testset "$label single-index slicing" begin
+        @test size(ts) == (1, 3, 1)
+        mid = 2:(length(ts) - 1)
+        # a single index must stay as inferrable as the equivalent full-rank slice
+        @test @inferred(ts[mid]) == ts[:, mid, :]
+        @test space(ts[mid]) == space(ts[:, mid, :])
+        @test @inferred(ts[1]) isa TensorMap
+        @test @inferred(ts[:]) == ts
+        @test size(@inferred(ts[[1, 3]])) == (1, 2, 1)
+        @test ts[[1, 3]][1, 2, 1] == ts[1, 3, 1]
+        mask = [true, false, true]
+        @test @inferred(ts[mask]) == ts[:, mask, :]
+    end
+end
+
+# a widened return type propagates into everything built from a slice, so keep every form concrete
+infers_concretely(f, args...) =
+    isconcretetype(Core.Compiler.return_type(f, Tuple{map(typeof, args)...}))
+
+@testset "indexing inference ($label)" for (label, t) in (
+        ("dense", rand(W ⊗ V ← W)), ("sparse", sprand(W ⊗ V ← W, 0.8)),
+        ("dense square", rand(V ⊗ V ← V)), ("sparse square", sprand(V ⊗ V ← V, 0.5)),
+        ("dense vector", rand(V ← one(V))), ("sparse vector", sprand(V ← one(V), 0.8)),
+    )
+    for ind in (1, 2:3, 1:2:3, :, [1, 3], [true, false, true])
+        @test infers_concretely(getindex, t, ind)
+    end
+    ndims(t) == 3 || continue
+    for inds in (
+            (1, 2, 1), (:, :, :), (1, 2:3, 1), (1, [1, 3], 1), (1, [true, false, true], 1),
+            (1:1, 1:2:3, 1:1), ([1, 1, 3], :, :),
+        )
+        @test infers_concretely(getindex, t, inds...)
+    end
+end
+
 # the parent array has its own slicing implementation
 @testset "parent array slicing" begin
     st = sprand(V ⊗ V ⊗ V, 0.5)