From 9e24b8199be62128ea00cc4c8d9de78985d20f53 Mon Sep 17 00:00:00 2001
From: Andrew Kille <68079167+apkille@users.noreply.github.com>
Date: Sat, 3 Aug 2024 21:32:11 -0400
Subject: [PATCH] Clean up metadata addition to struct definitions (#70)

* withmetadata update
---
 CHANGELOG.md                         |  4 ++
 Project.toml                         |  4 +-
 src/QSymbolicsBase/QSymbolicsBase.jl | 57 ++++++++--------------------
 test/Project.toml                    |  1 +
 4 files changed, 24 insertions(+), 42 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f5ef5cb..7d71bd9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # News
 
+## v0.3.5 - dev
+
+- Cleaned up metadata decoration of struct definitions.
+  
 ## v0.3.4 - 2024-07-22
 
 - Added `tr` and `ptrace` functionalities.
diff --git a/Project.toml b/Project.toml
index afd08a8..1ed77ea 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,11 +1,12 @@
 name = "QuantumSymbolics"
 uuid = "efa7fd63-0460-4890-beb7-be1bbdfbaeae"
 authors = ["QuantumSymbolics.jl contributors"]
-version = "0.3.4"
+version = "0.3.5-dev"
 
 [deps]
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
@@ -24,6 +25,7 @@ QuantumOpticsExt = "QuantumOpticsBase"
 [compat]
 Latexify = "0.16"
 LinearAlgebra = "1.9"
+MacroTools = "0.5.13"
 PrecompileTools = "1.2"
 QuantumClifford = "0.8.19, 0.9"
 QuantumInterface = "0.3.3"
diff --git a/src/QSymbolicsBase/QSymbolicsBase.jl b/src/QSymbolicsBase/QSymbolicsBase.jl
index 57bae11..985d02e 100644
--- a/src/QSymbolicsBase/QSymbolicsBase.jl
+++ b/src/QSymbolicsBase/QSymbolicsBase.jl
@@ -4,6 +4,7 @@ using SymbolicUtils
 import SymbolicUtils: Symbolic,_isone,flatten_term,isnotflat,Chain,Fixpoint,Prewalk,sorted_arguments
 using TermInterface
 import TermInterface: isexpr,head,iscall,children,operation,arguments,metadata,maketerm
+import MacroTools: namify, @capture
 
 using LinearAlgebra
 import LinearAlgebra: eigvecs,ishermitian,conj,transpose,inv,exp,vec,tr
@@ -66,52 +67,26 @@ Metadata() = Metadata(CacheType())
 
 """Decorate a struct definition in order to add a metadata dict which would be storing cached `express` results."""
 macro withmetadata(strct)
-    withmetadata(strct)
-end
-function withmetadata(strct) # TODO this should really use MacroTools instead of this mess
-    struct_name = strct.args[2]
-    constructor = :($struct_name() = new())
-    if struct_name isa Expr # if Struct{T<:QObj} <: Symbolic{T}
-        struct_name = struct_name.args[1]
-        constructor = :($struct_name() = new())
-        if struct_name isa Expr # if Struct{T<:QObj}
-            struct_name = struct_name.args[1] # now it is just Struct
-            constructor = :($struct_name{S}() where S = new{S}())
-        end
+    ex = quote $strct end
+    if @capture(ex, (struct T_{params__} fields__ end) | (struct T_{params__} <: A_ fields__ end))
+        struct_name = namify(T)
+        args = (namify(i) for i in fields)
+        constructor = :($struct_name{S}($(args...)) where S = new{S}($((args..., :(Metadata()))...)))
+    elseif @capture(ex, struct T_ fields__ end)
+        struct_name = namify(T)
+        args = (namify(i) for i in fields)
+        constructor = :($struct_name($(args...)) = new($((args..., :(Metadata()))...)))
+    else @capture(ex, struct T_ end)
+        struct_name = namify(T)
+        constructor = :($struct_name() = new($:(Metadata())))
     end
     struct_args = strct.args[end].args
-    if all(x->x isa Symbol || x isa LineNumberNode || x isa String || x.head==:(::), struct_args)
-        # add constructor
-        args = [x for x in struct_args if x isa Symbol || x isa Expr] # the arguments required for the constructor
-        args = [a isa Symbol ? a : (a.head==:(::) ? a.args[1] : a) for a in args] # drop typeasserts
-        declaring_line = constructor.args[1] # :(Constructor{}()) or :(Constructor{}() where {})
-        if declaring_line.head == :where
-            declaring_line = declaring_line.args[1]
-        end
-        append!(declaring_line.args, args) # Adding them to the line declaring the constructor, i.e. adding them at the location of ? in `Constructor(?) = new(...)`
-        new_call_args = constructor.args[end].args[end].args # The ? in `new(?)`
-        append!(new_call_args, args) # Adding them to the `new` call
-        push!(new_call_args, :(Metadata()))
-        push!(struct_args, constructor)
-    else
-        # modify constructor
-        newwithmetadata.(struct_args)
-    end
-    # add metadata slot
-    push!(struct_args, :(metadata::Metadata))
+    push!(struct_args, constructor, :(metadata::Metadata))
     esc(quote
-        Base.@__doc__ $strct
-        metadata(x::$struct_name)=x.metadata
+    Base.@__doc__ $strct
+    metadata(x::$struct_name)=x.metadata
     end)
 end
-function newwithmetadata(expr::Expr)
-    if expr.head==:call && (expr.args[1]==:new || expr.args[1]==:(new{S}))
-        push!(expr.args, :(Metadata()))
-    else
-        newwithmetadata.(expr.args)
-    end
-end
-newwithmetadata(x) = x
 
 ##
 # Basic Types
diff --git a/test/Project.toml b/test/Project.toml
index 449d3c3..c3a674b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -6,6 +6,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 QuantumClifford = "0525e862-1e90-11e9-3e4d-1b39d7109de1"
 QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5"
 QuantumOptics = "6e0679c1-51ea-5a7c-ac74-d61b76210b0c"