kago-macros 0.1.0

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Ident, LitInt, Type};

pub fn impl_add(struct_name: &Ident) -> TokenStream {
    quote! {
        impl std::ops::Add for #struct_name {
            type Output = Self;

            fn add(self, rhs: Self) -> Self {
                Self::from_bits(self.0 + rhs.0)
            }
        }
    }
}

pub fn impl_checked_add(struct_name: &Ident) -> TokenStream {
    quote! {
        impl num_traits::ops::checked::CheckedAdd for #struct_name {
            fn checked_add(&self, rhs: &Self) -> Option<Self> {
                Some(Self::from_bits(self.0.checked_add(rhs.0)?))
            }
        }
    }
}

pub fn impl_sub(struct_name: &Ident) -> TokenStream {
    quote! {
        impl std::ops::Sub for #struct_name {
            type Output = Self;

            fn sub(self, rhs: Self) -> Self {
                Self::from_bits(self.0 - rhs.0)
            }
        }
    }
}

pub fn impl_checked_sub(struct_name: &Ident) -> TokenStream {
    quote! {
        impl num_traits::ops::checked::CheckedSub for #struct_name {
            fn checked_sub(&self, rhs: &Self) -> Option<Self> {
                Some(Self::from_bits(self.0.checked_sub(rhs.0)?))
            }
        }
    }
}

pub fn impl_mul(struct_name: &Ident, shift: &Option<LitInt>) -> TokenStream {
    if let Some(shift) = shift {
        quote! {
            impl std::ops::Mul for #struct_name {
                type Output = Self;

                fn mul(self, rhs: Self) -> Self {
                    Self::from_bits(self.0 * (rhs.0 >> #shift))
                }
            }
        }
    } else {
        quote! {
            impl std::ops::Mul for #struct_name {
                type Output = Self;

                fn mul(self, rhs: Self) -> Self {
                    Self::from_bits(self.0 * rhs.0)
                }
            }
        }
    }
}

pub fn impl_checked_mul(struct_name: &Ident, ty: &Type, shift: &Option<LitInt>) -> TokenStream {
    if let Some(shift) = shift {
        quote! {
            impl num_traits::ops::checked::CheckedMul for #struct_name {
                fn checked_mul(&self, rhs: &Self) -> Option<Self> {
                    Some(Self::from_bits(
                            <#ty as num_traits::ops::checked::CheckedMul>::checked_mul(&self.0, &(rhs.0 >> #shift))?
                    ))
                }
            }
        }
    } else {
        quote! {
            impl num_traits::ops::checked::CheckedMul for #struct_name {
                fn checked_mul(&self, rhs: &Self) -> Option<Self> {
                    Some(Self::from_bits(
                        <#ty as num_traits::ops::checked::CheckedMul>::checked_mul(&self.0, &rhs.0)?
                    ))
                }
            }
        }
    }
}

pub fn impl_div(struct_name: &Ident, shift: &Option<LitInt>) -> TokenStream {
    if let Some(shift) = shift {
        quote! {
            impl std::ops::Div for #struct_name {
                type Output = Self;

                fn div(self, rhs: Self) -> Self {
                    Self::from_bits((self.0 / rhs.0) << #shift)
                }
            }
        }
    } else {
        quote! {
            impl std::ops::Div for #struct_name {
                type Output = Self;

                fn div(self, rhs: Self) -> Self {
                    Self::from_bits(self.0 / rhs.0)
                }
            }
        }
    }
}

pub fn impl_checked_div(struct_name: &Ident, ty: &Type, shift: &Option<LitInt>) -> TokenStream {
    if let Some(shift) = shift {
        quote! {
            impl num_traits::ops::checked::CheckedDiv for #struct_name {
                fn checked_div(&self, rhs: &Self) -> Option<Self> {
                    let value = <#ty as num_traits::ops::checked::CheckedDiv>::checked_div(&self.0, &rhs.0)?;
                    if (<Self as num_traits::Bounded>::min_value().0 >> #shift) <= value && value <= (<Self as num_traits::Bounded>::max_value().0 >> #shift) {
                        Some(Self::new(value))
                    } else {
                        None
                    }
                }
            }
        }
    } else {
        quote! {
            impl num_traits::ops::checked::CheckedDiv for #struct_name {
                fn checked_div(&self, rhs: &Self) -> Option<Self> {
                    Some(Self::from_bits(
                        <#ty as num_traits::ops::checked::CheckedDiv>::checked_div(&self.0, &rhs.0)?
                    ))
                }
            }
        }
    }
}

pub fn impl_rem(struct_name: &Ident) -> TokenStream {
    quote! {
        impl std::ops::Rem for #struct_name {
            type Output = Self;

            fn rem(self, rhs: Self) -> Self {
                Self::from_bits(self.0 % rhs.0)
            }
        }
    }
}

pub fn impl_checked_rem(struct_name: &Ident, ty: &Type) -> TokenStream {
    quote! {
        impl num_traits::ops::checked::CheckedRem for #struct_name {
            fn checked_rem(&self, rhs: &Self) -> Option<Self> {
                Some(Self::from_bits(
                    <#ty as num_traits::ops::checked::CheckedRem>::checked_rem(&self.0, &rhs.0)?
                ))
            }
        }
    }
}

pub fn impl_saturating(struct_name: &Ident, ty: &Type, masking: &TokenStream) -> TokenStream {
    quote! {
        impl num_traits::ops::saturating::Saturating for #struct_name {
            fn saturating_add(self, rhs: Self) -> Self {
                Self::from_bits(
                    <#ty as num_traits::ops::saturating::Saturating>::saturating_add(self.0, rhs.0) #masking
                )
            }

            fn saturating_sub(self, rhs: Self) -> Self {
                Self::from_bits(
                    <#ty as num_traits::ops::saturating::Saturating>::saturating_sub(self.0, rhs.0) #masking
                )
            }
        }
    }
}