refactor: reduce coupling

evaluator/nixjit_macros/src/ir.rs

@@ -0,0 +1,54 @@
+use syn::{
+    FieldsNamed, Ident, Token, Type, parenthesized,
+    parse::{Parse, ParseStream, Result},
+    punctuated::Punctuated,
+    token,
+};
+
+pub enum VariantInput {
+    Unit(Ident),
+    Tuple(Ident, Type),
+    Struct(Ident, FieldsNamed),
+}
+
+pub struct MacroInput {
+    pub base_name: Ident,
+    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) {
+            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) {
+            let fields: FieldsNamed = input.parse()?;
+
+            Ok(VariantInput::Struct(name, fields))
+        } else {
+            Ok(VariantInput::Unit(name))
+        }
+    }
+}
+
+impl Parse for MacroInput {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let base_name = input.parse()?;
+        input.parse::<Token![,]>()?;
+        let variants = Punctuated::parse_terminated(input)?;
+
+        Ok(MacroInput {
+            base_name,
+            variants,
+        })
+    }
+}

evaluator/nixjit_macros/src/lib.rs

@@ -0,0 +1,121 @@
+use convert_case::{Case, Casing};
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+
+mod ir;
+
+#[proc_macro]
+pub fn ir(input: TokenStream) -> TokenStream {
+    use ir::*;
+    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) }
+                    }
+                });
+            }
+        }
+    }
+
+    let expanded = quote! {
+        #[derive(Debug, IsVariant, Unwrap, TryUnwrap)]
+        pub enum #base_name {
+            #( #enum_variants ),*
+        }
+
+        #( #struct_defs )*
+
+        #[derive(Debug, IsVariant, Unwrap, TryUnwrap)]
+        pub enum #ref_name<'a> {
+            #( #ref_variants ),*
+        }
+
+        #[derive(Debug, IsVariant, Unwrap, TryUnwrap)]
+        pub enum #mut_name<'a> {
+            #( #mut_variants ),*
+        }
+
+        impl #base_name {
+            pub fn as_ref(&self) -> #ref_name<'_> {
+                match self {
+                    #( #as_ref_arms ),*
+                }
+            }
+
+            pub fn as_mut(&mut self) -> #mut_name<'_> {
+                match self {
+                    #( #as_mut_arms ),*
+                }
+            }
+        }
+
+        #( #from_impls )*
+
+        pub trait #to_trait_name {
+            fn #to_trait_fn_name(self) -> #base_name;
+        }
+
+        #( #to_trait_impls )*
+    };
+
+    TokenStream::from(expanded)
+}