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

oximo_macros/
lib.rs
