bitpiece_macros 2.5.0

use convert_case::Casing;
use quote::{format_ident, quote, ToTokens};
use syn::{DeriveInput, Field, FieldsNamed};

use crate::{
    newtypes::{BitLenExpr, BitOffsetExpr, StorageTypeExpr, TypeExpr},
    utils::{
        bitpiece_gen_impl, gen_explicit_bit_length_assertion, not_supported_err,
        BitPieceGenImplParams,
    },
    MacroArgs, OptIn,
};

pub fn bitpiece_named_struct(
    input: &DeriveInput,
    fields: &FieldsNamed,
    macro_args: MacroArgs,
) -> proc_macro::TokenStream {
    if fields.named.is_empty() {
        return not_supported_err("empty structs");
    }

    let ident = &input.ident;

    let bit_len_calc = calc_bit_len(fields);
    let bit_len_ident = proc_macro2::Ident::new(
        &format!(
            "{}_BIT_LEN",
            input
                .ident
                .to_string()
                .to_case(convert_case::Case::Constant)
        ),
        input.ident.span(),
    );
    let bit_len = BitLenExpr(bit_len_ident.to_token_stream());

    let storage_type_calc = bit_len.storage_type();
    let storage_type_ident = format_ident!("{}StorageTy", input.ident);
    let storage_type = StorageTypeExpr(storage_type_ident.to_token_stream());

    let fields_struct_ident = format_ident!("{}Fields", input.ident);

    let mut_type_ident = format_ident!("{}MutRef", input.ident);
    let fields_type = TypeExpr(quote! { #fields_struct_ident });

    let fields_offsets_and_lens_consts = gen_fields_offsets_and_lens_consts(ident, fields);

    let explicit_bit_len_assertion =
        gen_explicit_bit_length_assertion(macro_args.explicit_bit_length, &bit_len);

    let bitpiece_impl = bitpiece_gen_impl(BitPieceGenImplParams {
        type_ident: ident,
        type_vis: &input.vis,
        macro_args: &macro_args,
        bit_len: &bit_len,
        storage_type: &storage_type,
        to_bits_code: quote! { self.storage },
        try_from_bits_code: gen_try_from_bits_code(ident, fields, &storage_type),
        mut_type_ident: &mut_type_ident,
        fields_type: &fields_type,
        zeroes: gen_const_instantiation(ident, fields, &storage_type, "ZEROES"),
        ones: gen_const_instantiation(ident, fields, &storage_type, "ONES"),
        min: gen_const_instantiation(ident, fields, &storage_type, "MIN"),
        max: gen_const_instantiation(ident, fields, &storage_type, "MAX"),
        to_fields_code: gen_to_fields(fields, &fields_struct_ident),
        from_fields_code: gen_from_fields(fields, input, &storage_type),
    });

    let field_access_fns = macro_args.filter_opt_in_code(
        OptIn::Get,
        gen_field_access_fns(ident, fields, &storage_type),
    );
    let field_access_noshift_fns = macro_args.filter_opt_in_code(
        OptIn::GetNoshift,
        gen_field_access_noshift_fns(ident, fields, &storage_type),
    );
    let field_with_fns = macro_args.filter_opt_in_code(
        OptIn::With,
        gen_field_with_fns(ident, fields, &storage_type),
    );
    let field_set_fns =
        macro_args.filter_opt_in_code(OptIn::Set, gen_field_set_fns(ident, fields, &storage_type));
    let field_mut_fns = macro_args.filter_opt_in_code(
        OptIn::GetMut,
        gen_field_mut_fns(ident, fields, &storage_type),
    );

    let mut_struct_field_access_fns = macro_args.filter_opt_in_code(
        OptIn::MutStructFieldGet,
        gen_mut_struct_field_access_fns(ident, fields),
    );
    let mut_struct_field_access_noshift_fns = macro_args.filter_opt_in_code(
        OptIn::GetNoshift,
        gen_mut_struct_field_access_noshift_fns(ident, fields),
    );
    let mut_struct_field_set_fns = macro_args.filter_opt_in_code(
        OptIn::MutStructFieldSet,
        gen_mut_struct_field_set_fns(ident, fields),
    );
    let mut_struct_field_mut_fns = macro_args.filter_opt_in_code(
        OptIn::MutStructFieldMut,
        gen_mut_struct_field_mut_fns(ident, fields),
    );

    let vis = &input.vis;
    let attrs = &input.attrs;
    let base_code = quote! {
        #vis const #bit_len_ident: usize = #bit_len_calc;
        #vis type #storage_type_ident = #storage_type_calc;

        #explicit_bit_len_assertion

        #(#attrs)*
        #[derive(Clone, Copy)]
        #[repr(transparent)]
        #vis struct #ident {
            pub storage: #storage_type,
        }

        #bitpiece_impl

        impl #ident {
            #fields_offsets_and_lens_consts
            #field_access_fns
            #field_access_noshift_fns
            #field_with_fns
            #field_set_fns
            #field_mut_fns
        }
    };

    let opt_mut_struct_code = macro_args.filter_opt_in_code(
        OptIn::MutStruct,
        quote! {
            impl<'s> #mut_type_ident<'s> {
                #mut_struct_field_access_fns
                #mut_struct_field_access_noshift_fns
                #mut_struct_field_set_fns
                #mut_struct_field_mut_fns
            }
        },
    );

    let opt_fields_struct_code = macro_args.filter_opt_in_code(
        OptIn::FieldsStruct,
        quote! {
            #(#attrs)*
            #[derive(Clone, Copy)]
            #vis struct #fields_struct_ident #fields

            impl ::core::convert::From<#fields_struct_ident> for #ident {
                fn from(fields: #fields_struct_ident) -> Self {
                    Self::from_fields(fields)
                }
            }
            impl ::core::convert::From<#ident> for #fields_struct_ident {
                fn from(value: #ident) -> Self {
                    value.to_fields()
                }
            }
        },
    );

    quote! {
        #base_code
        #opt_mut_struct_code
        #opt_fields_struct_code
    }
    .into()
}

