bitstructs_macro 0.1.2

#![allow(dead_code)]
use std::num::NonZeroU32;

use quote::quote;
use syn::{Token, spanned::Spanned};

#[derive(Clone)]
pub struct Bitstruct {
    pub attrs: Vec<syn::Attribute>,
    pub vis: syn::Visibility,
    struct_token: Token![struct],
    pub ident: syn::Ident,
    brace_token: syn::token::Brace,
    pub fields: syn::punctuated::Punctuated<BitstructField, Token![,]>,
}

impl Bitstruct {
    pub fn span(&self) -> proc_macro2::Span {
        self.ident.span()
    }
    pub fn bit_width(&self) -> Result<u32, syn::Error> {
        let mut bit_width = 0;
        for field in &self.fields {
            bit_width += field.bit_width()?;
        }
        Ok(bit_width)
    }
    pub fn store_type(&self) -> Result<syn::Type, syn::Error> {
        let bit_width = self.bit_width()?;
        BitstructField::infer_store_type_by_field_bit_width(bit_width)
    }
}

impl syn::parse::Parse for Bitstruct {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let content;
        Ok(Bitstruct {
            attrs: input.call(syn::Attribute::parse_outer)?,
            vis: input.parse()?,
            struct_token: input.parse()?,
            ident: input.parse()?,
            brace_token: syn::braced!(content in input),
            fields: content.parse_terminated(BitstructField::parse, Token![,])?,
        })
    }
}

#[derive(Clone)]
pub struct BitstructField {
    pub attrs: Vec<syn::Attribute>,
    pub ident: syn::Ident,
    colon_token: Token![:],
    pub size: BitstructFieldSize,
    arrow_token: Option<Token![=>]>,
    pub target_type: Option<BitstructFieldTargetType>,
}
impl BitstructField {
    pub fn span(&self) -> proc_macro2::Span {
        self.ident.span()
    }
    pub fn doc_attrs(&self) -> &[syn::Attribute] {
        &self.attrs
    }
    pub fn infer_store_type_by_field_bit_width(bit_width: u32) -> Result<syn::Type, syn::Error> {
        if bit_width == 0 {
            return Err(syn::Error::new(
                proc_macro2::Span::call_site(),
                "field size must be greater than 0",
            ));
        } else if bit_width <= 8 {
            Ok(syn::parse_quote! {u8})
        } else if bit_width <= 16 {
            Ok(syn::parse_quote! {u16})
        } else if bit_width <= 32 {
            Ok(syn::parse_quote! {u32})
        } else if bit_width <= 64 {
            Ok(syn::parse_quote! {u64})
        } else if bit_width < 128 {
            Ok(syn::parse_quote! {u128})
        } else {
            Err(syn::Error::new(
                proc_macro2::Span::call_site(),
                "field size too large",
            ))
        }
    }
    pub fn store_type(&self) -> Result<syn::Type, syn::Error> {
        let size = self.bit_width()?;
        Self::infer_store_type_by_field_bit_width(size)
            .map_err(|_| syn::Error::new(self.size.span(), "field size too large"))
    }
    pub fn store_type_token_stream(&self) -> Result<proc_macro2::TokenStream, syn::Error> {
        let store_type = self.store_type()?;
        Ok(quote! {#store_type})
    }
    pub fn bit_width(&self) -> Result<u32, syn::Error> {
        Ok(self.size.value()?.get())
    }
    pub fn first_byte_index(&self, bit_offset: u32) -> u32 {
        bit_offset / 8
    }
    pub fn first_byte_bit_offset(&self, bit_offset: u32) -> u32 {
        bit_offset % 8
    }
    pub fn last_byte_index(&self, bit_offset: u32) -> Result<u32, syn::Error> {
        let last_bit = bit_offset + self.bit_width()? - 1;
        Ok(last_bit / 8)
    }
    pub fn bit_mask_litint(&self) -> Result<syn::LitInt, syn::Error> {
        let mask = (1_u128 << self.bit_width()?) - 1;
        Ok(syn::LitInt::new(
            &format!("{}_{}", mask, self.store_type_token_stream()?),
            self.size.span(),
        ))
    }
}
impl syn::parse::Parse for BitstructField {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let attrs = input.call(syn::Attribute::parse_outer)?;
        let ident = input.parse()?;
        let colon_token = input.parse()?;
        let size = input.parse()?;
        let arrow_token: Option<syn::token::FatArrow> = input.parse()?;
        let enum_type = if let Some(arrow_token) = &arrow_token {
            Some(input.parse().map_err(|_| {
                let span = arrow_token.span();
                syn::Error::new(span, "expected enum type after `=>`")
            })?)
        } else {
            None
        };
        Ok(BitstructField {
            attrs,
            ident,
            colon_token,
            size,
            arrow_token,
            target_type: enum_type,
        })
    }
}

#[derive(Clone)]
pub struct BitstructFieldSize {
    val: syn::LitInt,
}
impl BitstructFieldSize {
    pub fn value(&self) -> Result<NonZeroU32, syn::Error> {
        self.val.base10_parse::<NonZeroU32>()
    }

