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    integer_encoding_fn: Option<ItemFn>,
329    probe_score_fn: Option<ItemFn>,
330}
331
332impl Parse for ReprDefArgs {
333    fn parse(input: ParseStream) -> Result<Self> {
334        let ident = input.parse::<Ident>()?;
335        let short_name_content;
336        parenthesized!(short_name_content in input);
337        let short_name = short_name_content.parse::<LitStr>()?;
338        let state_ty = input.parse::<ReprStateType>()?;
339
340        // TODO: Exact syntax subject to change
341
342        let mut funcs = HashMap::<String, ItemFn>::new();
343        let mut errors: Vec<syn::Error> = Vec::new();
344        for _ in 0..9 {
345            match input.parse::<ItemFn>() {
346                Ok(func) => {
347                    let ident = func.sig.ident.to_string();
348                    funcs.insert(ident, func);
349                }
350                Err(e) => errors.push(e),
351            }
352        }
353
354        let fmt_errors = format!(
355            "\nErrors:\n{}",
356            errors.iter().map(syn::Error::to_string).join("\n")
357        );
358
359        let init_fn = funcs.remove("init").ok_or_else(|| {
360            input.error(format!("Expected `fn init(DomainPtr) -> Result<State<DomainPtr>, ReprInitError>`{fmt_errors}"))
361        })?;
362        let structural_fn = funcs.remove("structural").ok_or_else(|| {
363            input.error(format!(
364                "Expected `fn structural(&State<DeclarationPtr>) -> Vec<Expression>`{fmt_errors}"
365            ))
366        })?;
367        let down_fn = funcs.remove("down").ok_or_else(|| input.error(format!("Expected `fn down(&State<DomainPtr>, Literal) -> Result<State<Literal>, ReprDownError>`{fmt_errors}")))?;
368        let up_fn = funcs.remove("up").ok_or_else(|| {
369            input.error(format!(
370                "Expected `fn up(State<Literal>) -> Literal`{fmt_errors}"
371            ))
372        })?;
373        let repr_vars_fn = funcs.remove("repr_vars");
374        let compactness_fn = funcs.remove("compactness");
375        let applies_fn = funcs.remove("applies");
376        let integer_encoding_fn = funcs.remove("integer_encoding");
377        let probe_score_fn = funcs.remove("probe_score");
378
379        if repr_vars_fn.is_none() && matches!(state_ty, ReprStateType::Path(..)) {
380            return Err(input.error("A repr_vars implementation is required for external types"));
381        }
382
383        Ok(Self {
384            ident,
385            short_name,
386            state_ty,
387            init_fn,
388            structural_fn,
389            down_fn,
390            up_fn,
391            repr_vars_fn,
392            compactness_fn,
393            applies_fn,
394            integer_encoding_fn,
395            probe_score_fn,
396        })
397    }
398}
399
400#[proc_macro]
401pub fn register_representation(input: TokenStream) -> TokenStream {
402    let args = parse_macro_input!(input as ReprDefArgs);
403    let repr_ident = &args.ident;
404    let repr_short_name = &args.short_name;
405    let repr_name_str = repr_ident.to_string();
406
407    // prefix for generated names
408    let prefix = format!("CONJURE_GEN_REPR_{}_", repr_name_str);
409
410    let (user_state_ident, struct_def_tokens) = match &args.state_ty {
411        ReprStateType::Struct(item_struct, _) => {
412            // get ident and body of the struct
413            let ident = item_struct.ident.clone();
414            let prefixed_ident = Ident::new(
415                &format!("{}{}", repr_name_str, ident),
416                item_struct.ident.span(),
417            );
418            let tokens = generate_struct_def(item_struct, &prefixed_ident);
419            (ident, tokens)
420        }
421        ReprStateType::Path(type_path) => {
422            // for a path like `foo::MyState`, just use it directly; no struct to emit
423            let ident = type_path
424                .path
425                .segments
426                .last()
427                .expect("state type path must have at least one segment")
428                .ident
429                .clone();
430            (ident, quote! {})
431        }
432    };
433
434    // Actual ident of the "State<T>" type
435    let state_ident = match &args.state_ty {
436        ReprStateType::Struct(..) => {
437            // prefix user-defined struct's ident so it doesn't clash with anything
438            Ident::new(
439                &format!("{}{}", repr_name_str, user_state_ident),
440                user_state_ident.span(),
441            )
442        }
443        // otherwise use the provided name as is
444        ReprStateType::Path(_) => user_state_ident.clone(),
445    };
446
447    // Rename the idents in user-defined functions to their prefixed versions
448    // e.g:
449    // MyState<T> -> CONJURE_GEN_REPR_<Rule>_MyState<T>
450    // fn init(...) -> fn CONJURE_GEN_REPR_<Rule>_init(...)
451    let prefixed_init = Ident::new(&format!("{}init", prefix), args.init_fn.sig.ident.span());
452    let prefixed_structural = Ident::new(
453        &format!("{}structural", prefix),
454        args.structural_fn.sig.ident.span(),
455    );
456    let prefixed_down = Ident::new(&format!("{}down", prefix), args.down_fn.sig.ident.span());
457    let prefixed_up = Ident::new(&format!("{}up", prefix), args.up_fn.sig.ident.span());
458    let prefixed_repr_vars = args
459        .repr_vars_fn
460        .as_ref()
461        .map(|f| Ident::new(&format!("{}repr_vars", prefix), f.sig.ident.span()));
462    let prefixed_compactness = args
463        .compactness_fn
464        .as_ref()
465        .map(|f| Ident::new(&format!("{}compactness", prefix), f.sig.ident.span()));
466    let prefixed_applies = args
467        .applies_fn
468        .as_ref()
469        .map(|f| Ident::new(&format!("{}applies", prefix), f.sig.ident.span()));
470    let applies_fn = args.applies_fn.map(|f| {
471        let f = rename_fn(f, prefixed_applies.as_ref().unwrap());
472        quote! {
473            #[allow(non_snake_case)]
474            #f
475        }
476    });
477    let applies_impl = match &prefixed_applies {
478        Some(ident) => quote! {
479            fn applies_to(family: ::conjure_cp::settings::SolverFamily) -> bool {
480                #ident(family)
481            }
482        },
483        None => quote! {},
484    };
485
486    let probe_score_impl = args.probe_score_fn.map(|f| quote! { #f });
487
488    let integer_encoding_impl = args.integer_encoding_fn.map(|mut f| {
489        f.sig.ident = Ident::new("is_integer_encoding", f.sig.ident.span());
490        quote! { #f }
491    });
492
493    let mut init_fn = rename_fn(args.init_fn, &prefixed_init);
494    let mut structural_fn = rename_fn(args.structural_fn, &prefixed_structural);
495    let mut down_fn = rename_fn(args.down_fn, &prefixed_down);
496    let mut up_fn = rename_fn(args.up_fn, &prefixed_up);
497    let mut repr_vars_fn = args
498        .repr_vars_fn
499        .map(|f| rename_fn(f, prefixed_repr_vars.as_ref().unwrap()));
500    let mut compactness_fn = args
501        .compactness_fn
502        .map(|f| rename_fn(f, prefixed_compactness.as_ref().unwrap()));
503
504    if matches!(&args.state_ty, ReprStateType::Struct(..)) {
505        init_fn = rename_ident_in_fn(init_fn, &user_state_ident, &state_ident);
506        structural_fn = rename_ident_in_fn(structural_fn, &user_state_ident, &state_ident);
507        down_fn = rename_ident_in_fn(down_fn, &user_state_ident, &state_ident);
508        up_fn = rename_ident_in_fn(up_fn, &user_state_ident, &state_ident);
509        repr_vars_fn = repr_vars_fn.map(|f| rename_ident_in_fn(f, &user_state_ident, &state_ident));
510        compactness_fn =
511            compactness_fn.map(|f| rename_ident_in_fn(f, &user_state_ident, &state_ident));
512    }
513
514    // Rename idents in the user-provided impl
515    let renamed_impls = if let ReprStateType::Struct(_, impls) = args.state_ty {
516        impls
517            .into_iter()
518            .map(|imp| rename_ident_in_impl(imp, &user_state_ident, &state_ident))
519            .collect()
520    } else {
521        Vec::new()
522    };
523
524    let repr_vars_impl = if repr_vars_fn.is_some() {
525        quote! {#prefixed_repr_vars(self)}
526    } else {
527        quote! {self.__collect_t_children()}
528    };
529    let repr_vars_fn_toks = repr_vars_fn.map(|f| {
530        quote! {
531            #[allow(non_snake_case)]
532            #f
533        }
534    });
535    let compactness_impl = if compactness_fn.is_some() {
536        quote! {#prefixed_compactness(self)}
537    } else {
538        quote! {
539            self.__collect_t_children()
540                .iter()
541                .map(default_impls::domain_size)
542                .fold(1usize, usize::saturating_mul)
543        }
544    };
545    let compactness_fn_toks = compactness_fn.map(|f| {
546        quote! {
547            #[allow(non_snake_case)]
548            #f
549        }
550    });
551
552    // Static name for the registry entry
553    let dist_slice_name = format!("CONJURE_GEN_REPR_{}", repr_name_str).to_uppercase();
554    let dist_slice_ident = Ident::new(&dist_slice_name, repr_ident.span());
555
556    let init_cache_name = format!("CONJURE_GEN_REPR_{}_INIT_CACHE", repr_name_str).to_uppercase();
557    let init_cache_ident = Ident::new(&init_cache_name, repr_ident.span());
558
559    let expanded = quote! {
560        // -- Dependencies
561        use ::conjure_cp::representation::_dependencies::*;
562        use ::conjure_cp::ast::{
563            DeclarationPtr, DomainPtr, Expression, Literal, SymbolTable, Name
564        };
565
566        // -- User-provided struct definition
567        #struct_def_tokens
568
569        // -- User-provided struct impl
570        #(#renamed_impls)*
571
572        // -- User-provided functions
573        #[allow(non_snake_case)]
574        #init_fn
575        #[allow(non_snake_case)]
576        #structural_fn
577        #[allow(non_snake_case)]
578        #down_fn
579        #[allow(non_snake_case)]
580        #up_fn
581        #repr_vars_fn_toks
582        #compactness_fn_toks
583
584        static #init_cache_ident: std::sync::LazyLock<FrozenMap<DomainPtr, Box<#state_ident<DomainPtr>>>> = std::sync::LazyLock::new(|| FrozenMap::new());
585
586        // -- Trait implementations
587        impl ReprDomainLevel for #state_ident<DomainPtr> {
588            const RULE: &'static dyn ReprRuleStored = &#repr_ident;
589            type Assignment = #state_ident<Literal>;
590            type DeclLevel = #state_ident<DeclarationPtr>;
591
592            fn init(dom: DomainPtr) -> ::core::result::Result<Self, ReprInitError>
593            where
594                Self: Sized,
595            {
596                if let Some(res) = #init_cache_ident.get(&dom) {
597                    return Ok(res.clone());
598                }
599
600                let res = #prefixed_init(dom.clone())?;
601                let _ = #init_cache_ident.insert(dom, Box::new(res.clone()));
602                Ok(res)
603            }
604
605            fn compactness_score(&self) -> usize {
606                #compactness_impl
607            }
608
609            fn down(
610                &self,
611                value: Literal,
612            ) -> ::core::result::Result<Self::Assignment, ReprDownError> {
613                #prefixed_down(self, value)
614            }
615
616            fn instantiate(self, decl: DeclarationPtr) -> ReprInstantiateResult<Self::DeclLevel> {
617                default_impls::instantiate_default_impl(self, decl, #prefixed_structural)
618            }
619        }
620
621        impl ReprDeclLevel for #state_ident<DeclarationPtr> {
622            const RULE: &'static dyn ReprRuleStored = &#repr_ident;
623            type Assignment = #state_ident<Literal>;
624            type DomainLevel = #state_ident<DomainPtr>;
625
626            fn to_domain_level(self) -> Self::DomainLevel {
627                default_impls::to_domain_level_default_impl(self)
628            }
629
630            fn lookup_via(
631                &self,
632                lookup: &LookupFn<'_>,
633            ) -> ::core::result::Result<Self::Assignment, ReprUpError> {
634                default_impls::lookup_via_default_impl(self, lookup)
635            }
636
637            fn repr_vars(&self) -> ::std::collections::VecDeque<DeclarationPtr> {
638                #repr_vars_impl
639            }
640        }
641
642        impl ReprAssignment for #state_ident<Literal> {
643            fn up(self) -> Literal {
644                #prefixed_up(self)
645            }
646        }
647
648        // -- ReprRule marker struct
649        pub struct #repr_ident;
650
651        #applies_fn
652
653        impl ReprRule for #repr_ident {
654            const STORED: &'static dyn ReprRuleStored = &#repr_ident;
655            const NAME: &'static str = #repr_name_str;
656            const SHORT_NAME: &'static str = #repr_short_name;
657            type Assignment = #state_ident<Literal>;
658            type DeclLevel = #state_ident<DeclarationPtr>;
659            type DomainLevel = #state_ident<DomainPtr>;
660
661            #applies_impl
662            #integer_encoding_impl
663            #probe_score_impl
664        }
665
666        // -- Registry entry
667        pub static #dist_slice_ident: &'static dyn ReprRuleStored = &#repr_ident;
668
669        ::conjure_cp::representation::_dependencies::inventory::submit! {
670            #dist_slice_ident
671        }
672    };
673
674    TokenStream::from(expanded)
675}
676
677/// Generates the struct definition with the necessary derive macros and serde attributes.
678/// The struct is always emitted as `pub` with the given `prefixed_ident` as its name.
679fn generate_struct_def(
680    item_struct: &ItemStruct,
681    prefixed_ident: &Ident,
682) -> proc_macro2::TokenStream {
683    // Find the generic type parameter name (e.g. `T`)
684    let generic_param_ident = item_struct
685        .generics
686        .params
687        .iter()
688        .find_map(|p| {
689            if let GenericParam::Type(tp) = p {
690                Some(tp.ident.clone())
691            } else {
692                None
693            }
694        })
695        .expect("state struct must have exactly one type parameter");
696
697    let generics = &item_struct.generics;
698
699    let serde_bound = format!(
700        "ReprStateSerde: SerializeAs<{0}> + for<'d> DeserializeAs<'d, {0}>",
701        generic_param_ident
702    );
703
704    let mut collect_children_exprs: Vec<proc_macro2::TokenStream> = Vec::new();
705    let mut biplate_bounds: Vec<proc_macro2::TokenStream> = vec![quote! {
706        #generic_param_ident: ::conjure_cp::representation::_dependencies::uniplate::Uniplate
707    }];
708
709    let fields = match &item_struct.fields {
710        syn::Fields::Named(named) => {
711            let field_tokens: Vec<_> = named
712                .named
713                .iter()
714                .map(|f| {
715                    let field_attrs = &f.attrs;
716                    let field_vis = &f.vis;
717                    let field_ident = f
718                        .ident
719                        .as_ref()
720                        .expect("named field must have an identifier");
721                    let field_ty = &f.ty;
722
723                    if type_contains_ident(&f.ty, &generic_param_ident) {
724                        if type_is_ident(&f.ty, &generic_param_ident) {
725                            collect_children_exprs.push(quote! {
726                                children.push_back(self.#field_ident.clone());
727                            });
728                        } else if type_is_option_of_ident(&f.ty, &generic_param_ident) {
729                            collect_children_exprs.push(quote! {
730                                children.extend(self.#field_ident.iter().cloned());
731                            });
732                        } else {
733                            collect_children_exprs.push(quote! {
734                                children.extend(
735                                    ::conjure_cp::representation::_dependencies::uniplate::Biplate::<#generic_param_ident>::children_bi(&self.#field_ident)
736                                );
737                            });
738                            biplate_bounds.push(quote! {
739                                #field_ty: ::conjure_cp::representation::_dependencies::uniplate::Biplate<#generic_param_ident>
740                            });
741                        }
742
743                        let serde_as_ty =
744                            build_serde_as_type(field_ty, &generic_param_ident, "ReprStateSerde");
745                        let serde_as_str = serde_as_ty.to_token_stream().to_string();
746                        quote! {
747                            #(#field_attrs)*
748                            #[serde_as(as = #serde_as_str)]
749                            #field_vis #field_ident: #field_ty
750                        }
751                    } else {
752                        quote! {
753                            #(#field_attrs)*
754                            #field_vis #field_ident: #field_ty
755                        }
756                    }
757                })
758                .collect();
759
760            quote! { { #(#field_tokens),* } }
761        }
762        syn::Fields::Unnamed(unnamed) => {
763            let field_tokens: Vec<_> = unnamed
764                .unnamed
765                .iter()
766                .enumerate()
767                .map(|(idx, f)| {
768                    let field_attrs = &f.attrs;
769                    let field_vis = &f.vis;
770                    let field_ty = &f.ty;
771
772                    if type_contains_ident(&f.ty, &generic_param_ident) {
773                        let index = syn::Index::from(idx);
774                        if type_is_ident(&f.ty, &generic_param_ident) {
775                            collect_children_exprs.push(quote! {
776                                children.push_back(self.#index.clone());
777                            });
778                        } else if type_is_option_of_ident(&f.ty, &generic_param_ident) {
779                            collect_children_exprs.push(quote! {
780                                children.extend(self.#index.iter().cloned());
781                            });
782                        } else {
783                            collect_children_exprs.push(quote! {
784                                children.extend(
785                                    ::conjure_cp::representation::_dependencies::uniplate::Biplate::<#generic_param_ident>::children_bi(&self.#index)
786                                );
787                            });
788                            biplate_bounds.push(quote! {
789                                #field_ty: ::conjure_cp::representation::_dependencies::uniplate::Biplate<#generic_param_ident>
790                            });
791                        }
792
793                        let serde_as_ty =
794                            build_serde_as_type(field_ty, &generic_param_ident, "ReprStateSerde");
795                        let serde_as_str = serde_as_ty.to_token_stream().to_string();
796                        quote! {
797                            #(#field_attrs)*
798                            #[serde_as(as = #serde_as_str)]
799                            #field_vis #field_ty
800                        }
801                    } else {
802                        quote! {
803                            #(#field_attrs)*
804                            #field_vis #field_ty
805                        }
806                    }
807                })
808                .collect();
809
810            quote! { ( #(#field_tokens),* ); }
811        }
812        syn::Fields::Unit => quote! { ; },
813    };
814
815    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
816    let method_where_clause = if biplate_bounds.is_empty() {
817        quote! {}
818    } else {
819        quote! { where #(#biplate_bounds),* }
820    };
821
822    quote! {
823        #[allow(non_camel_case_types)]
824        #[::conjure_cp::representation::_dependencies::serde_with::serde_as(
825            crate = "::conjure_cp::representation::_dependencies::serde_with"
826        )]
827        #[derive(
828            Debug,
829            Clone,
830            PartialEq,
831            Eq,
832            ::conjure_cp::representation::_dependencies::funcmap::FuncMap,
833            ::conjure_cp::representation::_dependencies::funcmap::TryFuncMap,
834            ::conjure_cp::representation::_dependencies::serde::Serialize,
835            ::conjure_cp::representation::_dependencies::serde::Deserialize
836        )]
837        #[serde(
838            crate = "::conjure_cp::representation::_dependencies::serde",
839            bound = #serde_bound
840        )]
841        #[funcmap(crate = "::conjure_cp::representation::_dependencies::funcmap")]
842        pub struct #prefixed_ident #generics #fields
843
844        impl #impl_generics #prefixed_ident #ty_generics #where_clause {
845            fn __collect_t_children(&self) -> ::std::collections::VecDeque<#generic_param_ident>
846            #method_where_clause
847            {
848                let mut children = ::std::collections::VecDeque::new();
849                #(#collect_children_exprs)*
850                children
851            }
852        }
853    }
854}
855
856fn parse_parenthesized<T: Parse>(input: ParseStream) -> Result<Vec<T>> {
857    let content;
858    parenthesized!(content in input);
859
860    let mut paths = Vec::new();
861    while !content.is_empty() {
862        let path = content.parse()?;
863        paths.push(path);
864        if content.is_empty() {
865            break;
866        }
867        content.parse::<Comma>()?;
868    }
869
870    Ok(paths)
871}
872
873struct RuleSetArgs {
874    name: LitStr,
875    dependencies: Vec<LitStr>,
876    applies_fn: Option<ExprClosure>,
877}
878
879impl Parse for RuleSetArgs {
880    fn parse(input: ParseStream) -> Result<Self> {
881        let name = input.parse()?;
882
883        if input.is_empty() {
884            return Ok(Self {
885                name,
886                dependencies: Vec::new(),
887                applies_fn: None,
888            });
889        }
890
891        input.parse::<Comma>()?;
892        let dependencies = parse_parenthesized::<LitStr>(input)?;
893
894        if input.is_empty() {
895            return Ok(Self {
896                name,
897                dependencies,
898                applies_fn: None,
899            });
900        }
901
902        input.parse::<Comma>()?;
903        let applies_fn = input.parse::<ExprClosure>()?;
904
905        Ok(Self {
906            name,
907            dependencies,
908            applies_fn: Some(applies_fn),
909        })
910    }
911}
912
913/**
914* Register a rule set with the given name, dependencies, and metadata.
915*
916* # Example
917* ```rust
918 * use conjure_cp_rule_macros::register_rule_set;
919 * register_rule_set!("MyRuleSet", ("DependencyRuleSet", "AnotherRuleSet"));
920* ```
921 */
922#[proc_macro]
923pub fn register_rule_set(args: TokenStream) -> TokenStream {
924    let RuleSetArgs {
925        name,
926        dependencies,
927        applies_fn,
928    } = parse_macro_input!(args as RuleSetArgs);
929
930    let static_name = format!("CONJURE_GEN_RULE_SET_{}", name.value()).to_uppercase();
931    let static_ident = Ident::new(&static_name, Span::call_site());
932
933    let dependencies = quote! {
934        #(#dependencies),*
935    };
936
937    let applies_to_family = match applies_fn {
938        // Does not apply by default, e.g. only used as a dependency
939        None => quote! { |_: &::conjure_cp::settings::SolverFamily| false },
940        Some(func) => quote! { #func },
941    };
942
943    let expanded = quote! {
944        pub static #static_ident: ::conjure_cp::rule_engine::RuleSet<'static> =
945            ::conjure_cp::rule_engine::RuleSet::new(#name, &[#dependencies], #applies_to_family);
946
947        ::conjure_cp::rule_engine::_dependencies::inventory::submit! {
948            &#static_ident
949        }
950    };
951
952    TokenStream::from(expanded)
953}
conjure_cp_rule_macros/lib.rs

conjure_cp_rule_macros/
lib.rs
