From 544dc819862be46a72e932cdb1a71d18906b1537 Mon Sep 17 00:00:00 2001
From: Ole Poeth <olepoeth@gmail.com>
Date: Wed, 26 Mar 2025 10:07:13 +0100
Subject: [PATCH] Add specification field and parsing in add_rule!

---
 src/csg/csg.jl      |  8 +++++---
 src/grammar_base.jl | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/csg/csg.jl b/src/csg/csg.jl
index c10cae5..1afbc3e 100644
--- a/src/csg/csg.jl
+++ b/src/csg/csg.jl
@@ -34,6 +34,7 @@ mutable struct ContextSensitiveGrammar <: AbstractGrammar
 	bychildtypes::Vector{BitVector}
 	log_probabilities::Union{Vector{Real}, Nothing}
 	constraints::Vector{AbstractConstraint}
+	specification::Vector{Vector{Expr}}
 end
 
 ContextSensitiveGrammar(
@@ -45,10 +46,11 @@ ContextSensitiveGrammar(
 	domains::Dict{Symbol, BitVector},
 	childtypes::Vector{Vector{Symbol}},
 	bychildtypes::Vector{BitVector},
-	log_probabilities::Union{Vector{<:Real}, Nothing}
-) = ContextSensitiveGrammar(rules, types, isterminal, iseval, bytype, domains, childtypes, bychildtypes, log_probabilities, AbstractConstraint[])
+	log_probabilities::Union{Vector{<:Real}, Nothing},
+	specification::Vector{Vector{Expr}}
+) = ContextSensitiveGrammar(rules, types, isterminal, iseval, bytype, domains, childtypes, bychildtypes, log_probabilities, AbstractConstraint[], specification)
 
-ContextSensitiveGrammar() = ContextSensitiveGrammar([], [], BitVector[], BitVector[], Dict{Symbol, Vector{Int}}(), Dict{Symbol, BitVector}(), Vector{Vector{Symbol}}(), Vector{BitVector}(), nothing, AbstractConstraint[])
+ContextSensitiveGrammar() = ContextSensitiveGrammar([], [], BitVector[], BitVector[], Dict{Symbol, Vector{Int}}(), Dict{Symbol, BitVector}(), Vector{Vector{Symbol}}(), Vector{BitVector}(), nothing, AbstractConstraint[], Vector{Vector{Expr}}())
 
 """
 	expr2csgrammar(ex::Expr)::ContextSensitiveGrammar
diff --git a/src/grammar_base.jl b/src/grammar_base.jl
index c86cba6..8a8ab1d 100644
--- a/src/grammar_base.jl
+++ b/src/grammar_base.jl
@@ -165,6 +165,20 @@ function add_rule!(g::AbstractGrammar, e::Expr)
             # with strict equality so that true and 1 are not considered
             # equal. that means we can't use `in` or `∈` for equality checking.
             if !any(s == type && (r === rule || typeof(r)==Expr && r == rule) for (type, rule) ∈ zip(g.types, g.rules))
+                # Extract specification if defined.
+                if (isa(r, Expr) && r.head == :(:=))
+                    specification = r.args[2]
+                    if isa(specification, Expr) && specification.head == :tuple
+                        specification = specification.args
+                    else
+                        specification = [specification]
+                    end
+                    push!(g.specification, specification)
+                    r = r.args[1]
+                else
+                    push!(g.specification, [])
+                end
+                
                 push!(g.rules, r)
                 push!(g.iseval, iseval(rule))
                 push!(g.types, s)