From 9cca1371e0dd9bac566ce93c1ba24d7f45269025 Mon Sep 17 00:00:00 2001
From: "Viktor G. Gumennyy" <20798349+sairus7@users.noreply.github.com>
Date: Tue, 23 May 2023 20:43:16 +0300
Subject: [PATCH] Add macro for subtypes lowering with additional `subtypekey`
 field (#94)

* add macro `register_struct_subtype`

* version bump

* Update Project.toml

* fix non-exported name

---------

Co-authored-by: Jacob Quinn <quinn.jacobd@gmail.com>
---
 Project.toml     |  2 +-
 src/macros.jl    | 32 ++++++++++++++++++++++++++++++++
 test/runtests.jl | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 4c15ed4..9264413 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "StructTypes"
 uuid = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 authors = ["Jacob Quinn"]
-version = "1.10.1"
+version = "1.11.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
diff --git a/src/macros.jl b/src/macros.jl
index 6bcf38e..6acfece 100644
--- a/src/macros.jl
+++ b/src/macros.jl
@@ -55,4 +55,36 @@ function macro_constructor(expr::Symbol, structtype)
             throw(ArgumentError("StructType macros can only be used with a struct definition or a Type"))
         end
     end)
+end
+
+"""
+Macro to add subtypes for an abstract type without the need for type field.
+For a given `abstract_type`` and `struct_subtype` generates custom lowered NameTuple 
+with all subtype fields and additional `StructTypes.subtypekey` field
+used for identifying the appropriate concrete subtype.
+
+usage:
+```julia
+abstract type Vehicle end
+struct Car <: Vehicle
+    make::String
+end
+StructTypes.subtypekey(::Type{Vehicle}) = :type
+StructTypes.subtypes(::Type{Vehicle}) = (car=Car, truck=Truck)
+StructTypes.@register_struct_subtype Vehicle Car
+```
+"""
+macro register_struct_subtype(abstract_type, struct_subtype)
+    AT = Core.eval(__module__, abstract_type)
+    T = Core.eval(__module__, struct_subtype)
+    field_name = StructTypes.subtypekey(AT)
+    field_value = findfirst(x->x === T, StructTypes.subtypes(AT))
+    x_names = [:(x.$(e)) for e in fieldnames(T)] # x.a, x.b, ...
+    name_types = [:($n::$(esc(t))) for (n, t) in zip(fieldnames(T), fieldtypes(T))] # a::Int, b::String, ...
+    quote
+        StructTypes.StructType(::Type{$(esc(struct_subtype))}) = StructTypes.CustomStruct()
+        StructTypes.lower(x::$(esc(struct_subtype))) = ($(esc(field_name)) = $(QuoteNode(field_value)), $(x_names...))
+        StructTypes.lowertype(::Type{$(esc(struct_subtype))}) = @NamedTuple{$field_name::$(esc(Symbol)), $(name_types...)}
+        $(esc(struct_subtype))(x::@NamedTuple{$field_name::$(esc(Symbol)), $(name_types...)}) = $(esc(struct_subtype))($(x_names...))
+    end
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index a9c2b5b..4d992a7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -824,3 +824,35 @@ StructTypes.defaults(::Type{Defaultable}) = (b="b",)
     @test StructTypes.constructfrom(Defaultable, Dict(:a=>"a")) == Defaultable("a", "b")
     @test StructTypes.constructfrom(Defaultable, (a="a", b="c")) == Defaultable("a", "c")
 end
+
+# CustomStruct subtypes
+abstract type Vehicle2 end
+
+struct Car2 <: Vehicle2
+    make::String
+    model::String
+    seatingCapacity::Int
+    topSpeed::Float64
+end
+
+struct Truck2 <: Vehicle2
+    # type::String no need of this
+    make::String
+    model::String
+    payloadCapacity::Float64
+end
+
+StructTypes.StructType(::Type{Vehicle2}) = StructTypes.AbstractType()
+StructTypes.subtypekey(::Type{Vehicle2}) = :type
+StructTypes.subtypes(::Type{Vehicle2}) = (car=Car2, truck=Truck2)
+StructTypes.@register_struct_subtype Vehicle2 Car2
+StructTypes.@register_struct_subtype Vehicle2 Truck2
+
+@testset "register_struct_subtype" begin
+    nt = (; :type => :car, :make => "Mercedes-Benz", :model => "S500", :seatingCapacity => 5, :topSpeed => 250.1)
+    car = Car2(nt)
+    @test StructTypes.lower(car) === nt
+    @test StructTypes.lowertype(Car2) === typeof(nt)
+    @test typeof(car) == Car2
+    @test car.make == "Mercedes-Benz"
+end
\ No newline at end of file