use proc_macro2::{Spacing, TokenStream as TokenStream2, TokenTree};
use quote::quote;
use crate::bind::{Binds, filtered_set};
use crate::{build_set, oximo_root};
pub(crate) fn expand(input: TokenStream2) -> syn::Result<TokenStream2> {
let tts: Vec<TokenTree> = input.into_iter().collect();
let span = tts.first().map_or_else(proc_macro2::Span::call_site, TokenTree::span);
let TokenTree::Ident(name) =
tts.first().cloned().ok_or_else(|| syn::Error::new(span, "expected `set!(name = ...)`"))?
else {
return Err(syn::Error::new(span, "expected a set name identifier"));
};
match tts.get(1) {
Some(TokenTree::Punct(p)) if p.as_char() == '=' && p.spacing() == Spacing::Alone => {}
_ => {
return Err(syn::Error::new(name.span(), "expected `=` after the set name"));
}
}
let rhs: TokenStream2 = tts[2..].iter().cloned().collect();
if rhs.is_empty() {
return Err(syn::Error::new(name.span(), "expected a set expression after `=`"));
}
let root = oximo_root();
let set = if has_top_level_in(&rhs) {
let binds: Binds = syn::parse2(rhs)?;
if binds.binds.is_empty() {
return Err(syn::Error::new(
name.span(),
"set comprehension needs at least one `pat in domain`",
));
}
let built = build_set(&binds.binds, &root);
filtered_set(built, &binds.binds, binds.cond.as_ref(), &root)
} else {
let mut iter = split_top_mul(rhs)
.into_iter()
.map(|seg| quote!( #root::__macro_support::as_set(&(#seg)) ));
let first = iter.next().expect("non-empty rhs yields at least one operand");
iter.fold(first, |acc, seg| quote!( #root::__macro_support::product(&(#acc), &(#seg)) ))
};
Ok(quote!( let #name = #set; ))
}
fn split_top_mul(ts: TokenStream2) -> Vec<TokenStream2> {
let mut out = Vec::new();
let mut cur = Vec::new();
for tt in ts {
if matches!(&tt, TokenTree::Punct(p) if p.as_char() == '*') {
out.push(cur.drain(..).collect());
continue;
}
cur.push(tt);
}
out.push(cur.into_iter().collect());
out
}
fn has_top_level_in(ts: &TokenStream2) -> bool {
ts.clone().into_iter().any(|tt| matches!(&tt, TokenTree::Ident(id) if *id == "in"))
}