fn gen_const_instantiation(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
    const_name: &str,
) -> proc_macro2::TokenStream {
    let const_name_ident = syn::Ident::new(const_name, proc_macro2::Span::mixed_site());
    let bitwise_or_each_field = fields.named.iter().map(|f| {
        let field_ty = &f.ty;
        let offset = get_field_offset(type_ident, f);
        quote! {
            | (
                (
                    <#field_ty as ::bitpiece::BitPiece>::Converter::to_bits(
                        <#field_ty as ::bitpiece::BitPiece>::#const_name_ident
                    ) as #storage_type
                ) << #offset
            )
        }
    });
    quote! {
        #type_ident::from_bits(
            0 #(#bitwise_or_each_field)*
        )
    }
}

fn gen_field_const_value_ident(field: &Field, const_name: &str) -> syn::Ident {
    let field_ident = field.ident.as_ref().unwrap();
    let field_name_const_case = field_ident
        .to_string()
        .to_case(convert_case::Case::Constant);
    syn::Ident::new(
        &format!("{}_{}", field_name_const_case, const_name),
        field_ident.span(),
    )
}

fn get_field_offset_const_ident(field: &Field) -> syn::Ident {
    gen_field_const_value_ident(field, "OFFSET")
}

fn get_field_len_const_ident(field: &Field) -> syn::Ident {
    gen_field_const_value_ident(field, "LEN")
}

fn get_field_offset(type_ident: &syn::Ident, field: &Field) -> BitOffsetExpr {
    let const_ident = get_field_offset_const_ident(field);
    BitOffsetExpr(quote! {
        #type_ident::#const_ident
    })
}

fn get_field_len(type_ident: &syn::Ident, field: &Field) -> BitLenExpr {
    let const_ident = get_field_len_const_ident(field);
    BitLenExpr(quote! {
        #type_ident::#const_ident
    })
}

