fix create_result on Enum (#2835)

* add `create_result` method for enums

* remove legacy `testset` wrapper

* test

Author: mofeing

src/Compiler.jl

@@ -277,6 +277,37 @@ Base.@nospecializeinfer function create_result(
     return result_cache[tocopy]
 end
 
+Base.@nospecializeinfer function create_result(
+    @nospecialize(tocopy::Enum),
+    @nospecialize(path::Tuple),
+    result_stores,
+    path_to_shard_info,
+    to_unreshard_results,
+    _unresharded_code::Vector{Expr},
+    _unresharded_arrays_cache,
+    used_shardinfo,
+    result_cache,
+    var_idx,
+    resultgen_code,
+)
+    if !haskey(result_cache, tocopy)
+        sym = Symbol("result", var_idx[])
+        var_idx[] += 1
+
+        result = Meta.quot(tocopy)
+
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
+        result_cache[tocopy] = sym
+    end
+
+    return result_cache[tocopy]
+end
+
 function create_result(
     tocopy::ConcretePJRTNumber{T,D},
     @nospecialize(path::Tuple),

test/core/compile.jl

@@ -13,76 +13,85 @@ intout_caller(vis) = @noinline intout(vis)
     @test_throws MethodError @compile intout_caller(visr)
 end
 
-@testset "compile" begin
-    @testset "create_result" begin
-        @testset "NamedTuple" begin
-            x = (; a=Reactant.TestUtils.construct_test_array(Float64, 4, 3))
-            x2 = Reactant.to_rarray(x)
-
-            res = @jit sum(x2)
-            @test res isa NamedTuple
-            @test res.a isa ConcreteRNumber{Float64}
-            @test isapprox(res.a, sum(x.a))
-        end
-
-        @testset "Array" begin
-            x = [1 2; 3 4; 5 6]
-            f = Reactant.compile(() -> x, ())
-            @test f() ≈ x
-        end
+@testset "create_result" begin
+    @testset "NamedTuple" begin
+        x = (; a=Reactant.TestUtils.construct_test_array(Float64, 4, 3))
+        x2 = Reactant.to_rarray(x)
+
+        res = @jit sum(x2)
+        @test res isa NamedTuple
+        @test res.a isa ConcreteRNumber{Float64}
+        @test isapprox(res.a, sum(x.a))
     end
 
-    @testset "world-age" begin
-        a = ones(2, 10)
-        b = ones(10, 2)
-        a_ra = Reactant.to_rarray(a)
-        b_ra = Reactant.to_rarray(b)
+    @testset "Array" begin
+        x = [1 2; 3 4; 5 6]
+        f = Reactant.compile(() -> x, ())
+        @test f() ≈ x
+    end
 
-        fworld(x, y) = @jit(x * y)
+    @testset "Enum" begin
+        @enum MyEnum begin
+            MyEnumA = 1
+            MyEnumB = 2
+        end
 
-        @test fworld(a_ra, b_ra) ≈ ones(2, 2) * 10
+        x = MyEnumA
+        f = @compile identity(x)
+        @test f(x) == x
     end
+end
 
-    @testset "type casting & optimized out returns" begin
-        a = ones(2, 10)
-        a_ra = Reactant.to_rarray(a)
+@testset "world-age" begin
+    a = ones(2, 10)
+    b = ones(10, 2)
+    a_ra = Reactant.to_rarray(a)
+    b_ra = Reactant.to_rarray(b)
 
-        ftype1(x) = Float64.(x)
-        ftype2(x) = Float32.(x)
+    fworld(x, y) = @jit(x * y)
 
-        y1 = @jit ftype1(a_ra)
-        y2 = @jit ftype2(a_ra)
+    @test fworld(a_ra, b_ra) ≈ ones(2, 2) * 10
+end
 
-        @test y1 isa Reactant.ConcreteRArray{Float64,2}
-        @test y2 isa Reactant.ConcreteRArray{Float32,2}
+@testset "type casting & optimized out returns" begin
+    a = ones(2, 10)
+    a_ra = Reactant.to_rarray(a)
 
-        @test y1 ≈ Float64.(a)
-        @test y2 ≈ Float32.(a)
-    end
+    ftype1(x) = Float64.(x)
+    ftype2(x) = Float32.(x)
 
-    @testset "no variable name collisions in compile macros (#237)" begin
-        f(x) = x
-        g(x) = f(x)
-        x = Reactant.TestUtils.construct_test_array(Float64, 2, 2)
-        y = Reactant.to_rarray(x)
-        @test (@jit g(y); true)
-    end
+    y1 = @jit ftype1(a_ra)
+    y2 = @jit ftype2(a_ra)
 
-    # disabled due to long test time (core tests go from 2m to 7m just with this test)
-    # @testset "resource exhaustation bug (#190)" begin
-    #     x = rand(2, 2)
-    #     y = Reactant.to_rarray(x)
-    #     @test try
-    #         for _ in 1:10_000
-    #             f = @compile sum(y)
-    #         end
-    #         true
-    #     catch e
-    #         false
-    #     end
-    # end
+    @test y1 isa Reactant.ConcreteRArray{Float64,2}
+    @test y2 isa Reactant.ConcreteRArray{Float32,2}
+
+    @test y1 ≈ Float64.(a)
+    @test y2 ≈ Float32.(a)
 end
 
+@testset "no variable name collisions in compile macros (#237)" begin
+    f(x) = x
+    g(x) = f(x)
+    x = Reactant.TestUtils.construct_test_array(Float64, 2, 2)
+    y = Reactant.to_rarray(x)
+    @test (@jit g(y); true)
+end
+
+# disabled due to long test time (core tests go from 2m to 7m just with this test)
+# @testset "resource exhaustation bug (#190)" begin
+#     x = rand(2, 2)
+#     y = Reactant.to_rarray(x)
+#     @test try
+#         for _ in 1:10_000
+#             f = @compile sum(y)
+#         end
+#         true
+#     catch e
+#         false
+#     end
+# end
+
 @testset "Module export" begin
     f(x) = sin.(cos.(x))
     x_ra = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float64, 3))