add SubblockIterator

Author: lkdvos

src/tensors/blockiterator.jl

@@ -76,3 +76,69 @@ function Base.show(io::IO, mime::MIME"text/plain", b::BlockIterator)
     show_blocks(io, mime, b)
     return nothing
 end
+
+"""
+    struct SubblockIterator{T <: AbstractTensorMap, S}
+
+Iterator over the subblocks of a tensor of type `T`, possibly holding some pre-computed data of type `S`.
+This is typically constructed through of [`subblocks`](@ref).
+"""
+struct SubblockIterator{T <: AbstractTensorMap, S}
+    t::T
+    structure::S
+end
+
+Base.IteratorSize(::SubblockIterator) = Base.HasLength()
+Base.IteratorEltype(::SubblockIterator) = Base.HasEltype()
+Base.eltype(::Type{<:SubblockIterator{T}}) where {T} = Pair{fusiontreetype(T), subblocktype(T)}
+Base.length(iter::SubblockIterator) = length(iter.structure)
+Base.isdone(iter::SubblockIterator, state...) = Base.isdone(iter.structure, state...)
+
+# default implementation assumes `structure = fusiontrees(t)`
+function Base.iterate(iter::SubblockIterator, state...)
+    next = Base.iterate(iter.structure, state...)
+    isnothing(next) && return nothing
+    (f₁, f₂), state = next
+    @inbounds data = subblock(iter.t, (f₁, f₂))
+    return (f₁, f₂) => data, state
+end
+
+
+function Base.showarg(io::IO, iter::SubblockIterator, toplevel::Bool)
+    print(io, "subblocks(")
+    Base.showarg(io, iter.t, false)
+    print(io, ")")
+    return nothing
+end
+function Base.summary(io::IO, iter::SubblockIterator)
+    Base.showarg(io, iter, true)
+    return nothing
+end
+
+function show_subblocks(io::IO, mime::MIME"text/plain", iter::SubblockIterator)
+    if FusionStyle(sectortype(iter.t)) isa UniqueFusion
+        first = true
+        for ((f₁, f₂), b) in iter
+            first || print(io, "\n\n")
+            print(io, " * ", f₁.uncoupled, " ← ", f₂.uncoupled, " => ")
+            show(io, mime, b)
+            first = false
+        end
+    else
+        first = true
+        for ((f₁, f₂), b) in iter
+            first || print(io, "\n\n")
+            print(io, " * ", (f₁, f₂), " => ")
+            show(io, mime, b)
+            first = false
+        end
+    end
+    return nothing
+end
+
+function Base.show(io::IO, mime::MIME"text/plain", iter::SubblockIterator)
+    summary(io, iter)
+    println(io, ":")
+    show_subblocks(io, mime, iter)
+    return nothing
+end