fn gen_fields_offsets_and_lens_consts(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
) -> proc_macro2::TokenStream {
    // iterator over each fields and its previous field. in the first iteration, the previous fields is `None`.
    let fields_with_prev = fields.named.iter().enumerate().map(|(i, field)| {
        let prev = if i == 0 {
            None
        } else {
            Some(&fields.named[i - 1])
        };
        (prev, field)
    });

    fields_with_prev
        .map(|(prev, cur)| {
            let offset_const_ident = get_field_offset_const_ident(cur);
            let len_const_ident = get_field_len_const_ident(cur);
            let offset = match prev {
                Some(prev) => {
                    let prev_offset = get_field_offset(type_ident, prev);
                    let prev_len = get_field_len(type_ident, prev);
                    quote! {
                        (#prev_offset) + (#prev_len)
                    }
                }
                None => quote! { 0 },
            };
            let len = TypeExpr::from_type(&cur.ty).bit_len();
            quote! {
                pub const #len_const_ident: usize = #len;
                pub const #offset_const_ident: usize = #offset;
            }
        })
        .collect()
}

fn calc_bit_len(fields: &FieldsNamed) -> BitLenExpr {
    let field_types = fields
        .named
        .iter()
        .map(|field| TypeExpr::from_type(&field.ty));
    field_types.clone().map(|field_ty| field_ty.bit_len()).sum()
}

/// returns an iterator over the extracted bits of each field.
fn fields_extracted_bits<'a>(
    type_ident: &'a syn::Ident,
    fields: &'a FieldsNamed,
    storage_type: &'a StorageTypeExpr,
    storage_bits_expr: proc_macro2::TokenStream,
) -> impl Iterator<Item = proc_macro2::TokenStream> + 'a {
    fields.named.iter().map(move |field| {
        let len = get_field_len(type_ident, field);
        let offset = get_field_offset(type_ident, field);
        extract_bits(ExtractBitsParams {
            value: storage_bits_expr.clone(),
            value_type: storage_type.clone(),
            extract_offset: offset,
            extract_len: len,
        })
    })
}

/// returns an iterator over the extracted bits (mask only, no shift) of each field.
fn fields_extracted_bits_noshift<'a>(
    type_ident: &'a syn::Ident,
    fields: &'a FieldsNamed,
    storage_type: &'a StorageTypeExpr,
) -> impl Iterator<Item = proc_macro2::TokenStream> + 'a {
    fields.named.iter().map(move |field| {
        let len = get_field_len(type_ident, field);
        let offset = get_field_offset(type_ident, field);
        extract_bits_noshift(ExtractBitsParams {
            value: quote! { self.storage },
            value_type: storage_type.clone(),
            extract_offset: offset,
            extract_len: len,
        })
    })
}

/// parameters for extracting some range of bits from a value
struct ExtractBitsParams {
    /// the value to extract the bits from
    value: proc_macro2::TokenStream,
    /// the type of the value to extract the bits from
    value_type: StorageTypeExpr,
    /// the offset at which to start extracting
    extract_offset: BitOffsetExpr,
    /// the amount of bits to extract
    extract_len: BitLenExpr,
}

/// parameters for modifying some range of bits of a value
struct ModifyBitsParams {
    /// the parameters used for extracting the range of bits to be modified.
    extract_params: ExtractBitsParams,
    /// the new value of the specified bit range.
    new_value: proc_macro2::TokenStream,
}