    pub fn span(&self) -> proc_macro2::Span {
        self.val.span()
    }
}
impl syn::parse::Parse for BitstructFieldSize {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        Ok(BitstructFieldSize {
            val: input.parse()?,
        })
    }
}

#[derive(Clone)]
pub(crate) enum BitstructFieldTargetType {
    InlineEnum(BitstructFieldEnumType),
    PrimitiveType(SupportedPrimitiveType),
    NonInlineEnum(syn::Type),
}

impl BitstructFieldTargetType {
    pub fn span(&self) -> proc_macro2::Span {
        match self {
            BitstructFieldTargetType::InlineEnum(enum_type) => enum_type.span(),
            BitstructFieldTargetType::PrimitiveType(ty) => ty.span(),
            BitstructFieldTargetType::NonInlineEnum(ty) => ty.span(),
        }
    }
    pub fn get_type(&self) -> syn::Type {
        match self {
            BitstructFieldTargetType::InlineEnum(enum_type) => syn::Type::Path(syn::TypePath {
                qself: None,
                path: enum_type.ident.clone().into(),
            }),
            BitstructFieldTargetType::PrimitiveType(ty) => ty.ty().clone(),
            BitstructFieldTargetType::NonInlineEnum(ty) => ty.clone(),
        }
    }
    fn is_inline_enum_by_lookahead1(next_token: &syn::parse::Lookahead1) -> bool {
        if next_token.peek(Token![enum]) {
            true
        } else if next_token.peek(Token![#]) {
            true
        } else if next_token.peek(Token![pub]) {
            true
        } else {
            false
        }
    }
}

impl syn::parse::Parse for BitstructFieldTargetType {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let lookahead: syn::parse::Lookahead1<'_> = input.lookahead1();
        if BitstructFieldTargetType::is_inline_enum_by_lookahead1(&lookahead) {
            let enum_type = input.parse()?;
            Ok(BitstructFieldTargetType::InlineEnum(enum_type))
        } else {
            let ty = input.parse()?;
            if let Ok(ty) = SupportedPrimitiveType::new(&ty) {
                Ok(BitstructFieldTargetType::PrimitiveType(ty))
            } else {
                Ok(BitstructFieldTargetType::NonInlineEnum(ty))
            }
        }
    }
}

#[derive(Clone)]
pub struct BitstructFieldEnumType {
    pub attrs: Vec<syn::Attribute>,
    pub vis: syn::Visibility,
    enum_token: Token![enum],
    pub ident: syn::Ident,
    brace_token: syn::token::Brace,
    pub variants: syn::punctuated::Punctuated<BitstructFieldEnumVariant, Token![,]>,
}
impl BitstructFieldEnumType {
    /// The maximum bit size supported by the library.
    /// The bigger bit size will create a lot of enum variants, which is not recommended.
    pub const SUPPORT_MAX_BIT_WIDTH: u32 = 10;
    pub fn span(&self) -> proc_macro2::Span {
        self.ident.span()
    }
}
impl syn::parse::Parse for BitstructFieldEnumType {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let attrs = input.call(syn::Attribute::parse_outer)?;
        let vis = input.parse()?;
        let enum_token: Token![enum] = input.parse()?;
        let ident: syn::Ident = input.parse()?;
        let content;
        let brace_token = syn::braced!(content in input);
        let variants = content.parse_terminated(BitstructFieldEnumVariant::parse, Token![,])?;
        Ok(BitstructFieldEnumType {
            attrs,
            vis,
            enum_token,
            ident,
            brace_token,
            variants,
        })
    }
}

