codama-korok-visitors 0.9.2

use crate::{CombineTypesVisitor, KorokVisitor};
use codama_attributes::{
    Attribute, Attributes, ErrorDirective, ProgramDirective, TryFromFilter, UnsupportedAttribute,
};
use codama_errors::CodamaResult;
use codama_nodes::{Docs, ErrorNode, Node, ProgramNode};
use codama_syn_helpers::extensions::*;

pub struct SetErrorsVisitor {
    combine_types: CombineTypesVisitor,
    enum_current_discriminator: usize,
}

impl Default for SetErrorsVisitor {
    fn default() -> Self {
        Self {
            combine_types: CombineTypesVisitor::strict(),
            enum_current_discriminator: 0,
        }
    }
}

impl SetErrorsVisitor {
    pub fn new() -> Self {
        Self::default()
    }
}

impl KorokVisitor for SetErrorsVisitor {
    fn visit_enum(&mut self, korok: &mut codama_koroks::EnumKorok) -> CodamaResult<()> {
        // No overrides.
        if korok.node.is_some() {
            return Ok(());
        };

        // Ensure the struct has the `CodamaErrors` attribute.
        if !korok.attributes.has_codama_derive("CodamaErrors") {
            return Ok(());
        };

        // Create a `DefinedTypeNode` from the enum.
        self.combine_types.visit_enum(korok)?;

        // Transform each variant into an `ErrorNode`.
        self.enum_current_discriminator = 0;
        self.visit_children(korok)?;
        self.enum_current_discriminator = 0;

        // Gather all errors in a `ProgramNode`.
        let errors = korok
            .variants
            .iter()
            .filter_map(|variant| match &variant.node {
                Some(Node::Error(error)) => Some(error.clone()),
                _ => None,
            })
            .collect::<Vec<_>>();

        let node: Node = ProgramNode {
            errors,
            ..ProgramNode::default()
        }
        .into();
        korok.node = Some(ProgramDirective::apply(&korok.attributes, node));

        Ok(())
    }

    fn visit_enum_variant(
        &mut self,
        korok: &mut codama_koroks::EnumVariantKorok,
    ) -> CodamaResult<()> {
        // Update current discriminator.
        let current_discriminator = match &korok.ast.discriminant {
            Some((_, expr)) => expr.as_unsigned_integer()?,
            _ => self.enum_current_discriminator,
        };
        self.enum_current_discriminator = current_discriminator + 1;

        // Skip variants with #[codama(skip)] directive.
        if korok.attributes.has_codama_attribute("skip") {
            return Ok(());
        };

        // Get #[codama(error)] attribute.
        let codama_error = korok
            .attributes
            .get_first(ErrorDirective::filter)
            .cloned()
            .unwrap_or_default();

        // Get the message from `#[codama(error)]` and then `#[error]` if available.
        let message = codama_error
            .message
            .or(get_message_from_thiserror(&korok.attributes))
            .unwrap_or_default();

        // Get the code from `#[codama(error)]` and then the current discriminator.
        let code = codama_error.code.unwrap_or(current_discriminator);

        korok.node = Some(
            ErrorNode {
                name: korok.name(),
                code,
                message,
                docs: Docs::default(),
            }
            .into(),
        );

        Ok(())
    }
}

pub fn get_message_from_thiserror(attributes: &Attributes) -> Option<String> {
    attributes.iter().find_map(|attr| {
        // Ensure the attribute is a meta list.
        let Attribute::Unsupported(UnsupportedAttribute {
            ast:
                syn::Attribute {
                    meta: syn::Meta::List(list),
                    ..
                },
        }) = attr
        else {
            return None;
        };

        // Ensure the path is `#[error("...")]`.
        if !list.path.is("thiserror::error") {
            return None;
        };

        // Get the first meta as a string, if possible.
        let metas = list.parse_metas().ok()?;
        metas.first()?.as_expr().ok()?.as_string().ok()
    })
}