nix-js/nix-js-macros/src/ir.rs

//! Implements the `ir!` procedural macro.
//!
//! This macro is designed to reduce the boilerplate associated with defining
//! an Intermediate Representation (IR) that follows a specific pattern. It generates:
//! 1. An enum representing the different kinds of IR nodes.
//! 2. Structs for each of the variants that have fields.
//! 3. `Ref` and `Mut` versions of the main enum for ergonomic pattern matching on references.
//! 4. `From` implementations to easily convert from a struct variant (e.g., `BinOp`) to the main enum (`Ir::BinOp`).
//! 5. A `To[IrName]` trait to provide a convenient `.to_ir()` method on the variant structs.

use convert_case::{Case, Casing};
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{
    FieldsNamed, Ident, Token, Type, parenthesized,
    parse::{Parse, ParseStream, Result},
    punctuated::Punctuated,
    token,
};

/// Represents one of the variants passed to the `ir!` macro.
pub enum VariantInput {
    /// A unit-like variant, e.g., `Arg`.
    Unit(Ident),
    /// A tuple-like variant with one unnamed field, e.g., `ExprRef(ExprId)`.
    Tuple(Ident, Type),
    /// A struct-like variant with named fields, e.g., `BinOp { lhs: ExprId, rhs: ExprId, kind: BinOpKind }`.
    Struct(Ident, FieldsNamed),
}

/// The top-level input for the `ir!` macro.
pub struct MacroInput {
    /// The name of the main IR enum to be generated (e.g., `Ir`).
    pub base_name: Ident,
    /// The list of variants for the enum.
    pub variants: Punctuated<VariantInput, Token![,]>,
}

impl Parse for VariantInput {
    fn parse(input: ParseStream) -> Result<Self> {
        let name: Ident = input.parse()?;

        if input.peek(token::Paren) {
            // Parse a tuple-like variant: `Variant(Type)`
            let content;
            parenthesized!(content in input);
            let ty: Type = content.parse()?;

            if !content.is_empty() {
                return Err(content.error("Expected a single type inside parentheses"));
            }

            Ok(VariantInput::Tuple(name, ty))
        } else if input.peek(token::Brace) {
            // Parse a struct-like variant: `Variant { field: Type, ... }`
            let fields: FieldsNamed = input.parse()?;
            Ok(VariantInput::Struct(name, fields))
        } else {
            // Parse a unit-like variant: `Variant`
            Ok(VariantInput::Unit(name))
        }
    }
}

impl Parse for MacroInput {
    fn parse(input: ParseStream) -> Result<Self> {
        // The macro input is expected to be: `IrName, Variant1, Variant2, ...`
        let base_name = input.parse()?;
        input.parse::<Token![,]>()?;
        let variants = Punctuated::parse_terminated(input)?;

        Ok(MacroInput {
            base_name,
            variants,
        })
    }
}

/// The implementation of the `ir!` macro.
pub fn ir_impl(input: TokenStream) -> TokenStream {
    let parsed_input = syn::parse_macro_input!(input as MacroInput);

    let base_name = &parsed_input.base_name;
    let ref_name = format_ident!("{}Ref", base_name);
    let mut_name = format_ident!("{}Mut", base_name);
    let to_trait_name = format_ident!("To{}", base_name);
    let to_trait_fn_name = format_ident!("to_{}", base_name.to_string().to_case(Case::Snake));

    let mut enum_variants = Vec::new();
    let mut struct_defs = Vec::new();
    let mut ref_variants = Vec::new();
    let mut mut_variants = Vec::new();
    let mut as_ref_arms = Vec::new();
    let mut as_mut_arms = Vec::new();
    let mut from_impls = Vec::new();
    let mut to_trait_impls = Vec::new();

    for variant in parsed_input.variants {
        match variant {
            VariantInput::Unit(name) => {
                let inner_type = name.clone();
                enum_variants.push(quote! { #name(#inner_type) });
                ref_variants.push(quote! { #name(&'a #inner_type) });
                mut_variants.push(quote! { #name(&'a mut #inner_type) });
                as_ref_arms.push(quote! { Self::#name(inner) => #ref_name::#name(inner) });
                as_mut_arms.push(quote! { Self::#name(inner) => #mut_name::#name(inner) });
                from_impls.push(quote! {
                    impl From<#inner_type> for #base_name {
                        fn from(val: #inner_type) -> Self { #base_name::#name(val) }
                    }
                });
                to_trait_impls.push(quote! {
                    impl #to_trait_name for #name {
                        fn #to_trait_fn_name(self) -> #base_name { #base_name::from(self) }
                    }
                });
            }
            VariantInput::Tuple(name, ty) => {
                enum_variants.push(quote! { #name(#ty) });
                ref_variants.push(quote! { #name(&'a #ty) });
                mut_variants.push(quote! { #name(&'a mut #ty) });
                as_ref_arms.push(quote! { Self::#name(inner) => #ref_name::#name(inner) });
                as_mut_arms.push(quote! { Self::#name(inner) => #mut_name::#name(inner) });
            }
            VariantInput::Struct(name, fields) => {
                let inner_type = name.clone();
                struct_defs.push(quote! {
                    #[derive(Debug)]
                    pub struct #name #fields
                });
                enum_variants.push(quote! { #name(#inner_type) });
                ref_variants.push(quote! { #name(&'a #inner_type) });
                mut_variants.push(quote! { #name(&'a mut #inner_type) });
                as_ref_arms.push(quote! { Self::#name(inner) => #ref_name::#name(inner) });
                as_mut_arms.push(quote! { Self::#name(inner) => #mut_name::#name(inner) });
                from_impls.push(quote! {
                    impl From<#inner_type> for #base_name {
                        fn from(val: #inner_type) -> Self { #base_name::#name(val) }
                    }
                });
                to_trait_impls.push(quote! {
                    impl #to_trait_name for #name {
                        fn #to_trait_fn_name(self) -> #base_name { #base_name::from(self) }
                    }
                });
            }
        }
    }

    // Assemble the final generated code.
    let expanded = quote! {
        /// The main IR enum, generated by the `ir!` macro.
        #[derive(Debug, IsVariant, Unwrap, TryUnwrap)]
        pub enum #base_name {
            #( #enum_variants ),*
        }

        // The struct definitions for the enum variants.
        #( #struct_defs )*

        /// An immutable reference version of the IR enum.
        #[derive(Debug, IsVariant, Unwrap, TryUnwrap)]
        pub enum #ref_name<'a> {
            #( #ref_variants ),*
        }

        /// A mutable reference version of the IR enum.
        #[derive(Debug, IsVariant, Unwrap, TryUnwrap)]
        pub enum #mut_name<'a> {
            #( #mut_variants ),*
        }

        impl #base_name {
            /// Converts a `&Ir` into a `IrRef`.
            pub fn as_ref(&self) -> #ref_name<'_> {
                match self {
                    #( #as_ref_arms ),*
                }
            }

            /// Converts a `&mut Ir` into a `IrMut`.
            pub fn as_mut(&mut self) -> #mut_name<'_> {
                match self {
                    #( #as_mut_arms ),*
                }
            }
        }

        // `From` implementations for converting variant structs into the main enum.
        #( #from_impls )*

        /// A trait for converting a variant struct into the main IR enum.
        pub trait #to_trait_name {
            /// Performs the conversion.
            fn #to_trait_fn_name(self) -> #base_name;
        }

        // Implement the `ToIr` trait for each variant struct.
        #( #to_trait_impls )*
    };

    TokenStream::from(expanded)
}