Merge pull request #735 from AayushSabharwal/as/cse-fix

fix: tuples and arrays of `Symbols` should not be considered symbolic in codegen

Author: AayushSabharwal

src/code.jl

@@ -176,22 +176,29 @@ function substitute_name(O, st)
     end
 end
 
+function _is_tuple_or_array_of_symbolics(O)
+    return O isa CodegenPrimitive ||
+        (symbolic_type(O) != NotSymbolic() && !(O isa Union{Symbol, Expr})) ||
+        _is_array_of_symbolics(O) ||
+        _is_tuple_of_symbolics(O)
+end
+
 function _is_array_of_symbolics(O)
     # O is an array, not a symbolic array, and either has a non-symbolic eltype or contains elements that are
     # symbolic or arrays of symbolics
     return O isa AbstractArray && symbolic_type(O) == NotSymbolic() &&
-        (symbolic_type(eltype(O)) != NotSymbolic() ||
-        any(x -> x isa CodegenPrimitive || symbolic_type(x) != NotSymbolic() || _is_array_of_symbolics(x), O))
+        (symbolic_type(eltype(O)) != NotSymbolic() && !(eltype(O) <: Union{Symbol, Expr}) ||
+        any(_is_tuple_or_array_of_symbolics, O))
 end
 
 # workaround for https://github.com/JuliaSparse/SparseArrays.jl/issues/599
 function _is_array_of_symbolics(O::SparseMatrixCSC)
-    return symbolic_type(eltype(O)) != NotSymbolic() ||
-        any(x -> x isa CodegenPrimitive || symbolic_type(x) != NotSymbolic() || _is_array_of_symbolics(x), findnz(O)[3])
+    return symbolic_type(eltype(O)) != NotSymbolic() && !(eltype(O) <: Union{Symbol, Expr}) ||
+        any(_is_tuple_or_array_of_symbolics, findnz(O)[3])
 end
 
 function _is_tuple_of_symbolics(O::Tuple)
-    return any(x -> x isa CodegenPrimitive || symbolic_type(x) != NotSymbolic() || _is_array_of_symbolics(x) || _is_tuple_of_symbolics(x), O)
+    return any(_is_tuple_or_array_of_symbolics, O)
 end
 _is_tuple_of_symbolics(O) = false
 

test/cse.jl

@@ -229,3 +229,22 @@ end
     fn2 = @RuntimeGeneratedFunction(cse(toexpr(fnexpr)))
     @test fn1(ones(5)) == fn2(ones(5))
 end
+
+@testset "Tuples and arrays of `Symbol`s aren't symbolic" begin
+    @syms x y
+    f(a, b, c) = a + b + length(c)
+    ex = term(f, x, y, (:a, :b, :c))
+    expr = quote
+        let x = 1, y = 2
+            $(toexpr(cse(ex)))
+        end
+    end
+    @test eval(expr) == 6
+    ex = term(f, x, y, [:a, :b, :c])
+    expr = quote
+        let x = 1, y = 2
+            $(toexpr(cse(ex)))
+        end
+    end
+    @test eval(expr) == 6
+end