Improve type inference for adjoint/transpose of nested matrices (#1181)

* Improve type inference for adjoint/transpose of nested matrices

* Add comment

Author: jishnub

src/linalg.jl

@@ -62,6 +62,15 @@ end
 #--------------------------------------------------
 # Matrix algebra
 
+# _adjointtype returns the eltype of the container when computing the adjoint/transpose
+# of a static array. Using this method instead of calling `Base.promote_op` directly
+# helps with type-inference, particularly for nested static arrays,
+# where the adjoint is applied recursively.
+@inline _adjointtype(f, ::Type{T}) where {T} = Base.promote_op(f, T)
+for S in (:SMatrix, :MMatrix)
+    @eval @inline _adjointtype(f, ::Type{$S{M,N,T,L}}) where {M,N,T,L} = $S{N,M,_adjointtype(f, T),L}
+end
+
 # Transpose, etc
 @inline transpose(m::StaticMatrix) = _transpose(Size(m), m)
 # note: transpose of StaticVector is a Transpose, handled by Base
@@ -74,7 +83,7 @@ end
     return quote
         $(Expr(:meta, :inline))
         elements = tuple($(exprs...))
-        @inbounds return similar_type($m, Base.promote_op(transpose, T), Size($(n2,n1)))(elements)
+        @inbounds return similar_type($m, _adjointtype(transpose, T), Size($(n2,n1)))(elements)
     end
 end
 
@@ -88,7 +97,7 @@ end
     return quote
         $(Expr(:meta, :inline))
         elements = tuple($(exprs...))
-        @inbounds return similar_type($m, Base.promote_op(adjoint, T), Size($(n2,n1)))(elements)
+        @inbounds return similar_type($m, _adjointtype(adjoint, T), Size($(n2,n1)))(elements)
     end
 end
 

test/linalg.jl

@@ -237,6 +237,15 @@ end
         end
         @test adjoint(SMatrix{0,0,Vector{Int}}()) isa SMatrix{0,0,Adjoint{Int,Vector{Int}}}
         @test transpose(SMatrix{0,0,Vector{Int}}()) isa SMatrix{0,0,Transpose{Int,Vector{Int}}}
+
+        @testset "inference for nested matrices" begin
+            A = reshape([reshape([complex(i,2i)*j for i in 1:2], 1, 2) for j in 1:6], 3, 2)
+            for TA in (SMatrix, MMatrix), TB in (SMatrix, MMatrix)
+                S = TA{3,2}(TB{1,2}.(A)) # static matrix of static matrices
+                @test @inferred(transpose(S)) == transpose(A)
+                @test @inferred(adjoint(S)) == adjoint(A)
+            end
+        end
     end
 
     @testset "normalization" begin