#[derive(Clone)]
pub struct BitstructFieldEnumVariant {
    pub attrs: Vec<syn::Attribute>,
    pub ident: syn::Ident,
    eq_token: Token![=],
    pub value: syn::LitInt,
}

impl syn::parse::Parse for BitstructFieldEnumVariant {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        Ok(BitstructFieldEnumVariant {
            attrs: input.call(syn::Attribute::parse_outer)?,
            ident: input.parse()?,
            eq_token: input.parse()?,
            value: input.parse()?,
        })
    }
}

pub struct BitstructRepr {
    // store_type can be `u8`, `u16`, `u32`, `u64`, `u128`
    pub bit_numbering: BitNumbering,
    pub store_type: Option<syn::Ident>,
}

impl BitstructRepr {
    pub fn sotre_type_bit_width(store_type: &syn::Ident) -> Result<u32, syn::Error> {
        let bit_width = match store_type.to_string().as_str() {
            "u8" => 8,
            "u16" => 16,
            "u32" => 32,
            "u64" => 64,
            "u128" => 128,
            _ => return Err(syn::Error::new(store_type.span(), "invalid store type")),
        };
        Ok(bit_width)
    }
    pub fn span(&self) -> proc_macro2::Span {
        self.store_type
            .as_ref()
            .map(|store_type| store_type.span())
            .unwrap_or_else(|| self.bit_numbering.span())
    }
}

impl syn::parse::Parse for BitstructRepr {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let ident = input.parse::<syn::Ident>()?;
        match ident.to_string().as_str() {
            "MSB0" => Ok(BitstructRepr {
                store_type: None,
                bit_numbering: BitNumbering::MSB0(ident),
            }),
            "LSB0" => Ok(BitstructRepr {
                store_type: None,
                bit_numbering: BitNumbering::LSB0(ident),
            }),
            "u8" | "u16" | "u32" | "u64" | "u128" => Ok(BitstructRepr {
                store_type: Some(ident),
                bit_numbering: BitNumbering::DEFAULT,
            }),
            _ => Err(syn::Error::new(
                ident.span(),
                "expected `MSB0`, `LSB0`, `u8`, `u16`, `u32`, `u64`, `u128`",
            )),
        }
    }
}

syn::custom_keyword!(MSB0);
syn::custom_keyword!(LSB0);

pub enum BitNumbering {
    MSB0(syn::Ident),
    LSB0(syn::Ident),
    DEFAULT,
}
impl BitNumbering {
    pub fn is_msb0(&self) -> bool {
        matches!(self, BitNumbering::MSB0(_))
    }
    pub fn is_lsb0(&self) -> bool {
        matches!(self, BitNumbering::LSB0(_)) || matches!(self, BitNumbering::DEFAULT)
    }
    pub fn is_default(&self) -> bool {
        matches!(self, BitNumbering::DEFAULT)
    }
    pub fn span(&self) -> proc_macro2::Span {
        match self {
            BitNumbering::MSB0(msb0) => msb0.span(),
            BitNumbering::LSB0(lsb0) => lsb0.span(),
            BitNumbering::DEFAULT => proc_macro2::Span::call_site(),
        }
    }
}

// impl syn:: for BitNumbering {
//     fn span(&self) -> proc_macro2::Span {
//         self.span()
//     }
//     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
//         match self {
//             BitNumbering::MSB0(msb0) => msb0.to_tokens(tokens),
//             BitNumbering::LSB0(lsb0) => lsb0.to_tokens(tokens),
//         }
//     }
// }

#[derive(Clone)]
pub(crate) enum SupportedPrimitiveType {
    Bool(syn::Type),
    U8(syn::Type),
    U16(syn::Type),
    U32(syn::Type),
    U64(syn::Type),
    U128(syn::Type),
}

