#![forbid(unsafe_code)]
use proc_macro::{Delimiter, Group, Ident, Punct, Spacing, TokenStream, TokenTree};
#[proc_macro_attribute]
pub fn odin_union(arguments: TokenStream, item: TokenStream) -> TokenStream {
if !arguments.is_empty() {
return compile_error("`odin_union` does not accept arguments");
}
expand(item).unwrap_or_else(|message| compile_error(&message))
}
#[derive(Debug)]
struct Variant {
name: Ident,
conditional_attributes: TokenStream,
}
fn expand(item: TokenStream) -> Result<TokenStream, String> {
let mut item_tokens: Vec<TokenTree> = item.into_iter().collect();
let enum_index = item_tokens
.iter()
.position(|token| is_ident(token, "enum"))
.ok_or_else(|| "`odin_union` can only be applied to an enum".to_owned())?;
let name_index = enum_index + 1;
let enum_name = match item_tokens.get(name_index) {
Some(TokenTree::Ident(name)) => name.clone(),
_ => return Err("expected an enum name after `enum`".to_owned()),
};
let body_index = item_tokens
.iter()
.enumerate()
.skip(name_index + 1)
.find_map(|(index, token)| match token {
TokenTree::Group(group) if group.delimiter() == Delimiter::Brace => Some(index),
_ => None,
})
.ok_or_else(|| "expected an enum body".to_owned())?;
if body_index != name_index + 1 {
return Err(
"generic enums are not supported because payload type arguments cannot be inferred"
.to_owned(),
);
}
let original_body = match &item_tokens[body_index] {
TokenTree::Group(group) => group,
_ => unreachable!("the enum body index always points to a group"),
};
let (expanded_body, variants) = expand_variants(original_body.stream())?;
let mut replacement_body = Group::new(Delimiter::Brace, expanded_body);
replacement_body.set_span(original_body.span());
item_tokens[body_index] = TokenTree::Group(replacement_body);
let mut output: TokenStream = item_tokens.into_iter().collect();
for variant in variants {
output.extend(from_implementation(&enum_name, &variant)?);
}
Ok(output)
}
fn expand_variants(body: TokenStream) -> Result<(TokenStream, Vec<Variant>), String> {
let variant_groups = split_variants(body)?;
let mut expanded_body = TokenStream::new();
let mut variants = Vec::with_capacity(variant_groups.len());
for variant_tokens in variant_groups {
let (attributes, conditional_attributes, name) = parse_variant(variant_tokens)?;
expanded_body.extend(attributes);
expanded_body.extend([TokenTree::Ident(name.clone())]);
let payload: TokenStream = [TokenTree::Ident(name.clone())].into_iter().collect();
let mut payload_group = Group::new(Delimiter::Parenthesis, payload);
payload_group.set_span(name.span());
expanded_body.extend([TokenTree::Group(payload_group)]);
let mut comma = Punct::new(',', Spacing::Alone);
comma.set_span(name.span());
expanded_body.extend([TokenTree::Punct(comma)]);
variants.push(Variant {
name,
conditional_attributes,
});
}
Ok((expanded_body, variants))
}
fn split_variants(body: TokenStream) -> Result<Vec<Vec<TokenTree>>, String> {
let mut variants = Vec::new();
let mut current = Vec::new();
for token in body {
if matches!(&token, TokenTree::Punct(punctuation) if punctuation.as_char() == ',') {
if current.is_empty() {
return Err("expected a variant before `,`".to_owned());
}
variants.push(current);
current = Vec::new();
} else {
current.push(token);
}
}
if !current.is_empty() {
variants.push(current);
}
Ok(variants)
}
fn parse_variant(tokens: Vec<TokenTree>) -> Result<(TokenStream, TokenStream, Ident), String> {
let mut index = 0;
let mut attributes = TokenStream::new();
let mut conditional_attributes = TokenStream::new();
while matches!(tokens.get(index), Some(TokenTree::Punct(punctuation)) if punctuation.as_char() == '#')
{
let hash = tokens[index].clone();
let attribute = match tokens.get(index + 1) {
Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Bracket => {
group.clone()
}
_ => return Err("expected an attribute after `#`".to_owned()),
};
attributes.extend([hash.clone(), TokenTree::Group(attribute.clone())]);
if attribute_is_cfg(&attribute) {
conditional_attributes.extend([hash, TokenTree::Group(attribute)]);
}
index += 2;
}
let name = match tokens.get(index) {
Some(TokenTree::Ident(name)) => name.clone(),
_ => return Err("expected a bare enum variant".to_owned()),
};
if index + 1 != tokens.len() {
return Err(format!(
"variant `{name}` must be bare; write `{name}`, not `{name}(...)`, `{name} {{ ... }}`, or `{name} = ...`"
));
}
Ok((attributes, conditional_attributes, name))
}
fn attribute_is_cfg(attribute: &Group) -> bool {
matches!(
attribute.stream().into_iter().next(),
Some(TokenTree::Ident(name)) if name.to_string() == "cfg"
)
}
fn from_implementation(enum_name: &Ident, variant: &Variant) -> Result<TokenStream, String> {
let name = &variant.name;
let implementation = format!(
"impl ::core::convert::From<{name}> for {enum_name} {{
fn from(value: {name}) -> Self {{
Self::{name}(value)
}}
}}"
);
let implementation: TokenStream = implementation
.parse()
.map_err(|_| format!("failed to generate `From<{name}> for {enum_name}`"))?;
let mut output = variant.conditional_attributes.clone();
output.extend(implementation);
Ok(output)
}
fn is_ident(token: &TokenTree, expected: &str) -> bool {
matches!(token, TokenTree::Ident(identifier) if identifier.to_string() == expected)
}
fn compile_error(message: &str) -> TokenStream {
format!("compile_error!({message:?});")
.parse()
.expect("a quoted string always produces a valid `compile_error!` invocation")
}