bitstructs_macro 0.2.1

use std::collections::HashMap;

use crate::{
    BITSTRUCT_REPR,
    tokens::{self, Bitstruct, BitstructField, BitstructFieldEnumType, SupportedPrimitiveType},
};
use proc_macro2::TokenStream as TokenStream2;
use quote::{ToTokens, format_ident, quote};

type Result<T> = std::result::Result<T, syn::Error>;

mod gen_bitstruct;
mod gen_field_accessor;
mod gen_field_enum;

#[derive(Debug, Default)]
struct Cache {
    /// Cache for match arms of enum type values and variants.
    ///
    /// 枚举类型的value和variant的match arms缓存
    ///
    /// Map<Enum type identifier, (value => variant, variant => value)>
    field_enum_type_match_arms: HashMap<syn::Ident, (Vec<TokenStream2>, Vec<TokenStream2>)>,
    /// bitstruct的store type, only used in bitstruct_small
    bitstruct_store_type: Option<TokenStream2>,
}

#[derive(Default, Clone, Copy, PartialEq, Eq)]
enum BitstructGenType {
    /// Generate code with owned
    /// 生成拥有所有权的代码
    #[default]
    WithOwned,
    /// Generate code with Cow
    /// 生成使用Cow的代码
    WithCow,
    /// Generate code for small data structure
    /// 生成用于小型数据结构的代码
    ForSmall,
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
enum BitNumbering {
    MSB0,
    #[default]
    LSB0,
}

#[derive(Default)]
struct BistructGenOptions {
    gen_type: BitstructGenType,
    bit_numbering: BitNumbering,
    given_store_type: Option<syn::Ident>,
}
impl BistructGenOptions {
    fn with_cow(&self) -> bool {
        match self.gen_type {
            BitstructGenType::WithCow => true,
            _ => false,
        }
    }
    fn for_small(&self) -> bool {
        match self.gen_type {
            BitstructGenType::ForSmall => true,
            _ => false,
        }
    }
}

pub fn gen_code_with_owned(bitstruct: Bitstruct) -> Result<TokenStream2> {
    let mut cache = Cache::default();
    let mut options = BistructGenOptions {
        gen_type: BitstructGenType::WithOwned,
        ..Default::default()
    };
    check_and_parse_bitstruct_repr_attr(&mut options, &bitstruct)?;
    check_fields_of_primitive_type(&bitstruct)?;

    let struct_def = gen_bitstruct::gen_struct_def(&bitstruct)?;
    let struct_field_enum_types_def =
        gen_field_enum::gen_bitstruct_field_enum_types_def(&mut cache, &bitstruct)?;
    let field_accessors =
        gen_field_accessor::gen_field_accessors(&mut cache, &bitstruct, &options)?;
    Ok(quote! {
        #struct_def
        #struct_field_enum_types_def
        #field_accessors
    })
}

pub fn gen_code_with_cow(bitstruct: Bitstruct) -> Result<TokenStream2> {
    let mut cache = Cache::default();
    let mut options = BistructGenOptions {
        gen_type: BitstructGenType::WithCow,
        ..Default::default()
    };
    check_and_parse_bitstruct_repr_attr(&mut options, &bitstruct)?;
    check_fields_of_primitive_type(&bitstruct)?;

    let struct_def = gen_bitstruct::gen_struct_def_with_cow(&bitstruct)?;
    let struct_field_enum_types_def =
        gen_field_enum::gen_bitstruct_field_enum_types_def(&mut cache, &bitstruct)?;
    let field_accessors =
        gen_field_accessor::gen_field_accessors(&mut cache, &bitstruct, &options)?;
    Ok(quote! {
        #struct_def
        #struct_field_enum_types_def
        #field_accessors
    })
}

pub fn gen_code_for_small(bitstruct: Bitstruct) -> Result<TokenStream2> {
    let mut cache = Cache::default();
    let mut options = BistructGenOptions {
        gen_type: BitstructGenType::ForSmall,
        ..Default::default()
    };
    check_and_parse_bitstruct_repr_attr(&mut options, &bitstruct)?;
    check_fields_of_primitive_type(&bitstruct)?;

    let struct_def = gen_bitstruct::gen_struct_def_for_small(&mut cache, &options, &bitstruct)?;
    let struct_field_enum_types_def =
        gen_field_enum::gen_bitstruct_field_enum_types_def(&mut cache, &bitstruct)?;
    let field_accessors =
        gen_field_accessor::gen_field_accessors(&mut cache, &bitstruct, &options)?;
    Ok(quote! {
        #struct_def
        #struct_field_enum_types_def
        #field_accessors
    })
}

/// Check whether the bitstruct has bitstruct_repr attribute, if it has, parse bitstruct_repr attribute
/// 检查是否有bitstruct_repr属性，如果有，解析bitstruct_repr属性
fn check_and_parse_bitstruct_repr_attr(
    options: &mut BistructGenOptions,
    bitstruct: &Bitstruct,
) -> Result<()> {
    // CHECK bitstruct_repr attribute
    let bitstruct_repr_attr = bitstruct
        .attrs
        .iter()
        .find(|attr| attr.path().is_ident(BITSTRUCT_REPR));
    if let Some(bitstruct_repr_attr) = bitstruct_repr_attr {
        let bitstruct_repr: tokens::BitstructRepr =
            bitstruct_repr_attr.parse_args::<tokens::BitstructRepr>()?;
        match options.gen_type {
            BitstructGenType::WithCow | BitstructGenType::WithOwned => {
                if let Some(store_type) = &bitstruct_repr.store_type {
                    return Err(syn::Error::new(
                        store_type.span(),
                        "store_type is only supported in bitstruct_small",
                    ));
                }
            }
            BitstructGenType::ForSmall => {
                if !bitstruct_repr.bit_numbering.is_default() {
                    return Err(syn::Error::new(
                        bitstruct_repr.bit_numbering.span(),
                        "bit_numbering is not supported in bitstruct_small now",
                    ));
                }
            }
        }

        // PARSE bitstruct_repr attribute and set options
        if bitstruct_repr.bit_numbering.is_lsb0() {
            options.bit_numbering = BitNumbering::LSB0;
        } else {
            options.bit_numbering = BitNumbering::MSB0;
        }
        if let Some(store_type) = &bitstruct_repr.store_type {
            options.given_store_type = Some(store_type.clone());
        }
    }
    Ok(())
}

/// Generate a field's enum type definition, if the number of enum type variants is insufficient,
/// it will be automatically padded.
///
/// 生成一个字段的枚举类型定义，如果枚举类型的变体数量不足的情况，会自动补齐。
pub(crate) fn gen_code_of_enum_type(
    field_enum_type: &BitstructFieldEnumType,
    field_bit_width: u32,
    impl_trait: bool,
) -> Result<TokenStream2> {
    let mut dummpy = Cache::default();
    gen_field_enum::check_bitstruct_enum_type_def(field_enum_type, field_bit_width)?;
    unsafe {
        gen_field_enum::unchecked_gen_bitstruct_field_enum_type_def(
            &mut dummpy,
            field_enum_type,
            field_bit_width,
            impl_trait,
        )
    }
}

fn check_fields_of_primitive_type(bitstruct: &Bitstruct) -> Result<()> {
    for field in &bitstruct.fields {
        check_field_of_primitive_type(field)?;
    }
    Ok(())
}

fn check_field_of_primitive_type(field: &BitstructField) -> Result<()> {
    let field_bit_width = field.bit_width()?;
    if field_bit_width == 0 {
        return Err(syn::Error::new(
            field.size.span(),
            "bit width of field must be greater than 0",
        ));
    } else if field_bit_width > SupportedPrimitiveType::SUPPORT_PRIMITIVE_TYPE_MAX_BIT_WIDTH {
        return Err(syn::Error::new(
            field.size.span(),
            format!(
                "bit width of field is too large, the maximum bit width is {}",
                SupportedPrimitiveType::SUPPORT_PRIMITIVE_TYPE_MAX_BIT_WIDTH
            ),
        ));
    }
    if let Some(tokens::BitstructFieldTargetType::PrimitiveType(ty)) = &field.target_type {
        if !ty.is_containable(field_bit_width) {
            return Err(syn::Error::new(
                field.size.span(),
                format!(
                    "bit width of field is too large for {}",
                    ty.to_token_stream()
                ),
            ));
        }
    }
    Ok(())
}