impl SupportedPrimitiveType {
    pub const SUPPORT_PRIMITIVE_TYPE_MAX_BIT_WIDTH: u32 = 128;
    fn new(ty: &syn::Type) -> Result<Self, syn::Error> {
        let err = Err(syn::Error::new(
            ty.span(),
            "expected `bool`, `u8`, `u16`, `u32`, `u64`, `u128`",
        ));
        if let syn::Type::Path(syn::TypePath { path, .. }) = ty {
            if let Some(ident) = path.get_ident() {
                let ret = match ident.to_string().as_str() {
                    "bool" => Ok(SupportedPrimitiveType::Bool(ty.clone())),
                    "u8" => Ok(SupportedPrimitiveType::U8(ty.clone())),
                    "u16" => Ok(SupportedPrimitiveType::U16(ty.clone())),
                    "u32" => Ok(SupportedPrimitiveType::U32(ty.clone())),
                    "u64" => Ok(SupportedPrimitiveType::U64(ty.clone())),
                    "u128" => Ok(SupportedPrimitiveType::U128(ty.clone())),
                    _ => err,
                };
                return ret;
            }
        }
        err
    }

    pub fn span(&self) -> proc_macro2::Span {
        match self {
            SupportedPrimitiveType::Bool(ident) => ident.span(),
            SupportedPrimitiveType::U8(ident) => ident.span(),
            SupportedPrimitiveType::U16(ident) => ident.span(),
            SupportedPrimitiveType::U32(ident) => ident.span(),
            SupportedPrimitiveType::U64(ident) => ident.span(),
            SupportedPrimitiveType::U128(ident) => ident.span(),
        }
    }
    pub fn infer_primitive_type_by_bit_width(bit_width: u32) -> Result<syn::Type, syn::Error> {
        if bit_width == 0 {
            return Err(syn::Error::new(
                proc_macro2::Span::call_site(),
                "field size must be greater than 0",
            ));
        } else if bit_width <= 8 {
            Ok(syn::parse_quote! {u8})
        } else if bit_width <= 16 {
            Ok(syn::parse_quote! {u16})
        } else if bit_width <= 32 {
            Ok(syn::parse_quote! {u32})
        } else if bit_width <= 64 {
            Ok(syn::parse_quote! {u64})
        } else if bit_width < 128 {
            Ok(syn::parse_quote! {u128})
        } else {
            Err(syn::Error::new(
                proc_macro2::Span::call_site(),
                "field size too large",
            ))
        }
    }
    pub fn is_containable(&self, bit_width: u32) -> bool {
        match self {
            SupportedPrimitiveType::Bool(_) => bit_width == 1,
            SupportedPrimitiveType::U8(_) => bit_width <= 8,
            SupportedPrimitiveType::U16(_) => bit_width <= 16,
            SupportedPrimitiveType::U32(_) => bit_width <= 32,
            SupportedPrimitiveType::U64(_) => bit_width <= 64,
            SupportedPrimitiveType::U128(_) => bit_width <= 128,
        }
    }
    pub fn is_fullfilled_the_ty(&self, bit_width: u32) -> bool {
        match self {
            SupportedPrimitiveType::Bool(_) => bit_width == 1,
            SupportedPrimitiveType::U8(_) => bit_width == 8,
            SupportedPrimitiveType::U16(_) => bit_width == 16,
            SupportedPrimitiveType::U32(_) => bit_width == 32,
            SupportedPrimitiveType::U64(_) => bit_width == 64,
            SupportedPrimitiveType::U128(_) => bit_width == 128,
        }
    }
    pub fn ty(&self) -> &syn::Type {
        match self {
            SupportedPrimitiveType::Bool(ty) => ty,
            SupportedPrimitiveType::U8(ty) => ty,
            SupportedPrimitiveType::U16(ty) => ty,
            SupportedPrimitiveType::U32(ty) => ty,
            SupportedPrimitiveType::U64(ty) => ty,
            SupportedPrimitiveType::U128(ty) => ty,
        }
    }
}

impl quote::ToTokens for SupportedPrimitiveType {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        self.ty().to_tokens(tokens)
    }
}