Refactor create_result for Tuple types to include NamedTuples (#11

)

* Refactor `create_result` for `Tuple` types to include `NamedTuple`s

* Fix `NamedTuple` codegen in `create_result`

* Add test for `create_result` on `NamedTuple`

src/Reactant.jl

@@ -771,9 +771,9 @@ function generate_jlfunc(
         end
         if T <: Tuple
             elems = Symbol[]
-            for (i, v) in enumerate(tocopy)
-                sym = Symbol(string(resname) * "_" * string(i))
-                create_result(v, sym, (path..., i))
+            for (k, v) in pairs(tocopy)
+                sym = Symbol(resname, :_, k)
+                create_result(v, sym, (path..., k))
                 push!(elems, sym)
             end
             push!(
@@ -784,6 +784,18 @@ function generate_jlfunc(
             )
             return nothing
         end
+        if T <: NamedTuple
+            elems = Symbol[]
+            for (k, v) in pairs(tocopy)
+                sym = Symbol(resname, :_, k)
+                create_result(v, sym, (path..., k))
+                push!(elems, sym)
+            end
+            push!(concrete_result_maker, quote
+                    $resname = NamedTuple{$(keys(tocopy))}($elems)
+            end)
+            return
+        end
         if T <: Array
             elems = Symbol[]
             for (i, v) in enumerate(tocopy)

test/compile.jl

@@ -0,0 +1,20 @@
+using Reactant
+using Test
+
+@testset "compile" begin
+    @testset "create_result" begin
+        @testset "NamedTuple" begin
+            const MockNamedTuple{T} = @NamedTuple{a::T}
+            MockNamedTuple(x::T) where {T} = MockNamedTuple{T}((x,))
+
+            x = MockNamedTuple(rand(4, 3))
+            x2 = MockNamedTuple(Reactant.ConcreteRArray(x.a))
+
+            Base.sum(x::MockNamedTuple) = MockNamedTuple((sum(x.a),))
+            f = Reactant.compile(sum, (x2,))
+
+            @test f(x2) isa MockNamedTuple{Reactant.ConcreteRArray}
+            @test isapprox(f(x2).a, sum(x).a)
+        end
+    end
+end

test/runtests.jl

@@ -44,3 +44,4 @@ include("basic.jl")
 include("bcast.jl")
 include("nn.jl")
 include("struct.jl")
+include("compile.jl")