nits from JET (#1095)

* nits from JET

* add tests for non-`:generator` comprehension

Author: thchr

src/SArray.jl

@@ -159,10 +159,13 @@ function static_array_gen(::Type{SA}, @nospecialize(ex), mod::Module) where {SA}
         args = parse_cat_ast(ex)
         return :($SA{$Tuple{$(size(args)...)}}($tuple($(escall(args)...))))
     elseif head === :comprehension
-        if length(ex.args) != 1 || !isa(ex.args[1], Expr) || ex.args[1].head != :generator
+        if length(ex.args) != 1
             error("Expected generator in comprehension, e.g. [f(i,j) for i = 1:3, j = 1:3]")
         end
         ex = ex.args[1]
+        if !isa(ex, Expr) || (ex::Expr).head != :generator
+            error("Expected generator in comprehension, e.g. [f(i,j) for i = 1:3, j = 1:3]")
+        end
         n_rng = length(ex.args) - 1
         rng_args = (ex.args[i+1].args[1] for i = 1:n_rng)
         rngs = Any[Core.eval(mod, ex.args[i+1].args[2]) for i = 1:n_rng]
@@ -174,11 +177,14 @@ function static_array_gen(::Type{SA}, @nospecialize(ex), mod::Module) where {SA}
             end
         end
     elseif head === :typed_comprehension
-        if length(ex.args) != 2 || !isa(ex.args[2], Expr) || ex.args[2].head != :generator
+        if length(ex.args) != 2
             error("Expected generator in typed comprehension, e.g. Float64[f(i,j) for i = 1:3, j = 1:3]")
         end
         T = esc(ex.args[1])
         ex = ex.args[2]
+        if !isa(ex, Expr) || (ex::Expr).head != :generator
+            error("Expected generator in typed comprehension, e.g. Float64[f(i,j) for i = 1:3, j = 1:3]")
+        end
         n_rng = length(ex.args) - 1
         rng_args = (ex.args[i+1].args[1] for i = 1:n_rng)
         rngs = Any[Core.eval(mod, ex.args[i+1].args[2]) for i = 1:n_rng]

src/SVector.jl

@@ -39,10 +39,13 @@ function static_vector_gen(::Type{SV}, @nospecialize(ex), mod::Module) where {SV
         len = check_vector_length(size(args))
         return :($SV{$len}($tuple($(escall(args)...))))
     elseif head === :comprehension
-        if length(ex.args) != 1 || !isa(ex.args[1], Expr) || ex.args[1].head != :generator
+        if length(ex.args) != 1
             error("Expected generator in comprehension, e.g. [f(i) for i = 1:3]")
         end
         ex = ex.args[1]
+        if !isa(ex, Expr) || (ex::Expr).head != :generator
+            error("Expected generator in comprehension, e.g. [f(i) for i = 1:3]")
+        end
         if length(ex.args) != 2
             error("Use a one-dimensional comprehension for @$SV")
         end
@@ -55,11 +58,14 @@ function static_vector_gen(::Type{SV}, @nospecialize(ex), mod::Module) where {SV
             end
         end
     elseif head === :typed_comprehension
-        if length(ex.args) != 2 || !isa(ex.args[2], Expr) || ex.args[2].head != :generator
+        if length(ex.args) != 2
             error("Expected generator in typed comprehension, e.g. Float64[f(i) for i = 1:3]")
         end
         T = esc(ex.args[1])
         ex = ex.args[2]
+        if !isa(ex, Expr) || (ex::Expr).head != :generator
+            error("Expected generator in typed comprehension, e.g. Float64[f(i) for i = 1:3]")
+        end
         if length(ex.args) != 2
             error("Use a one-dimensional comprehension for @$SV")
         end

src/matrix_multiply.jl

@@ -128,7 +128,6 @@ for TWR in [Adjoint, Transpose, Symmetric, Hermitian, LowerTriangular, UpperTria
 end
 
 @generated function _mul(Sa::Size{sa}, Sb::Size{sb}, a::StaticMatMulLike{<:Any, <:Any, Ta}, b::StaticMatMulLike{<:Any, <:Any, Tb}) where {sa, sb, Ta, Tb}
-    S = Size(sa[1], sb[2])
     # Heuristic choice for amount of codegen
     a_tri_mul = a <: LinearAlgebra.AbstractTriangular ? 4 : 1
     b_tri_mul = b <: LinearAlgebra.AbstractTriangular ? 4 : 1

test/SArray.jl

@@ -91,6 +91,10 @@
         test_expand_error(:(@SArray fill()))
         test_expand_error(:(@SArray [1; 2; 3; 4]...))
 
+        # (typed-)comprehension LoadError for `ex.args[1].head != :generator`
+        test_expand_error(:(@SArray [i+j for i in 1:2 for j in 1:2]))
+        test_expand_error(:(@SArray Int[i+j for i in 1:2 for j in 1:2]))
+
         @test ((@SArray fill(1))::SArray{Tuple{},Int}).data === (1,)
         @test ((@SArray ones())::SArray{Tuple{},Float64}).data === (1.,)