edlc_layout 0.2.16

/*
 *     EDLc, a compiler for the EDL programming language.
 *     Copyright (C) 2026  Adrian Paskert
 *
 *     This program is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU Affero 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 Affero General Public License for more details.
 *
 *     You should have received a copy of the GNU Affero General Public License
 *     along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#![feature(offset_of_enum)]


extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::{Ident, Literal};
use quote::{quote, quote_spanned};
use syn::{parse_macro_input, spanned::Spanned, Data, DeriveInput, Fields};


#[proc_macro_derive(MirLayout)]
pub fn derive_mir_layout(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;
    // let name_str = Literal::string(&name.to_string());

    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    let layout = layout_from_type(&name, &input.data);

    let expanded = quote! {
        impl #impl_generics ::edlc_core::prelude::mir_type::layout::MirLayout for #name #ty_generics #where_clause {
            fn layout(types: &::edlc_core::prelude::mir_type::MirTypeRegistry) -> ::edlc_core::prelude::mir_type::layout::Layout {
                #layout
            }
        }
    };

    TokenStream::from(expanded)
}

fn layout_from_type(struct_name: &Ident, data: &Data) -> proc_macro2::TokenStream {
    match data {
        Data::Struct(data) => {
            let layout = process_struct_fields(&data.fields, struct_name);
            quote! {
                #[allow(unused_mut)]
                let mut builder = ::edlc_core::prelude::mir_type::layout::OffsetStructLayoutBuilder::default();
                #layout
                ::edlc_core::prelude::mir_type::layout::OffsetStructLayoutBuilder::make::<Self>(builder, types)
            }
        },
        Data::Enum(data) => {
            let variants = data.variants.iter().map(|variant| {
                let variant_name = &variant.ident;
                let variant_name_str = Literal::string(&variant_name.to_string());

                let v = process_enum_fields(&variant.fields, struct_name, variant_name);
                quote_spanned! { variant.span() => ::edlc_core::prelude::mir_type
                    ::layout::EnumLayoutBuilder::add_variant(&mut enum_builder, #variant_name_str.to_string(), {
                        #[allow(unused_mut)]
                        let mut builder = ::edlc_core::prelude::mir_type::layout::OffsetStructLayoutBuilder::default();
                        #v
                        ::edlc_core::prelude::mir_type::layout::OffsetStructLayoutBuilder::make_unchecked(builder, types)
                    });
                }
            });

            quote! {
                #[allow(unused_mut)]
                let mut enum_builder = ::edlc_core::prelude::mir_type::layout::EnumLayoutBuilder::new(types.u8());
                #(#variants)*
                ::edlc_core::prelude::mir_type::layout::EnumLayoutBuilder::make::<Self>(enum_builder, types)
            }
        }
        Data::Union(_data) => {
            todo!()
        }
    }
}

fn process_struct_fields(fields: &Fields, data_ty: &Ident) -> proc_macro2::TokenStream {
    match fields {
        Fields::Named(fields) => {
            let values = fields.named
                .iter()
                .map(|field| {
                    let field_name = field.ident.as_ref().unwrap();
                    let field_name_str = Literal::string(&field_name.to_string());
                    let field_ty = &field.ty;

                    quote_spanned! { field.span() => {
                        let offset = std::mem::offset_of!(#data_ty, #field_name);

                        ::edlc_core::prelude::mir_type
                            ::layout::OffsetStructLayoutBuilder::add_type::<#field_ty>(
                            &mut builder, #field_name_str.to_string(), types, offset);
                    }}
            });

            quote! {
                #(#values)*
            }
        }
        Fields::Unit => quote! {},
        Fields::Unnamed(fields) => {
            let values = fields.unnamed
                .iter()
                .enumerate()
                .map(|(name, field)| {
                    let field_name = name.to_string();
                    let field_name_ident = Literal::usize_unsuffixed(name);
                    let field_name_str = Literal::string(&field_name);
                    let field_ty = &field.ty;

                    quote_spanned! { field.span() => {
                        let offset = std::mem::offset_of!(#data_ty, #field_name_ident);

                        ::edlc_core::prelude::mir_type
                            ::layout::OffsetStructLayoutBuilder::add_type::<#field_ty>(
                            &mut builder, #field_name_str.to_string(), types, offset);
                    }}
            });

            quote! {
                #(#values)*
            }
        }
    }
}

fn process_enum_fields(fields: &Fields, data_ty: &Ident, variant: &Ident) -> proc_macro2::TokenStream {
    match fields {
        Fields::Named(fields) => {
            let values = fields.named
                .iter()
                .map(|field| {
                    let field_name = field.ident.as_ref().unwrap();
                    let field_name_str = Literal::string(&field_name.to_string());
                    let field_ty = &field.ty;

                    quote_spanned! { field.span() => {
                        let offset = std::mem::offset_of!(#data_ty, #variant.#field_name);

                        ::edlc_core::prelude::mir_type
                            ::layout::OffsetStructLayoutBuilder::add_type::<#field_ty>(
                            &mut builder, #field_name_str.to_string(), types, offset);
                    }}
                });

            quote! {
                #(#values)*
            }
        }
        Fields::Unit => quote! {},
        Fields::Unnamed(fields) => {
            let values = fields.unnamed
                .iter()
                .enumerate()
                .map(|(name, field)| {
                    let field_name = name.to_string();
                    let field_name_ident = Literal::usize_unsuffixed(name);
                    let field_name_str = Literal::string(&field_name);
                    let field_ty = &field.ty;

                    quote_spanned! { field.span() => {
                        let offset = std::mem::offset_of!(#data_ty, #variant.#field_name_ident);

                        ::edlc_core::prelude::mir_type
                            ::layout::OffsetStructLayoutBuilder::add_type::<#field_ty>(
                            &mut builder, #field_name_str.to_string(), types, offset);
                    }}
                });

            quote! {
                #(#values)*
            }
        }
    }
}