add more array types

Author: tmigot

Project.toml

@@ -11,6 +11,7 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

src/code.jl

@@ -1,6 +1,6 @@
 module Code
 
-using StaticArrays, LabelledArrays, SparseArrays
+using StaticArrays, LabelledArrays, SparseArrays, LinearAlgebra
 
 export toexpr, Assignment, (←), Let, Func, DestructuredArgs, LiteralExpr,
        SetArray, MakeArray, MakeSparseArray, MakeTuple, AtIndex,
@@ -375,7 +375,7 @@ function toexpr(a::MakeArray, st)
 end
 
 ## Array
-@inline function _create_array(::Union{Type{<:Array},Type{<:SubArray}}, T, ::Val{dims}, elems...) where dims
+@inline function _create_array(::Type{<:Array}, T, ::Val{dims}, elems...) where dims
     arr = Array{T}(undef, dims)
     @assert prod(dims) == nfields(elems)
     @inbounds for i=1:prod(dims)
@@ -384,15 +384,23 @@ end
     arr
 end
 
-@inline function create_array(A::Union{Type{<:Array},Type{<:SubArray}}, T, d::Val, elems...)
+@inline function create_array(A::Type{<:Array}, T, d::Val, elems...)
     _create_array(A, T, d, elems...)
 end
 
-@inline function create_array(A::Union{Type{<:Array},Type{<:SubArray}}, ::Nothing, d::Val{dims}, elems...) where dims
+@inline function create_array(A::Type{<:Array}, ::Nothing, d::Val{dims}, elems...) where dims
     T = promote_type(map(typeof, elems)...)
     _create_array(A, T, d, elems...)
 end
 
+@inline function create_array(A::Type{<:SubArray{T,N,P,I,L}}, S, d::Val, elems...) where {T,N,P,I,L}
+    create_array(P, S, d, elems...)
+end
+
+@inline function create_array(A::Type{<:PermutedDimsArray{T,N,perm,iperm,P}}, S, d::Val, elems...) where {T,N,perm,iperm,P}
+    create_array(P, S, d, elems...)
+end
+
 ## Matrix
 
 @inline function create_array(::Type{<:Matrix}, ::Nothing, ::Val{dims}, elems...) where dims
@@ -403,6 +411,14 @@ end
     Base.typed_hvcat(T, dims, elems...)
 end
 
+@inline function create_array(A::Type{<:Transpose{T,P}}, S, d::Val, elems...) where {T,P}
+    create_array(P, S, d, elems...)
+end
+
+@inline function create_array(A::Type{<:UpperTriangular{T,P}}, S, d::Val, elems...) where {T,P}
+    create_array(P, S, d, elems...)
+end
+
 ## SArray
 @inline function create_array(::Type{<:SArray}, ::Nothing, ::Val{dims}, elems...) where dims
     SArray{Tuple{dims...}}(elems...)

test/code.jl

@@ -4,6 +4,7 @@ using SymbolicUtils.Code: LazyState
 using StaticArrays
 using LabelledArrays
 using SparseArrays
+using LinearAlgebra
 
 test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linenums!(b))
 
@@ -92,6 +93,15 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
     @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
                           MakeArray(view([a,b,a+b,a/b], :), arr)))) == [1, 2, 3, 1/2]
 
+    @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
+                          MakeArray(PermutedDimsArray([a b;a+b a/b], (1,2)), arr)))) == [1 2 ; 3  1/2]
+
+    @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
+                          MakeArray(transpose([a b;a+b a/b]), arr)))) == [1 3;2 1/2]
+
+    @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
+                          MakeArray(UpperTriangular([a b;a+b a/b]), arr)))) == [1 2;0 1/2]
+
     @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
                           MakeArray([a b;a+b a/b], arr)))) == [1 2; 3 1/2]