/// extracts some bits from a value
fn extract_bits(params: ExtractBitsParams) -> proc_macro2::TokenStream {
    let ExtractBitsParams {
        value,
        value_type,
        extract_offset,
        extract_len,
    } = &params;
    quote! {
        (
            ::bitpiece::extract_bits(#value as u64, #extract_offset, #extract_len) as #value_type
        )
    }
}

/// extracts some bits (mask only, no shift) from a value
fn extract_bits_noshift(params: ExtractBitsParams) -> proc_macro2::TokenStream {
    let ExtractBitsParams {
        value,
        value_type,
        extract_offset,
        extract_len,
    } = &params;
    quote! {
        ::bitpiece::extract_bits_noshift(#value as u64, #extract_offset, #extract_len) as #value_type
    }
}

/// returns an expression for the provided value with the specified bit range modified to its new value.
fn modify_bits(params: ModifyBitsParams) -> proc_macro2::TokenStream {
    let ModifyBitsParams {
        extract_params:
            ExtractBitsParams {
                value,
                value_type,
                extract_offset,
                extract_len,
            },
        new_value,
    } = params;
    quote! {
        ::bitpiece::modify_bits(#value as u64, #extract_offset, #extract_len, #new_value as u64) as #value_type
    }
}

fn gen_from_fields(
    fields: &FieldsNamed,
    input: &DeriveInput,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    let type_ident = &input.ident;
    let bitwise_or_each_field = fields.named.iter().map(|f| {
        let field_ty = &f.ty;
        let field_ident = &f.ident;
        let offset = get_field_offset(type_ident, f);
        quote! {
            | (
                (
                    <#field_ty as ::bitpiece::BitPiece>::Converter::to_bits(
                        fields.#field_ident
                    ) as #storage_type
                ) << #offset
            )
        }
    });
    quote! {
        #type_ident::from_bits(
            0 #(#bitwise_or_each_field)*
        )
    }
}

fn gen_try_from_bits_code(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    // before constructing the type, make sure that the values of all fields are valid
    let per_field_call = fields_extracted_bits(type_ident, fields, storage_type, quote! { bits })
        .zip(fields.named.iter())
        .map(|(bits, field)| {
            let ty = &field.ty;
            quote! {
                if <#ty as ::bitpiece::BitPiece>::Converter::try_from_bits(#bits as <#ty as ::bitpiece::BitPiece>::Bits).is_none() {
                    return None;
                }
            }
        });
    quote! {
        let result = Self { storage: bits };
        #(#per_field_call)*
        Some(result)
    }
}

fn gen_to_fields<'a>(
    fields: &'a FieldsNamed,
    fields_struct_ident: &syn::Ident,
) -> proc_macro2::TokenStream {
    let field_initializers = fields.named.iter().map(|field| {
        let field_ident = field.ident.as_ref().unwrap();
        quote! {
            #field_ident: self.#field_ident(),
        }
    });
    quote! {
        #fields_struct_ident {
            #(#field_initializers)*
        }
    }
}

fn gen_field_access_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    fields_extracted_bits(type_ident, fields, storage_type, quote!{ self.storage })
        .zip(fields.named.iter())
        .map(|(bits, field)| {
            let vis = &field.vis;
            let ident = &field.ident;
            let ty = &field.ty;
            quote! {
                #vis const fn #ident (self) -> #ty {
                    <#ty as ::bitpiece::BitPiece>::Converter::from_bits(#bits as <#ty as ::bitpiece::BitPiece>::Bits)
                }
            }
        }).collect()
}

fn gen_field_with_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(|field| {
            let len = get_field_len(type_ident, field);
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ty = &field.ty;
            let ident = field.ident.as_ref().unwrap();
            let with_ident = format_ident!("with_{}", ident);
            let modified_value_expr = modify_bits(ModifyBitsParams {
                extract_params: ExtractBitsParams {
                    value: quote! { self.storage },
                    value_type: storage_type.clone(),
                    extract_offset: offset,
                    extract_len: len,
                },
                new_value: quote! { <#ty as ::bitpiece::BitPiece>::Converter::to_bits(new_value) },
            });

            quote! {
                #vis const fn #with_ident (mut self, new_value: #ty) -> Self {
                    self.storage = #modified_value_expr;
                    self
                }
            }
        })
        .collect()
}

