oximo_macros/
lib.rs

1#![doc = include_str!("../README.md")]
2#![forbid(unsafe_code)]
3
4use proc_macro::TokenStream;
5use proc_macro_crate::{FoundCrate, crate_name};
6use proc_macro2::{Delimiter, Spacing, TokenStream as TokenStream2, TokenTree};
7use quote::quote;
8use syn::Ident;
9
10mod bind;
11mod constraint;
12mod index;
13mod objective;
14mod param;
15mod set;
16mod sum;
17mod variable;
18
19use bind::{Binds, IndexBind};
20
21/// Resolve the path prefix used to reach `__macro_support`. Prefers the umbrella
22/// `oximo` crate (which re-exports the support module) and falls back to
23/// `oximo-core`.
24fn oximo_root() -> TokenStream2 {
25    fn to_path(found: &FoundCrate, fallback: &str) -> TokenStream2 {
26        let name = match found {
27            FoundCrate::Itself => fallback,
28            FoundCrate::Name(n) => n.as_str(),
29        };
30        let id = Ident::new(name, proc_macro2::Span::call_site());
31        quote!(::#id)
32    }
33
34    if let Ok(found) = crate_name("oximo") {
35        return to_path(&found, "oximo");
36    }
37    if let Ok(found) = crate_name("oximo-core") {
38        return to_path(&found, "oximo_core");
39    }
40    quote!(::oximo_core)
41}
42
43/// `variable!(model, spec)`, declare a decision variable (or an indexed family)
44/// and bind it to a local of the same name. See the crate docs for the grammar.
45#[proc_macro]
46pub fn variable(input: TokenStream) -> TokenStream {
47    variable::expand(input.into()).unwrap_or_else(syn::Error::into_compile_error).into()
48}
49
50/// `constraint!(model, [name|name[idx]], lhs <op> rhs)`, register a constraint,
51/// an auto-named anonymous constraint, or an indexed family of constraints.
52#[proc_macro]
53pub fn constraint(input: TokenStream) -> TokenStream {
54    constraint::expand(input.into()).unwrap_or_else(syn::Error::into_compile_error).into()
55}
56
57/// `objective!(model, Min|Max, expr)`, set the model objective and sense.
58#[proc_macro]
59pub fn objective(input: TokenStream) -> TokenStream {
60    objective::expand(input.into()).unwrap_or_else(syn::Error::into_compile_error).into()
61}
62
63/// `sum!(body for pat in domain[, pat in domain ...])`, algebraic summation,
64/// lowered to nested `sum_over` folds.
65#[proc_macro]
66pub fn sum(input: TokenStream) -> TokenStream {
67    sum::expand(input.into()).unwrap_or_else(syn::Error::into_compile_error).into()
68}
69
70/// `param!(model, name = value)`, declare a re-bindable scalar parameter and
71/// bind it to a local of the same name.
72#[proc_macro]
73pub fn param(input: TokenStream) -> TokenStream {
74    param::expand(input.into()).unwrap_or_else(syn::Error::into_compile_error).into()
75}
76
77/// `set!(name = domain)`, bind a local to an index `Set`. A plain right side
78/// (`0..5`, `a * b`) is normalized to an owned set (a top-level `*` is a borrowing
79/// Cartesian product). A `pat in domain[ if cond]` comprehension builds (and
80/// optionally filters) the set. See the crate docs.
81#[proc_macro]
82pub fn set(input: TokenStream) -> TokenStream {
83    set::expand(input.into()).unwrap_or_else(syn::Error::into_compile_error).into()
84}
85
86// ---------------------------------------------------------------------------
87// Shared token-walking helpers. The macros must accept forms that are not valid
88// `syn::Expr` (indexed `name[i in set]`, chained `lb <= x <= ub`), so a few
89// splits are done at the raw token-tree level.
90// ---------------------------------------------------------------------------
91
92/// Relational operator recognized inside `constraint!`/`variable!`.
93#[derive(Copy, Clone, PartialEq, Eq)]
94enum RelOp {
95    Le,
96    Ge,
97    Eq,
98}
99
100impl RelOp {
101    /// The `Relate` method this operator maps to.
102    fn method(self) -> Ident {
103        let name = match self {
104            RelOp::Le => "le",
105            RelOp::Ge => "ge",
106            RelOp::Eq => "eq",
107        };
108        Ident::new(name, proc_macro2::Span::call_site())
109    }
110}
111
112/// Split a token stream on top-level commas.
113fn split_top_commas(ts: TokenStream2) -> Vec<TokenStream2> {
114    let mut out = Vec::new();
115    let mut cur = Vec::new();
116    for tt in ts {
117        if let TokenTree::Punct(p) = &tt {
118            if p.as_char() == ',' {
119                out.push(cur.drain(..).collect());
120                continue;
121            }
122        }
123        cur.push(tt);
124    }
125    out.push(cur.into_iter().collect());
126    out
127}
128
129/// Split a token stream on top-level relational operators (`==`, `<=`, `>=`),
130/// returning the intervening segments and the operators between them.
131fn split_relops(ts: TokenStream2) -> (Vec<TokenStream2>, Vec<RelOp>) {
132    let tts: Vec<TokenTree> = ts.into_iter().collect();
133    let mut segs: Vec<TokenStream2> = Vec::new();
134    let mut ops: Vec<RelOp> = Vec::new();
135    let mut cur: Vec<TokenTree> = Vec::new();
136
137    let mut i = 0;
138    while i < tts.len() {
139        if let TokenTree::Punct(p1) = &tts[i] {
140            if p1.spacing() == Spacing::Joint && i + 1 < tts.len() {
141                if let TokenTree::Punct(p2) = &tts[i + 1] {
142                    let op = match (p1.as_char(), p2.as_char()) {
143                        ('<', '=') => Some(RelOp::Le),
144                        ('>', '=') => Some(RelOp::Ge),
145                        ('=', '=') => Some(RelOp::Eq),
146                        _ => None,
147                    };
148                    if let Some(op) = op {
149                        segs.push(cur.drain(..).collect());
150                        ops.push(op);
151                        i += 2;
152                        continue;
153                    }
154                }
155            }
156        }
157        cur.push(tts[i].clone());
158        i += 1;
159    }
160    segs.push(cur.into_iter().collect());
161    (segs, ops)
162}
163
164/// A parsed `name` or `name[binds]` "core" of a `variable!`/`constraint!`
165/// declaration. `cond` holds an optional `if` filter on the index family.
166struct Named {
167    name: Ident,
168    binds: Option<Vec<IndexBind>>,
169    cond: Option<syn::Expr>,
170}
171
172/// Parse a `name`/`name[i in dom, ...]` core out of a token segment.
173fn parse_named(seg: TokenStream2) -> syn::Result<Named> {
174    let tts: Vec<TokenTree> = seg.into_iter().collect();
175    let span = tts.first().map_or_else(proc_macro2::Span::call_site, TokenTree::span);
176    let TokenTree::Ident(name) = tts
177        .first()
178        .cloned()
179        .ok_or_else(|| syn::Error::new(span, "expected a variable/constraint name identifier"))?
180    else {
181        return Err(syn::Error::new(span, "expected a name identifier"));
182    };
183
184    let (binds, cond) = match tts.get(1) {
185        None => (None, None),
186        Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
187            let parsed: Binds = syn::parse2(g.stream())?;
188            if parsed.binds.is_empty() {
189                return Err(syn::Error::new(
190                    g.span(),
191                    "index family needs at least one binding, e.g. `name[i in domain]`",
192                ));
193            }
194            (Some(parsed.binds), parsed.cond)
195        }
196        Some(other) => {
197            return Err(syn::Error::new(other.span(), "expected `[index in domain, ...]`"));
198        }
199    };
200    Ok(Named { name, binds, cond })
201}
202
203/// Build an owned `Set` token expression from one or more index bindings.
204fn build_set(binds: &[IndexBind], root: &TokenStream2) -> TokenStream2 {
205    let mut iter = binds.iter().map(|b| {
206        let dom = &b.domain;
207        quote!(#root::__macro_support::as_set(&(#dom)))
208    });
209    let first = iter.next().expect("at least one index binding");
210    iter.fold(first, |acc, s| quote!(#root::__macro_support::product(&(#acc), &(#s))))
211}
oximo_macros/lib.rs

oximo_macros/
lib.rs
