bitflags2-derive 0.2.0

use std::collections::HashMap;

use proc_macro2::Span;
use syn::punctuated::Punctuated;
use syn::token::Comma;
use syn::{
    Attribute, BinOp, Error, Expr, ExprLit, ExprPath, ItemEnum, Lit, Meta, Path, Result, Visibility,
};

use crate::model::{FlagDirective, FlagVariant};

/// Derives implemented by the macro itself; user-provided duplicates are dropped.
const RESERVED_DERIVES: &[&str] = &["Copy", "Clone", "PartialEq", "Eq", "Debug"];

/// A fully parsed `#[flags]` input enum.
pub(crate) struct FlagsInput {
    pub(crate) vis: Visibility,
    pub(crate) ident: syn::Ident,
    pub(crate) variants: Vec<FlagVariant>,
    /// User-provided derives forwarded to the generated type (e.g. `Serialize`).
    pub(crate) forwarded_derives: Vec<Path>,
    /// Other container attributes forwarded as-is (e.g. `#[serde(...)]`, docs).
    pub(crate) forwarded_attrs: Vec<Attribute>,
}

struct RawFlagVariant {
    ident: syn::Ident,
    directive: Option<FlagDirective>,
    discriminant_expr: Option<Expr>,
}

#[derive(Clone, Copy)]
enum ResolveState {
    Unresolved,
    Resolving,
    Resolved(u128),
}

struct Resolver {
    raw_variants: Vec<RawFlagVariant>,
    indexes: HashMap<String, usize>,
    states: Vec<ResolveState>,
}

/// Parses the source enum and resolves every flag value.
pub(crate) fn parse_flags(input: ItemEnum) -> Result<FlagsInput> {
    if !input.generics.params.is_empty() || input.generics.where_clause.is_some() {
        return Err(Error::new_spanned(
            input.generics,
            "flags enum cannot be generic",
        ));
    }

    let vis = input.vis;
    let ident = input.ident;
    let (forwarded_derives, forwarded_attrs) = split_container_attrs(input.attrs)?;
    let mut raw_variants = Vec::new();

    for variant in input.variants {
        if !variant.fields.is_empty() {
            return Err(Error::new_spanned(
                variant.fields,
                "flags variants cannot have fields",
            ));
        }

        let flag_attr = flag_attr(&variant.attrs)?;
        let Some(flag_attr) = flag_attr else {
            return Err(Error::new_spanned(
                variant.ident,
                "flags enum variants must have #[flag] or #[flag(value)]",
            ));
        };

        raw_variants.push(RawFlagVariant {
            ident: variant.ident,
            directive: parse_flag_attr(flag_attr)?,
            discriminant_expr: variant.discriminant.map(|(_, expr)| expr),
        });
    }

    let ignored_count = raw_variants
        .iter()
        .filter(|variant| matches!(variant.directive, Some(FlagDirective::Ignore)))
        .count();
    if ignored_count == raw_variants.len() {
        return Err(Error::new(
            Span::call_site(),
            "flags enum must have at least one non-ignored variant",
        ));
    }

    let mut resolver = Resolver::new(raw_variants)?;
    let mut variants = Vec::new();

    for index in 0..resolver.raw_variants.len() {
        if matches!(
            resolver.raw_variants[index].directive,
            Some(FlagDirective::Ignore)
        ) {
            continue;
        }
        let directive = resolver.raw_variants[index].directive.clone();
        let ident = resolver.raw_variants[index].ident.clone();
        let value = resolver.resolve_variant(index)?;
        variants.push(FlagVariant {
            ident,
            directive: directive.unwrap_or(FlagDirective::Auto),
            value: Some(value),
        });
    }

    Ok(FlagsInput {
        vis,
        ident,
        variants,
        forwarded_derives,
        forwarded_attrs,
    })
}

/// Splits the enum's outer attributes into forwarded derives and other attributes.
///
/// Derives implemented directly by this macro (`Copy`, `Clone`, `PartialEq`,
/// `Eq`, `Debug`) are dropped to avoid conflicting impls; everything else is
/// forwarded to the generated newtype so that, for example, `#[derive(Serialize,
/// Deserialize)]` keeps working on the flag type directly.
fn split_container_attrs(attrs: Vec<Attribute>) -> Result<(Vec<Path>, Vec<Attribute>)> {
    let mut derives = Vec::new();
    let mut others = Vec::new();

    for attr in attrs {
        if attr.path().is_ident("derive") {
            let paths = attr.parse_args_with(Punctuated::<Path, Comma>::parse_terminated)?;
            for path in paths {
                if !is_reserved_derive(&path) {
                    derives.push(path);
                }
            }
        } else {
            others.push(attr);
        }
    }

    Ok((derives, others))
}

/// Returns true for derives the macro implements directly and must not forward.
fn is_reserved_derive(path: &Path) -> bool {
    let Some(last) = path.segments.last() else {
        return false;
    };
    RESERVED_DERIVES.contains(&last.ident.to_string().as_str())
}