fn gen_field_access_noshift_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    fields_extracted_bits_noshift(type_ident, fields, storage_type)
        .zip(fields.named.iter())
        .map(move |(bits, field)| {
            let vis = &field.vis;
            let ident = field.ident.as_ref().unwrap();
            let ident_noshift = format_ident!("{}_noshift", ident);
            quote! {
                #vis const fn #ident_noshift (self) -> #storage_type {
                    #bits
                }
            }
        })
        .collect()
}

fn gen_mut_struct_field_access_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(|field| {
            let len = get_field_len(type_ident, field);
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ident = &field.ident;
            let ty = &field.ty;
            quote! {
                #vis const fn #ident(&self) -> #ty {
                    <#ty as ::bitpiece::BitPiece>::Converter::from_bits(
                        self.0.get_bits(#offset, #len) as <#ty as ::bitpiece::BitPiece>::Bits
                    )
                }
            }
        })
        .collect()
}

fn gen_mut_struct_field_access_noshift_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(|field| {
            let len = get_field_len(type_ident, field);
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ty = &field.ty;
            let ident = field.ident.as_ref().unwrap();
            let ident_noshift = format_ident!("{}_noshift", ident);
            quote! {
                #vis const fn #ident_noshift(&self) -> #ty {
                    <#ty as ::bitpiece::BitPiece>::Converter::from_bits(
                        self.0.get_bits_noshift(#offset, #len) as <#ty as ::bitpiece::BitPiece>::Bits
                    )
                }
            }
        })
        .collect()
}

fn gen_mut_struct_field_set_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(|field| {
            let len = get_field_len(type_ident, field);
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ident = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            let set_ident = format_ident!("set_{}", ident);
            quote! {
                #vis const fn #set_ident(&mut self, new_value: #ty) {
                    let new_value_bits = <#ty as ::bitpiece::BitPiece>::Converter::to_bits(new_value);
                    self.0.set_bits(#offset, #len, new_value_bits as u64)
                }
            }
        })
        .collect()
}

fn gen_mut_struct_field_mut_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(|field| {
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ident = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            let ident_mut = format_ident!("{}_mut", ident);
            let mut_ty = quote! {
                <#ty as ::bitpiece::BitPieceHasMutRef>::MutRef
            };
            quote! {
                #vis const fn #ident_mut(&'s mut self) -> #mut_ty<'s> {
                    #mut_ty::new(self.0.storage.reborrow(), self.0.start_bit_index + #offset)
                }
            }
        })
        .collect()
}

fn gen_field_set_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(|field| {
            let len = get_field_len(type_ident, field);
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ident = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            let set_ident = format_ident!("set_{}", ident);
            let modified_value_expr = modify_bits(ModifyBitsParams {
                extract_params: ExtractBitsParams {
                    value: quote! { self.storage },
                    value_type: storage_type.clone(),
                    extract_offset: offset,
                    extract_len: len,
                },
                new_value: quote! { <#ty as ::bitpiece::BitPiece>::Converter::to_bits(new_value) },
            });
            quote! {
                #vis const fn #set_ident (&mut self, new_value: #ty) {
                    self.storage = #modified_value_expr;
                }
            }
        })
        .collect()
}

fn gen_field_mut_fns(
    type_ident: &syn::Ident,
    fields: &FieldsNamed,
    storage_type: &StorageTypeExpr,
) -> proc_macro2::TokenStream {
    fields
        .named
        .iter()
        .map(move |field| {
            let offset = get_field_offset(type_ident, field);
            let vis = &field.vis;
            let ident = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            let ident_mut = format_ident!("{}_mut", ident);
            let storage_type = storage_type.clone();
            let mut_ty = quote! {
                <#ty as ::bitpiece::BitPieceHasMutRef>::MutRef
            };
            let storage_mut_ref =
                storage_type.convert_mut_ref_to_storage_mut_ref(quote! { &mut self.storage });
            quote! {
                #vis const fn #ident_mut<'a>(&'a mut self) -> #mut_ty<'a> {
                    #mut_ty::new(#storage_mut_ref, #offset)
                }
            }
        })
        .collect()
}