conjure_cp_rule_macros/
lib.rs

1mod util;
2
3use itertools::Itertools;
4use proc_macro::TokenStream;
5
6use proc_macro2::Span;
7use quote::{ToTokens, quote};
8use std::collections::HashMap;
9use syn::token::Comma;
10use syn::{
11    ExprClosure, GenericParam, Ident, ItemFn, ItemImpl, ItemStruct, LitInt, LitStr, Result, Token,
12    TypePath, bracketed, parenthesized, parse::Parse, parse::ParseStream, parse_macro_input,
13};
14
15use crate::util::{rename_ident_in_impl, type_is_ident, type_is_option_of_ident};
16use util::{build_serde_as_type, rename_fn, rename_ident_in_fn, type_contains_ident};
17
18/// Parsed arguments for `#[register_rule(...)]`.
19///
20/// The attribute syntax is:
21///
22/// ```text
23/// #[register_rule("RuleSet", priority)]
24/// #[register_rule(["RuleSetA", "RuleSetB"], priority)]
25/// #[register_rule("RuleSet", priority, [prefilter, ...])]
26/// #[register_rule("RuleSet", priority, [prefilter, ...], symbols_only)]
27/// ```
28///
29/// The optional prefilter list narrows where the scheduler attempts the rule:
30///
31/// ```text
32/// Variant              // focused expression must have this Expression variant
33/// * / Variant          // focused expression must have an immediate child with this variant
34/// Variant / Variant    // focused expression must have the left variant and an immediate child
35///                      // with the right variant
36/// Atomic / Reference   // focused expression must be an atomic reference
37/// Atomic / Literal     // focused expression must be an atomic literal
38/// * / Atomic / Reference  // focused expression must have an immediate atomic-reference child
39/// * / Atomic / Literal    // focused expression must have an immediate atomic-literal child
40/// ```
41///
42/// Items in the list are alternatives. For example, `[And / Comprehension, Or / Comprehension]`
43/// means `(And with direct Comprehension child) OR (Or with direct Comprehension child)`.
44///
45/// Omitting the prefilter list makes the rule universal. A universal rule is considered at every
46/// focused expression for its priority. In Rust source, write slash filters with spaces around the
47/// slash, e.g. `[* / Bubble]` or `[Atomic / Reference]`, to keep the syntax visually distinct.
48enum ParsedPrefilter {
49    /// Focused-expression variant prefilter, e.g. `Add` or `Sub`.
50    Variant(Ident),
51    /// Immediate-child variant prefilter, parsed from `* / Child`.
52    Child { child: Ident },
53    /// Paired focused-expression and immediate-child variant prefilter, parsed from `Root / Child`.
54    VariantChild { variant: Ident, child: Ident },
55    /// Atomic subvariant prefilter, parsed from `Atomic / Reference` or `Atomic / Literal`.
56    Atom(Ident),
57    /// Immediate-child atomic subvariant prefilter, parsed from `* / Atomic / Reference`.
58    ChildAtom(Ident),
59}
60
61struct RegisterRuleArgs {
62    /// Rule set names this rule belongs to.
63    rule_sets: Vec<LitStr>,
64    /// Priority used for every listed rule set.
65    priority: LitInt,
66    /// Complete prefilter alternatives. Empty means the rule is universal.
67    prefilters: Vec<ParsedPrefilter>,
68    /// Whether a failed application can only be invalidated by a symbol-table change.
69    symbols_only: bool,
70}
71
72impl Parse for RegisterRuleArgs {
73    fn parse(input: ParseStream) -> Result<Self> {
74        if input.is_empty() {
75            return Ok(RegisterRuleArgs {
76                rule_sets: Vec::new(),
77                priority: LitInt::new("0", Span::call_site()),
78                prefilters: Vec::new(),
79                symbols_only: false,
80            });
81        }
82
83        let rule_sets = if input.peek(syn::token::Bracket) {
84            let content;
85            bracketed!(content in input);
86
87            let mut rule_sets = Vec::new();
88            while !content.is_empty() {
89                let rule_set: LitStr = content.parse()?;
90                rule_sets.push(rule_set);
91                if content.is_empty() {
92                    break;
93                }
94                let _: Comma = content.parse()?;
95            }
96            rule_sets
97        } else {
98            vec![input.parse()?]
99        };
100
101        let _: Comma = input.parse()?;
102        let priority: LitInt = input.parse()?;
103
104        // Parse optional variant names in brackets: "Minion", 4200, [Add, Sub]
105        let mut prefilters = Vec::new();
106        if input.peek(Comma) {
107            let _: Comma = input.parse()?;
108            let content;
109            bracketed!(content in input);
110            while !content.is_empty() {
111                if content.peek(Token![*]) {
112                    let _: Token![*] = content.parse()?;
113                    let _: Token![/] = content.parse()?;
114                    let variant: Ident = content.parse()?;
115                    if variant == "Atomic" && content.peek(Token![/]) {
116                        let _: Token![/] = content.parse()?;
117                        let subvariant: Ident = content.parse()?;
118                        match subvariant.to_string().as_str() {
119                            "Literal" | "Reference" => {
120                                prefilters.push(ParsedPrefilter::ChildAtom(subvariant));
121                            }
122                            _ => {
123                                return Err(syn::Error::new(
124                                    subvariant.span(),
125                                    "expected Literal or Reference after * / Atomic /",
126                                ));
127                            }
128                        }
129                    } else {
130                        prefilters.push(ParsedPrefilter::Child { child: variant });
131                    }
132                } else {
133                    let variant: Ident = content.parse()?;
134                    if content.peek(Token![/]) {
135                        let _: Token![/] = content.parse()?;
136                        let subvariant: Ident = content.parse()?;
137                        if variant != "Atomic" {
138                            prefilters.push(ParsedPrefilter::VariantChild {
139                                variant,
140                                child: subvariant,
141                            });
142                        } else {
143                            match subvariant.to_string().as_str() {
144                                "Literal" | "Reference" => {
145                                    prefilters.push(ParsedPrefilter::Atom(subvariant));
146                                }
147                                _ => {
148                                    return Err(syn::Error::new(
149                                        subvariant.span(),
150                                        "expected Literal or Reference after Atomic /",
151                                    ));
152                                }
153                            }
154                        }
155                    } else {
156                        prefilters.push(ParsedPrefilter::Variant(variant));
157                    }
158                }
159                if content.is_empty() {
160                    break;
161                }
162                let _: Comma = content.parse()?;
163            }
164        }
165
166        let symbols_only = if input.peek(Comma) {
167            let _: Comma = input.parse()?;
168            let dependency: Ident = input.parse()?;
169            if dependency != "symbols_only" {
170                return Err(syn::Error::new(
171                    dependency.span(),
172                    "expected `symbols_only`",
173                ));
174            }
175            true
176        } else {
177            false
178        };
179
180        Ok(RegisterRuleArgs {
181            rule_sets,
182            priority,
183            prefilters,
184            symbols_only,
185        })
186    }
187}
188
189/// Register a rule with the given rule sets and priorities.
190#[proc_macro_attribute]
191pub fn register_rule(arg_tokens: TokenStream, item: TokenStream) -> TokenStream {
192    let func = parse_macro_input!(item as ItemFn);
193    let rule_ident = &func.sig.ident;
194    let static_name = format!("CONJURE_GEN_RULE_{rule_ident}").to_uppercase();
195    let static_ident = Ident::new(&static_name, rule_ident.span());
196
197    let args = parse_macro_input!(arg_tokens as RegisterRuleArgs);
198
199    let rule_sets_token = if args.rule_sets.is_empty() {
200        quote! { &[] }
201    } else {
202        let rule_sets = &args.rule_sets;
203        let priority = &args.priority;
204        quote! { &[#((#rule_sets, #priority as u16)),*] }
205    };
206
207    let prefilters = if args.prefilters.is_empty() {
208        quote! { None }
209    } else {
210        let prefilter_tokens = args.prefilters.iter().map(|prefilter| match prefilter {
211            ParsedPrefilter::Variant(variant) => {
212                quote! {
213                    ::conjure_cp::rule_engine::RulePrefilter::Variant(
214                        ::conjure_cp::discriminant_from_name!(#variant)
215                    )
216                }
217            }
218            ParsedPrefilter::Child { child } => {
219                quote! {
220                    ::conjure_cp::rule_engine::RulePrefilter::Child {
221                        child: ::conjure_cp::discriminant_from_name!(#child),
222                    }
223                }
224            }
225            ParsedPrefilter::VariantChild { variant, child } => {
226                quote! {
227                    ::conjure_cp::rule_engine::RulePrefilter::VariantChild {
228                        variant: ::conjure_cp::discriminant_from_name!(#variant),
229                        child: ::conjure_cp::discriminant_from_name!(#child),
230                    }
231                }
232            }
233            ParsedPrefilter::Atom(atom_variant) => {
234                quote! {
235                    ::conjure_cp::rule_engine::RulePrefilter::Atom(
236                        ::conjure_cp::rule_engine::AtomKind::#atom_variant
237                    )
238                }
239            }
240            ParsedPrefilter::ChildAtom(atom_variant) => {
241                quote! {
242                    ::conjure_cp::rule_engine::RulePrefilter::ChildAtom(
243                        ::conjure_cp::rule_engine::AtomKind::#atom_variant
244                    )
245                }
246            }
247        });
248        quote! {
249            Some(&[#(#prefilter_tokens),*])
250        }
251    };
252    let failure_invalidation = if args.symbols_only {
253        quote! { ::conjure_cp::rule_engine::RuleFailureInvalidation::SymbolsOnly }
254    } else {
255        quote! { ::conjure_cp::rule_engine::RuleFailureInvalidation::ExpressionOrSymbols }
256    };
257
258    let expanded = quote! {
259        #func
260
261        pub static #static_ident: ::conjure_cp::rule_engine::Rule<'static> = ::conjure_cp::rule_engine::Rule {
262            name: stringify!(#rule_ident),
263            application: #rule_ident,
264            rule_sets: #rule_sets_token,
265            prefilters: #prefilters,
266            failure_invalidation: #failure_invalidation,
267        };
268
269        ::conjure_cp::rule_engine::_dependencies::inventory::submit! {
270            &#static_ident
271        }
272    };
273
274    TokenStream::from(expanded)
275}
276// this is only constructed once at comptime so this doesn't matter
277#[allow(clippy::large_enum_variant)]
278enum ReprStateType {
279    Struct(ItemStruct, Vec<ItemImpl>),
280    Path(TypePath),
281}
282
283impl Parse for ReprStateType {
284    fn parse(input: ParseStream) -> Result<Self> {
285        if let Ok(strct) = input.parse::<ItemStruct>() {
286            let mut imps = Vec::new();
287            while let Ok(imp) = input.parse::<ItemImpl>() {
288                imps.push(imp);
289            }
290            return Ok(ReprStateType::Struct(strct, imps));
291        }
292        if let Ok(path) = input.parse::<TypePath>() {
293            return Ok(ReprStateType::Path(path));
294        }
295        Err(input.error(
296            "Expected one of:\
297        - A struct definition, e.g `struct State<T> {...}`\
298        - A path to an existing type, e.g `Vec<T>`\
299        ",
300        ))
301    }
302}
303
304struct ReprDefArgs {
305    /// Name of this representation rule
306    ident: Ident,
307    /// Short name used in generated representation-variable names and traces.
308    short_name: LitStr,
309    /// Representation state type
310    state_ty: ReprStateType,
311    /// Initialisation at domain level: `init: (DomainPtr) -> Result<State<DomainPtr>, ReprInitError>`
312    init_fn: ItemFn,
313    /// Generating structural constraints: `structural: (&State<DeclarationPtr>) -> Vec<Expression>`
314    structural_fn: ItemFn,
315    /// Going down: `down: (&State<DeclarationPtr>, Literal) -> Result<State<Literal>, ReprDownError>`
316    down_fn: ItemFn,
317    /// Going up: `up: State<Literal> -> Literal`
318    up_fn: ItemFn,
319    /// Getting representation variables: `repr_vars: &State<DeclarationPtr> -> VecDeque<DeclarationPtr>`
320    /// If not provided, we attempt to codegen one using uniplate
321    repr_vars_fn: Option<ItemFn>,
322    /// Measuring the representation-domain size. If omitted, this is generated using
323    /// uniplate; representations with non-uniplate containers can provide it explicitly.
324    compactness_fn: Option<ItemFn>,
325    /// Restricting this representation to some solver families:
326    /// `applies: (SolverFamily) -> bool`. If omitted, the representation is solver-independent.
327    applies_fn: Option<ItemFn>,
328}
329
330impl Parse for ReprDefArgs {
331    fn parse(input: ParseStream) -> Result<Self> {
332        let ident = input.parse::<Ident>()?;
333        let short_name_content;
334        parenthesized!(short_name_content in input);
335        let short_name = short_name_content.parse::<LitStr>()?;
336        let state_ty = input.parse::<ReprStateType>()?;
337
338        // TODO: Exact syntax subject to change
339
340        let mut funcs = HashMap::<String, ItemFn>::new();
341        let mut errors: Vec<syn::Error> = Vec::new();
342        for _ in 0..7 {
343            match input.parse::<ItemFn>() {
344                Ok(func) => {
345                    let ident = func.sig.ident.to_string();
346                    funcs.insert(ident, func);
347                }
348                Err(e) => errors.push(e),
349            }
350        }
351
352        let fmt_errors = format!(
353            "\nErrors:\n{}",
354            errors.iter().map(syn::Error::to_string).join("\n")
355        );
356
357        let init_fn = funcs.remove("init").ok_or_else(|| {
358            input.error(format!("Expected `fn init(DomainPtr) -> Result<State<DomainPtr>, ReprInitError>`{fmt_errors}"))
359        })?;
360        let structural_fn = funcs.remove("structural").ok_or_else(|| {
361            input.error(format!(
362                "Expected `fn structural(&State<DeclarationPtr>) -> Vec<Expression>`{fmt_errors}"
363            ))
364        })?;
365        let down_fn = funcs.remove("down").ok_or_else(|| input.error(format!("Expected `fn down(&State<DomainPtr>, Literal) -> Result<State<Literal>, ReprDownError>`{fmt_errors}")))?;
366        let up_fn = funcs.remove("up").ok_or_else(|| {
367            input.error(format!(
368                "Expected `fn up(State<Literal>) -> Literal`{fmt_errors}"
369            ))
370        })?;
371        let repr_vars_fn = funcs.remove("repr_vars");
372        let compactness_fn = funcs.remove("compactness");
373        let applies_fn = funcs.remove("applies");
374
375        if repr_vars_fn.is_none() && matches!(state_ty, ReprStateType::Path(..)) {
376            return Err(input.error("A repr_vars implementation is required for external types"));
377        }
378
379        Ok(Self {
380            ident,
381            short_name,
382            state_ty,
383            init_fn,
384            structural_fn,
385            down_fn,
386            up_fn,
387            repr_vars_fn,
388            compactness_fn,
389            applies_fn,
390        })
391    }
392}
393
394#[proc_macro]
395pub fn register_representation(input: TokenStream) -> TokenStream {
396    let args = parse_macro_input!(input as ReprDefArgs);
397    let repr_ident = &args.ident;
398    let repr_short_name = &args.short_name;
399    let repr_name_str = repr_ident.to_string();
400
401    // prefix for generated names
402    let prefix = format!("CONJURE_GEN_REPR_{}_", repr_name_str);
403
404    let (user_state_ident, struct_def_tokens) = match &args.state_ty {
405        ReprStateType::Struct(item_struct, _) => {
406            // get ident and body of the struct
407            let ident = item_struct.ident.clone();
408            let prefixed_ident = Ident::new(
409                &format!("{}{}", repr_name_str, ident),
410                item_struct.ident.span(),
411            );
412            let tokens = generate_struct_def(item_struct, &prefixed_ident);
413            (ident, tokens)
414        }
415        ReprStateType::Path(type_path) => {
416            // for a path like `foo::MyState`, just use it directly; no struct to emit
417            let ident = type_path
418                .path
419                .segments
420                .last()
421                .expect("state type path must have at least one segment")
422                .ident
423                .clone();
424            (ident, quote! {})
425        }
426    };
427
428    // Actual ident of the "State<T>" type
429    let state_ident = match &args.state_ty {
430        ReprStateType::Struct(..) => {
431            // prefix user-defined struct's ident so it doesn't clash with anything
432            Ident::new(
433                &format!("{}{}", repr_name_str, user_state_ident),
434                user_state_ident.span(),
435            )
436        }
437        // otherwise use the provided name as is
438        ReprStateType::Path(_) => user_state_ident.clone(),
439    };
440
441    // Rename the idents in user-defined functions to their prefixed versions
442    // e.g:
443    // MyState<T> -> CONJURE_GEN_REPR_<Rule>_MyState<T>
444    // fn init(...) -> fn CONJURE_GEN_REPR_<Rule>_init(...)
445    let prefixed_init = Ident::new(&format!("{}init", prefix), args.init_fn.sig.ident.span());
446    let prefixed_structural = Ident::new(
447        &format!("{}structural", prefix),
448        args.structural_fn.sig.ident.span(),
449    );
450    let prefixed_down = Ident::new(&format!("{}down", prefix), args.down_fn.sig.ident.span());
451    let prefixed_up = Ident::new(&format!("{}up", prefix), args.up_fn.sig.ident.span());
452    let prefixed_repr_vars = args
453        .repr_vars_fn
454        .as_ref()
455        .map(|f| Ident::new(&format!("{}repr_vars", prefix), f.sig.ident.span()));
456    let prefixed_compactness = args
457        .compactness_fn
458        .as_ref()
459        .map(|f| Ident::new(&format!("{}compactness", prefix), f.sig.ident.span()));
460    let prefixed_applies = args
461        .applies_fn
462        .as_ref()
463        .map(|f| Ident::new(&format!("{}applies", prefix), f.sig.ident.span()));
464    let applies_fn = args.applies_fn.map(|f| {
465        let f = rename_fn(f, prefixed_applies.as_ref().unwrap());
466        quote! {
467            #[allow(non_snake_case)]
468            #f
469        }
470    });
471    let applies_impl = match &prefixed_applies {
472        Some(ident) => quote! {
473            fn applies_to(family: ::conjure_cp::settings::SolverFamily) -> bool {
474                #ident(family)
475            }
476        },
477        None => quote! {},
478    };
479
480    let mut init_fn = rename_fn(args.init_fn, &prefixed_init);
481    let mut structural_fn = rename_fn(args.structural_fn, &prefixed_structural);
482    let mut down_fn = rename_fn(args.down_fn, &prefixed_down);
483    let mut up_fn = rename_fn(args.up_fn, &prefixed_up);
484    let mut repr_vars_fn = args
485        .repr_vars_fn
486        .map(|f| rename_fn(f, prefixed_repr_vars.as_ref().unwrap()));
487    let mut compactness_fn = args
488        .compactness_fn
489        .map(|f| rename_fn(f, prefixed_compactness.as_ref().unwrap()));
490
491    if matches!(&args.state_ty, ReprStateType::Struct(..)) {
492        init_fn = rename_ident_in_fn(init_fn, &user_state_ident, &state_ident);
493        structural_fn = rename_ident_in_fn(structural_fn, &user_state_ident, &state_ident);
494        down_fn = rename_ident_in_fn(down_fn, &user_state_ident, &state_ident);
495        up_fn = rename_ident_in_fn(up_fn, &user_state_ident, &state_ident);
496        repr_vars_fn = repr_vars_fn.map(|f| rename_ident_in_fn(f, &user_state_ident, &state_ident));
497        compactness_fn =
498            compactness_fn.map(|f| rename_ident_in_fn(f, &user_state_ident, &state_ident));
499    }
500
501    // Rename idents in the user-provided impl
502    let renamed_impls = if let ReprStateType::Struct(_, impls) = args.state_ty {
503        impls
504            .into_iter()
505            .map(|imp| rename_ident_in_impl(imp, &user_state_ident, &state_ident))
506            .collect()
507    } else {
508        Vec::new()
509    };
510
511    let repr_vars_impl = if repr_vars_fn.is_some() {
512        quote! {#prefixed_repr_vars(self)}
513    } else {
514        quote! {self.__collect_t_children()}
515    };
516    let repr_vars_fn_toks = repr_vars_fn.map(|f| {
517        quote! {
518            #[allow(non_snake_case)]
519            #f
520        }
521    });
522    let compactness_impl = if compactness_fn.is_some() {
523        quote! {#prefixed_compactness(self)}
524    } else {
525        quote! {
526            self.__collect_t_children()
527                .iter()
528                .map(default_impls::domain_size)
529                .fold(1usize, usize::saturating_mul)
530        }
531    };
532    let compactness_fn_toks = compactness_fn.map(|f| {
533        quote! {
534            #[allow(non_snake_case)]
535            #f
536        }
537    });
538
539    // Static name for the registry entry
540    let dist_slice_name = format!("CONJURE_GEN_REPR_{}", repr_name_str).to_uppercase();
541    let dist_slice_ident = Ident::new(&dist_slice_name, repr_ident.span());
542
543    let init_cache_name = format!("CONJURE_GEN_REPR_{}_INIT_CACHE", repr_name_str).to_uppercase();
544    let init_cache_ident = Ident::new(&init_cache_name, repr_ident.span());
545
546    let expanded = quote! {
547        // -- Dependencies
548        use ::conjure_cp::representation::_dependencies::*;
549        use ::conjure_cp::ast::{
550            DeclarationPtr, DomainPtr, Expression, Literal, SymbolTable, Name
551        };
552
553        // -- User-provided struct definition
554        #struct_def_tokens
555
556        // -- User-provided struct impl
557        #(#renamed_impls)*
558
559        // -- User-provided functions
560        #[allow(non_snake_case)]
561        #init_fn
562        #[allow(non_snake_case)]
563        #structural_fn
564        #[allow(non_snake_case)]
565        #down_fn
566        #[allow(non_snake_case)]
567        #up_fn
568        #repr_vars_fn_toks
569        #compactness_fn_toks
570
571        static #init_cache_ident: std::sync::LazyLock<FrozenMap<DomainPtr, Box<#state_ident<DomainPtr>>>> = std::sync::LazyLock::new(|| FrozenMap::new());
572
573        // -- Trait implementations
574        impl ReprDomainLevel for #state_ident<DomainPtr> {
575            const RULE: &'static dyn ReprRuleStored = &#repr_ident;
576            type Assignment = #state_ident<Literal>;
577            type DeclLevel = #state_ident<DeclarationPtr>;
578
579            fn init(dom: DomainPtr) -> ::core::result::Result<Self, ReprInitError>
580            where
581                Self: Sized,
582            {
583                if let Some(res) = #init_cache_ident.get(&dom) {
584                    return Ok(res.clone());
585                }
586
587                let res = #prefixed_init(dom.clone())?;
588                let _ = #init_cache_ident.insert(dom, Box::new(res.clone()));
589                Ok(res)
590            }
591
592            fn compactness_score(&self) -> usize {
593                #compactness_impl
594            }
595
596            fn down(
597                &self,
598                value: Literal,
599            ) -> ::core::result::Result<Self::Assignment, ReprDownError> {
600                #prefixed_down(self, value)
601            }
602
603            fn instantiate(self, decl: DeclarationPtr) -> ReprInstantiateResult<Self::DeclLevel> {
604                default_impls::instantiate_default_impl(self, decl, #prefixed_structural)
605            }
606        }
607
608        impl ReprDeclLevel for #state_ident<DeclarationPtr> {
609            const RULE: &'static dyn ReprRuleStored = &#repr_ident;
610            type Assignment = #state_ident<Literal>;
611            type DomainLevel = #state_ident<DomainPtr>;
612
613            fn to_domain_level(self) -> Self::DomainLevel {
614                default_impls::to_domain_level_default_impl(self)
615            }
616
617            fn lookup_via(
618                &self,
619                lookup: &LookupFn<'_>,
620            ) -> ::core::result::Result<Self::Assignment, ReprUpError> {
621                default_impls::lookup_via_default_impl(self, lookup)
622            }
623
624            fn repr_vars(&self) -> ::std::collections::VecDeque<DeclarationPtr> {
625                #repr_vars_impl
626            }
627        }
628
629        impl ReprAssignment for #state_ident<Literal> {
630            fn up(self) -> Literal {
631                #prefixed_up(self)
632            }
633        }
634
635        // -- ReprRule marker struct
636        pub struct #repr_ident;
637
638        #applies_fn
639
640        impl ReprRule for #repr_ident {
641            const STORED: &'static dyn ReprRuleStored = &#repr_ident;
642            const NAME: &'static str = #repr_name_str;
643            const SHORT_NAME: &'static str = #repr_short_name;
644            type Assignment = #state_ident<Literal>;
645            type DeclLevel = #state_ident<DeclarationPtr>;
646            type DomainLevel = #state_ident<DomainPtr>;
647
648            #applies_impl
649        }
650
651        // -- Registry entry
652        pub static #dist_slice_ident: &'static dyn ReprRuleStored = &#repr_ident;
653
654        ::conjure_cp::representation::_dependencies::inventory::submit! {
655            #dist_slice_ident
656        }
657    };
658
659    TokenStream::from(expanded)
660}
661
662/// Generates the struct definition with the necessary derive macros and serde attributes.
663/// The struct is always emitted as `pub` with the given `prefixed_ident` as its name.
664fn generate_struct_def(
665    item_struct: &ItemStruct,
666    prefixed_ident: &Ident,
667) -> proc_macro2::TokenStream {
668    // Find the generic type parameter name (e.g. `T`)
669    let generic_param_ident = item_struct
670        .generics
671        .params
672        .iter()
673        .find_map(|p| {
674            if let GenericParam::Type(tp) = p {
675                Some(tp.ident.clone())
676            } else {
677                None
678            }
679        })
680        .expect("state struct must have exactly one type parameter");
681
682    let generics = &item_struct.generics;
683
684    let serde_bound = format!(
685        "ReprStateSerde: SerializeAs<{0}> + for<'d> DeserializeAs<'d, {0}>",
686        generic_param_ident
687    );
688
689    let mut collect_children_exprs: Vec<proc_macro2::TokenStream> = Vec::new();
690    let mut biplate_bounds: Vec<proc_macro2::TokenStream> = vec![quote! {
691        #generic_param_ident: ::conjure_cp::representation::_dependencies::uniplate::Uniplate
692    }];
693
694    let fields = match &item_struct.fields {
695        syn::Fields::Named(named) => {
696            let field_tokens: Vec<_> = named
697                .named
698                .iter()
699                .map(|f| {
700                    let field_attrs = &f.attrs;
701                    let field_vis = &f.vis;
702                    let field_ident = f
703                        .ident
704                        .as_ref()
705                        .expect("named field must have an identifier");
706                    let field_ty = &f.ty;
707
708                    if type_contains_ident(&f.ty, &generic_param_ident) {
709                        if type_is_ident(&f.ty, &generic_param_ident) {
710                            collect_children_exprs.push(quote! {
711                                children.push_back(self.#field_ident.clone());
712                            });
713                        } else if type_is_option_of_ident(&f.ty, &generic_param_ident) {
714                            collect_children_exprs.push(quote! {
715                                children.extend(self.#field_ident.iter().cloned());
716                            });
717                        } else {
718                            collect_children_exprs.push(quote! {
719                                children.extend(
720                                    ::conjure_cp::representation::_dependencies::uniplate::Biplate::<#generic_param_ident>::children_bi(&self.#field_ident)
721                                );
722                            });
723                            biplate_bounds.push(quote! {
724                                #field_ty: ::conjure_cp::representation::_dependencies::uniplate::Biplate<#generic_param_ident>
725                            });
726                        }
727
728                        let serde_as_ty =
729                            build_serde_as_type(field_ty, &generic_param_ident, "ReprStateSerde");
730                        let serde_as_str = serde_as_ty.to_token_stream().to_string();
731                        quote! {
732                            #(#field_attrs)*
733                            #[serde_as(as = #serde_as_str)]
734                            #field_vis #field_ident: #field_ty
735                        }
736                    } else {
737                        quote! {
738                            #(#field_attrs)*
739                            #field_vis #field_ident: #field_ty
740                        }
741                    }
742                })
743                .collect();
744
745            quote! { { #(#field_tokens),* } }
746        }
747        syn::Fields::Unnamed(unnamed) => {
748            let field_tokens: Vec<_> = unnamed
749                .unnamed
750                .iter()
751                .enumerate()
752                .map(|(idx, f)| {
753                    let field_attrs = &f.attrs;
754                    let field_vis = &f.vis;
755                    let field_ty = &f.ty;
756
757                    if type_contains_ident(&f.ty, &generic_param_ident) {
758                        let index = syn::Index::from(idx);
759                        if type_is_ident(&f.ty, &generic_param_ident) {
760                            collect_children_exprs.push(quote! {
761                                children.push_back(self.#index.clone());
762                            });
763                        } else if type_is_option_of_ident(&f.ty, &generic_param_ident) {
764                            collect_children_exprs.push(quote! {
765                                children.extend(self.#index.iter().cloned());
766                            });
767                        } else {
768                            collect_children_exprs.push(quote! {
769                                children.extend(
770                                    ::conjure_cp::representation::_dependencies::uniplate::Biplate::<#generic_param_ident>::children_bi(&self.#index)
771                                );
772                            });
773                            biplate_bounds.push(quote! {
774                                #field_ty: ::conjure_cp::representation::_dependencies::uniplate::Biplate<#generic_param_ident>
775                            });
776                        }
777
778                        let serde_as_ty =
779                            build_serde_as_type(field_ty, &generic_param_ident, "ReprStateSerde");
780                        let serde_as_str = serde_as_ty.to_token_stream().to_string();
781                        quote! {
782                            #(#field_attrs)*
783                            #[serde_as(as = #serde_as_str)]
784                            #field_vis #field_ty
785                        }
786                    } else {
787                        quote! {
788                            #(#field_attrs)*
789                            #field_vis #field_ty
790                        }
791                    }
792                })
793                .collect();
794
795            quote! { ( #(#field_tokens),* ); }
796        }
797        syn::Fields::Unit => quote! { ; },
798    };
799
800    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
801    let method_where_clause = if biplate_bounds.is_empty() {
802        quote! {}
803    } else {
804        quote! { where #(#biplate_bounds),* }
805    };
806
807    quote! {
808        #[allow(non_camel_case_types)]
809        #[::conjure_cp::representation::_dependencies::serde_with::serde_as(
810            crate = "::conjure_cp::representation::_dependencies::serde_with"
811        )]
812        #[derive(
813            Debug,
814            Clone,
815            PartialEq,
816            Eq,
817            ::conjure_cp::representation::_dependencies::funcmap::FuncMap,
818            ::conjure_cp::representation::_dependencies::funcmap::TryFuncMap,
819            ::conjure_cp::representation::_dependencies::serde::Serialize,
820            ::conjure_cp::representation::_dependencies::serde::Deserialize
821        )]
822        #[serde(
823            crate = "::conjure_cp::representation::_dependencies::serde",
824            bound = #serde_bound
825        )]
826        #[funcmap(crate = "::conjure_cp::representation::_dependencies::funcmap")]
827        pub struct #prefixed_ident #generics #fields
828
829        impl #impl_generics #prefixed_ident #ty_generics #where_clause {
830            fn __collect_t_children(&self) -> ::std::collections::VecDeque<#generic_param_ident>
831            #method_where_clause
832            {
833                let mut children = ::std::collections::VecDeque::new();
834                #(#collect_children_exprs)*
835                children
836            }
837        }
838    }
839}
840
841fn parse_parenthesized<T: Parse>(input: ParseStream) -> Result<Vec<T>> {
842    let content;
843    parenthesized!(content in input);
844
845    let mut paths = Vec::new();
846    while !content.is_empty() {
847        let path = content.parse()?;
848        paths.push(path);
849        if content.is_empty() {
850            break;
851        }
852        content.parse::<Comma>()?;
853    }
854
855    Ok(paths)
856}
857
858struct RuleSetArgs {
859    name: LitStr,
860    dependencies: Vec<LitStr>,
861    applies_fn: Option<ExprClosure>,
862}
863
864impl Parse for RuleSetArgs {
865    fn parse(input: ParseStream) -> Result<Self> {
866        let name = input.parse()?;
867
868        if input.is_empty() {
869            return Ok(Self {
870                name,
871                dependencies: Vec::new(),
872                applies_fn: None,
873            });
874        }
875
876        input.parse::<Comma>()?;
877        let dependencies = parse_parenthesized::<LitStr>(input)?;
878
879        if input.is_empty() {
880            return Ok(Self {
881                name,
882                dependencies,
883                applies_fn: None,
884            });
885        }
886
887        input.parse::<Comma>()?;
888        let applies_fn = input.parse::<ExprClosure>()?;
889
890        Ok(Self {
891            name,
892            dependencies,
893            applies_fn: Some(applies_fn),
894        })
895    }
896}
897
898/**
899* Register a rule set with the given name, dependencies, and metadata.
900*
901* # Example
902* ```rust
903 * use conjure_cp_rule_macros::register_rule_set;
904 * register_rule_set!("MyRuleSet", ("DependencyRuleSet", "AnotherRuleSet"));
905* ```
906 */
907#[proc_macro]
908pub fn register_rule_set(args: TokenStream) -> TokenStream {
909    let RuleSetArgs {
910        name,
911        dependencies,
912        applies_fn,
913    } = parse_macro_input!(args as RuleSetArgs);
914
915    let static_name = format!("CONJURE_GEN_RULE_SET_{}", name.value()).to_uppercase();
916    let static_ident = Ident::new(&static_name, Span::call_site());
917
918    let dependencies = quote! {
919        #(#dependencies),*
920    };
921
922    let applies_to_family = match applies_fn {
923        // Does not apply by default, e.g. only used as a dependency
924        None => quote! { |_: &::conjure_cp::settings::SolverFamily| false },
925        Some(func) => quote! { #func },
926    };
927
928    let expanded = quote! {
929        pub static #static_ident: ::conjure_cp::rule_engine::RuleSet<'static> =
930            ::conjure_cp::rule_engine::RuleSet::new(#name, &[#dependencies], #applies_to_family);
931
932        ::conjure_cp::rule_engine::_dependencies::inventory::submit! {
933            &#static_ident
934        }
935    };
936
937    TokenStream::from(expanded)
938}
conjure_cp_rule_macros/lib.rs

conjure_cp_rule_macros/
lib.rs
