odin_union/

lib.rs

// Copyright (C) 2026  Sisyphus1813
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#![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")
}