fix HCAT+Sum domain problem

Author: hakkelt

src/calculus/HCAT.jl

@@ -85,8 +85,23 @@ function HCAT(A::Vararg{AbstractOperator})
     return HCAT(AA, buf)
 end
 
-# compile-time domain ndoms for HCAT's sub-operators
-_ndoms_from_type(::Type{<:HCAT{N}}, dim::Int) where {N} = dim == 2 ? N : 1
+# Count actual domain slots from an HCAT's P (idxs) type.
+# Each entry in P is either an Int (1 slot) or a NTuple{n,Int} (n slots).
+_count_hcat_ndoms(::Type{<:Tuple{}}) = 0
+@generated function _count_hcat_ndoms(::Type{P}) where {P <: Tuple}
+    K = 0
+    for i in 1:fieldcount(P)
+        Pi = fieldtype(P, i)
+        K += Pi <: Integer ? 1 : fieldcount(Pi)
+    end
+    return :($K)
+end
+
+# compile-time domain ndoms for HCAT's sub-operators:
+# use the index-tuple type P (not N which only counts sub-operators) so that
+# sub-operators with multi-component domains are accounted for correctly.
+_ndoms_from_type(::Type{<:HCAT{N, L, P}}, dim::Int) where {N, L, P} =
+    dim == 2 ? _count_hcat_ndoms(P) : 1
 
 @generated function HCAT(AA::NTuple{N, AbstractOperator}, buf::C) where {N, C}
     N == 1 && return :(AA[1])

src/calculus/Sum.jl

@@ -162,6 +162,12 @@ codomain_storage_type(S::Sum{K, C, D, L}) where {K, C <: Tuple, D, L} = codomain
 fun_domain(S::Sum) = fun_domain(S.A[1])
 fun_codomain(S::Sum) = fun_codomain(S.A[1])
 
+# A Sum of multi-domain operators (e.g. Sum of HCATs) has the same domain arity
+# as its first constituent.  This extends the HCAT machinery so that HCAT can
+# correctly assign indices when a Sum appears as one of its sub-operators.
+_ndoms_from_type(::Type{<:Sum{K, C, D, L}}, dim::Int) where {K, C, D, L} =
+    _ndoms_from_type(fieldtype(L, 1), dim)
+
 fun_name(S::Sum) = length(S.A) == 2 ? fun_name(S.A[1]) * "+" * fun_name(S.A[2]) : "Σ"
 
 is_linear(L::Sum) = all(is_linear.(L.A))