impl Resolver {
    fn new(raw_variants: Vec<RawFlagVariant>) -> Result<Self> {
        let mut indexes = HashMap::new();

        for (index, variant) in raw_variants.iter().enumerate() {
            if indexes.insert(variant.ident.to_string(), index).is_some() {
                return Err(Error::new_spanned(
                    &variant.ident,
                    "duplicate flag variant name",
                ));
            }
        }

        let states = vec![ResolveState::Unresolved; raw_variants.len()];

        Ok(Self {
            raw_variants,
            indexes,
            states,
        })
    }

    fn resolve_variant(&mut self, index: usize) -> Result<u128> {
        if let Some(FlagDirective::Ignore) = self.raw_variants[index].directive {
            return Ok(0);
        }
        match self.states[index] {
            ResolveState::Resolved(value) => return Ok(value),
            ResolveState::Resolving => {
                return Err(Error::new_spanned(
                    &self.raw_variants[index].ident,
                    "cyclic flag value reference",
                ));
            }
            ResolveState::Unresolved => {}
        }

        self.states[index] = ResolveState::Resolving;

        let directive = self.raw_variants[index].directive.clone();
        let discriminant_expr = self.raw_variants[index].discriminant_expr.clone();
        let span = self.raw_variants[index].ident.span();

        let explicit_value = match directive {
            Some(FlagDirective::Auto) | None => None,
            Some(FlagDirective::Value(expr)) => Some(self.eval_flag_expr(&expr)?),
            Some(FlagDirective::Ignore) => unreachable!(),
        };
        let discriminant_value = match discriminant_expr.as_ref() {
            Some(expr) => Some(parse_int_expr(expr)?),
            None => None,
        };

        let value = match (explicit_value, discriminant_value) {
            (Some(attr), Some(discriminant)) if attr != discriminant => {
                return Err(Error::new(
                    span,
                    "#[flag(value)] and discriminant value differ",
                ));
            }
            (Some(attr), _) => attr,
            (None, Some(discriminant)) => discriminant,
            (None, None) => {
                let previous = if index == 0 {
                    None
                } else {
                    Some(self.resolve_variant(index - 1)?)
                };
                next_auto_value(previous, span)?
            }
        };

        self.states[index] = ResolveState::Resolved(value);
        Ok(value)
    }

    fn eval_flag_expr(&mut self, expr: &Expr) -> Result<u128> {
        match expr {
            Expr::Lit(ExprLit {
                lit: Lit::Int(lit), ..
            }) => lit.base10_parse::<u128>(),
            Expr::Path(path) => self.eval_path(path),
            Expr::Binary(binary) if matches!(binary.op, BinOp::BitOr(_)) => {
                Ok(self.eval_flag_expr(&binary.left)? | self.eval_flag_expr(&binary.right)?)
            }
            Expr::Paren(paren) => self.eval_flag_expr(&paren.expr),
            Expr::Group(group) => self.eval_flag_expr(&group.expr),
            _ => Err(Error::new_spanned(
                expr,
                "expected an integer literal, flag name, or `|` expression",
            )),
        }
    }

    fn eval_path(&mut self, path: &ExprPath) -> Result<u128> {
        if path.qself.is_some() {
            return Err(Error::new_spanned(path, "expected a flag variant name"));
        }

        let Some(ident) = path.path.get_ident() else {
            return Err(Error::new_spanned(path, "expected a flag variant name"));
        };

        let Some(index) = self.indexes.get(&ident.to_string()).copied() else {
            return Err(Error::new_spanned(ident, "unknown flag variant name"));
        };

        self.resolve_variant(index)
    }
}

fn flag_attr(attrs: &[Attribute]) -> Result<Option<&Attribute>> {
    let mut found = None;

    for attr in attrs {
        if attr.path().is_ident("flag") {
            if found.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[flag] attribute"));
            }
            found = Some(attr);
        }
    }

    Ok(found)
}

fn parse_flag_attr(attr: &Attribute) -> Result<Option<FlagDirective>> {
    match &attr.meta {
        Meta::Path(_) => Ok(Some(FlagDirective::Auto)),
        Meta::List(list) => Ok(Some(parse_flag_expr(list.parse_args::<Expr>()?)?)),
        Meta::NameValue(_) => Err(Error::new_spanned(
            attr,
            "expected #[flag], #[flag(ignore)], or #[flag(value)]",
        )),
    }
}

fn parse_flag_expr(expr: Expr) -> Result<FlagDirective> {
    if let Expr::Path(ExprPath {
        qself: None, path, ..
    }) = &expr
    {
        if path.is_ident("ignore") {
            return Ok(FlagDirective::Ignore);
        }
    }
    Ok(FlagDirective::Value(expr))
}

fn parse_int_expr(expr: &Expr) -> Result<u128> {
    match expr {
        Expr::Lit(ExprLit {
            lit: Lit::Int(lit), ..
        }) => lit.base10_parse::<u128>(),
        _ => Err(Error::new_spanned(expr, "expected integer literal")),
    }
}

fn next_auto_value(previous: Option<u128>, span: Span) -> Result<u128> {
    match previous {
        None => Ok(0),
        Some(0) => Ok(1),
        Some(value) if value.is_power_of_two() => Ok(value << 1),
        Some(_) => Err(Error::new(
            span,
            "cannot auto-assign flag value after a composite value",
        )),
    }
}