ecore 0.10.0

use core::{
    marker::PhantomData,
    mem::{MaybeUninit, transmute},
    num::NonZero,
};

use self::sealed::IntoBitField;
use crate::{
    Array, EnumMap, MapEnum,
    bitint::{BitInt, u1},
    int::{BasicInt, BasicUInt, BitsOp, PrimaryInt},
    nbool::NBool,
    repr::{Uncheckable, Unchecked},
};
use bytemuck::{Pod, Zeroable};
pub use ecore_macro::{BitsCast, bitfld};

pub use self::{
    dynref::DynBitField,
    seg2::BitField2,
    tag_union::{BitsEnumReLayout, ReLayoutBitsEnum},
};

pub mod prelude {
    #[doc(hidden)]
    pub use super::{BitField, BitField2, BitsCast, BitsEnumReLayout, PlainBitsCast, ReLayoutBitsEnum, bitfld};
    #[doc(hidden)]
    pub use crate::{
        int::{BasicInt, BasicUInt, PrimaryInt},
        repr::{ConvertEndian, PlainUncheckable, Uncheckable, Unchecked},
    };
    #[doc(hidden)]
    pub use bytemuck::{Pod, Zeroable};
}

/// type which can safe cast inter bits which maybe not align to bytes, if `T` only occupy bits but not exhaustive then can impl [Uncheckable] instead,
/// so [`Unchecked<T: Uncheckable<UncheckedRaw: BasicInt>>`] will auto impl [BitsCast], usually should use `#[bitfld(..)]` to gen impl
pub trait BitsCast: Copy {
    const BITS: u32;
    /// Convert from underlying int, data is store at lowest `Self::BITS`, higher bits can be any value
    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self;

    /// Convert into underlying int, data must store at lowest `Self::BITS`, higher bits can be any value
    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits;
}

/// [BitsCast] type with can safe transmute between [`Self::Bits`] and [`Self`], **will deprecated if `const_trait_impl` is stable**,
/// usually should use `#[bitfld(..)]` to gen impl
///
/// # Safety
///
/// Implementors must ensure that `Self` and `Self::Bits` have the same size and bit representation.
pub unsafe trait PlainBitsCast: BitsCast {
    /// stored raw value type, can be u1, u2, ... uN and i1, i2, ... iN
    type Bits: BasicInt;
    const ASSERT: () = assert!(size_of::<Self>() == size_of::<Self::Bits>() && Self::BITS == Self::Bits::BITS);
}

/// Bit-field type for bit range `START..=LAST` of bit-struct `STRUCT` with field type `FLD`.
/// This object is typically not constructed directly, but rather bound as a reference from the bit-struct.
/// It is primarily auto-generated by `#[bitfld(..)]` and generally not defined manually.
#[repr(transparent)]
pub struct BitField<STRUCT, FLD, const START: u32, const LAST: u32 = START> {
    raw: STRUCT,
    mark: PhantomData<FLD>,
}

impl<STRUCT: BitsCast, FLD: BitsCast, const START: u32, const LAST: u32> BitField<STRUCT, FLD, START, LAST> {
    const ASSERT: () = assert!(START <= LAST && LAST < STRUCT::BITS && Self::BITS == FLD::BITS);

    /// Reinterprets a shared reference to the bit-struct as a shared reference to this bit-field.
    pub const fn bind_ref(raw: &STRUCT) -> &Self {
        let () = Self::ASSERT;
        unsafe { transmute(raw) }
    }

    /// Reinterprets a mutable reference to the bit-struct as a mutable reference to this bit-field.
    pub const fn bind_mut(raw: &mut STRUCT) -> &mut Self {
        let () = Self::ASSERT;
        unsafe { transmute(raw) }
    }
}

impl<STRUCT: BitsCast, FLD: BitsCast, const START: u32, const LAST: u32> BitField<STRUCT, FLD, START, LAST> {
    pub const START: u32 = START;
    pub const LAST: u32 = LAST;
    pub const BITS: u32 = LAST - START + 1;

    /// Returns the bit width of this field, using a dummy function argument for type-guided access.
    pub const fn bits(_: fn(&STRUCT) -> &Self) -> u32 {
        Self::BITS
    }

    /// Returns the starting bit position of this field, using a dummy function argument for type-guided access.
    pub const fn bits_start(_: fn(&STRUCT) -> &Self) -> u32 {
        START
    }

    /// Returns the ending bit position of this field, using a dummy function argument for type-guided access.
    pub const fn bits_end(_: fn(&STRUCT) -> &Self) -> u32 {
        LAST
    }
}

impl<const START: u32, const LAST: u32, STRUCT: PlainBitsCast<Bits: BasicUInt>, FLD> BitField<STRUCT, FLD, START, LAST> {
    fn raw(&self) -> <STRUCT::Bits as BasicInt>::Primary {
        STRUCT::into_underlying(self.raw)
    }
}

impl<const START: u32, const LAST: u32, STRUCT: PlainBitsCast<Bits: BasicUInt>, FLD: BitsCast> BitField<STRUCT, FLD, START, LAST> {
    /// Reads the bit-field's value from the struct.
    pub fn read(&self) -> FLD {
        FLD::from_underlying(self.raw() >> START)
    }

    /// Returns a new bit-struct with this field set to the given value, leaving the original unchanged.
    #[must_use]
    pub fn with(&self, fld: impl IntoBitField<FLD>) -> STRUCT {
        let fld: <STRUCT::Bits as BasicInt>::Primary = IntoBitField::into_bit_field(fld);
        STRUCT::from_underlying(self.raw().with_bits::<START, LAST>(fld.cast_as()))
    }

    /// Writes the given value into this bit-field.
    pub fn write(&mut self, fld: impl IntoBitField<FLD>) {
        self.raw = self.with(fld);
    }
}

impl<const START: u32, const LAST: u32, STRUCT: PlainBitsCast<Bits: BasicUInt>, FLD: Uncheckable<UncheckedRaw: BasicInt>>
    BitField<STRUCT, Unchecked<FLD>, START, LAST>
{
    /// Validates and returns the checked value from the bit-field.
    pub fn get(&self) -> Result<FLD, FLD::UncheckedRaw> {
        self.read().get()
    }

    /// Returns the value without validation.
    ///
    /// # Safety
    ///
    /// The caller must ensure the contained raw value represents a valid `FLD`.
    pub unsafe fn get_unchecked(&self) -> FLD {
        unsafe { self.read().get_unchecked() }
    }

    /// Returns the value, falling back to the default if validation fails.
    pub fn get_or_default(&self) -> FLD
    where
        FLD: Default,
    {
        self.read().get_or_default()
    }
}

#[cfg(feature = "ranged-int")]
impl<const START: u32, const LAST: u32, STRUCT: PlainBitsCast<Bits: BasicUInt>, STORE: BasicUInt, RANGE: crate::int::ranged::RIntRange>
    BitField<STRUCT, Unchecked<crate::int::ranged::RInt<STORE, RANGE>>, START, LAST>
{
    /// Returns the ranged integer value if valid, or `None`.
    pub fn value(&self) -> Option<RANGE::Value> {
        self.read().value()
    }

    /// Returns the ranged integer value without validation.
    ///
    /// # Safety
    ///
    /// The caller must ensure the contained raw value represents a valid value within the range.
    pub unsafe fn value_unchecked(&self) -> RANGE::Value {
        unsafe { self.read().value_unchecked() }
    }

    /// Returns the ranged integer value, falling back to the default if invalid.
    pub fn value_or_default(&self) -> RANGE::Value {
        self.read().value_or_default()
    }
}

impl<STRUCT: Copy, FLD, const START: u32, const LAST: u32> Copy for BitField<STRUCT, FLD, START, LAST> {}
impl<STRUCT: Clone, FLD, const START: u32, const LAST: u32> Clone for BitField<STRUCT, FLD, START, LAST> {
    fn clone(&self) -> Self {
        Self { raw: self.raw.clone(), mark: PhantomData }
    }
}

unsafe impl<STRUCT: Zeroable, FLD, const START: u32, const LAST: u32> Zeroable for BitField<STRUCT, FLD, START, LAST> {}

unsafe impl<STRUCT: Pod, FLD: 'static, const START: u32, const LAST: u32> Pod for BitField<STRUCT, FLD, START, LAST> {}

impl<FLD: Uncheckable<UncheckedRaw: BasicInt>> BitsCast for Unchecked<FLD> {
    const BITS: u32 = FLD::UncheckedRaw::BITS;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        Unchecked::new_with_raw_value(v.cast_as())
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        sf.raw_value().cast_as()
    }
}
unsafe impl<FLD: Uncheckable<UncheckedRaw: BasicInt>> PlainBitsCast for Unchecked<FLD> {
    type Bits = FLD::UncheckedRaw;
}

impl BitsCast for bool {
    const BITS: u32 = 1;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        v.read_bit::<0>()
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        (sf as u8).cast_as()
    }
}
unsafe impl PlainBitsCast for bool {
    type Bits = u1;
}

impl BitsCast for NBool {
    const BITS: u32 = 1;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        Self::new_with_raw_value(v.cast_as())
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        sf.raw_value().cast_as()
    }
}
unsafe impl PlainBitsCast for NBool {
    type Bits = u1;
}

macro_rules! impl_primary {
    ($type:ty) => {
        impl BitsCast for $type {
            const BITS: u32 = Self::BITS;

            fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
                v.cast_as()
            }

            fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
                sf.cast_as()
            }
        }
        unsafe impl PlainBitsCast for $type {
            type Bits = Self;
        }

        impl BitsCast for Option<NonZero<$type>> {
            const BITS: u32 = <$type>::BITS;

            fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
                NonZero::new(v.cast_as())
            }

            fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
                let Some(v) = sf else { return Bits::ZERO };
                v.get().cast_as()
            }
        }
        unsafe impl PlainBitsCast for Option<NonZero<$type>> {
            type Bits = $type;
        }

        impl<const BITS: u32> BitsCast for Option<BitInt<NonZero<$type>, BITS>> {
            const BITS: u32 = BITS;

            fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
                BitInt::<NonZero<$type>, BITS>::new_non_zero(v.cast_as())
            }

            fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
                let Some(v) = sf else { return Bits::ZERO };
                v.get().cast_as()
            }
        }
        unsafe impl<const BITS: u32> PlainBitsCast for Option<BitInt<NonZero<$type>, BITS>> {
            type Bits = BitInt<$type, BITS>;
        }
    };
}
impl_primary!(u8);
impl_primary!(u16);
impl_primary!(u32);
impl_primary!(u64);
impl_primary!(i8);
impl_primary!(i16);
impl_primary!(i32);
impl_primary!(i64);

impl<T: PrimaryInt, const BITS: u32> BitsCast for BitInt<T, BITS> {
    const BITS: u32 = BITS;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        v.cast_as()
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        sf.cast_as()
    }
}
unsafe impl<T: PrimaryInt, const BITS: u32> PlainBitsCast for BitInt<T, BITS> {
    type Bits = Self;
}

#[cfg(feature = "ranged-int")]
impl<STORE: BasicUInt, RANGE: crate::int::ranged::RIntRange> BitsCast for crate::int::ranged::RInt<STORE, RANGE> {
    const BITS: u32 = STORE::BITS;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        const { assert!(Self::EXHAUSTIVE, "RInt as bitfield must exhaustive") }
        unsafe { Self::new_with_raw_value(v.cast_as()) }
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        sf.raw_value().cast_as()
    }
}

impl<Fld: BitsCast> BitsCast for Option<Fld> {
    const BITS: u32 = Fld::BITS + 1;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        if v.read_bit::<0>() { Some(Fld::from_underlying(v >> 1u32)) } else { None }
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        if let Some(v) = sf { Fld::into_underlying::<Bits>(v) << 1u32 | Bits::ONE } else { Bits::ZERO }
    }
}

impl<OkT: BitsCast, ErrT: BitsCast> BitsCast for Result<OkT, ErrT> {
    const BITS: u32 = if OkT::BITS > ErrT::BITS { OkT::BITS } else { ErrT::BITS } + 1;

    fn from_underlying<Bits: PrimaryInt>(v: Bits) -> Self {
        if v.read_bit::<0>() { Ok(OkT::from_underlying(v >> 1u32)) } else { Err(ErrT::from_underlying(v >> 1u32)) }
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        match sf {
            Ok(v) => OkT::into_underlying::<Bits>(v) << 1u32,
            Err(v) => ErrT::into_underlying::<Bits>(v) << 1u32 | Bits::ONE,
        }
    }
}

impl<Fld: BitsCast, const N: usize> BitsCast for [Fld; N] {
    const BITS: u32 = {
        assert!(N <= 128);
        Fld::BITS * N as u32
    };

    fn from_underlying<Bits: PrimaryInt>(mut v: Bits) -> Self {
        let mut out = MaybeUninit::<Self>::uninit();
        let buf = unsafe { out.assume_init_mut() };
        for fld in buf.iter_mut() {
            *fld = Fld::from_underlying(v);
            v >>= Fld::BITS;
        }
        unsafe { out.assume_init() }
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        let mut v = Bits::ZERO;
        let mut shift = 0;
        for fld in sf.into_iter() {
            v |= Fld::into_underlying::<Bits>(fld) << shift;
            shift *= Fld::BITS;
        }
        v
    }
}

impl<Ids: MapEnum<Array: Array<Map<Fld>: Copy>>, Fld: BitsCast> BitsCast for EnumMap<Ids, Fld> {
    const BITS: u32 = {
        assert!(Ids::Array::LEN <= 128);
        Ids::Array::LEN as u32 * Fld::BITS
    };

    fn from_underlying<Bits: PrimaryInt>(mut v: Bits) -> Self {
        let mut out = MaybeUninit::<Self>::uninit();
        let buf = unsafe { out.assume_init_mut() }.as_array_mut().as_mut();
        for fld in buf.iter_mut() {
            *fld = Fld::from_underlying(v);
            v >>= Fld::BITS;
        }
        unsafe { out.assume_init() }
    }

    fn into_underlying<Bits: PrimaryInt>(sf: Self) -> Bits {
        let mut v = Bits::ZERO;
        let mut shift = 0;
        for fld in sf.as_array().as_ref() {
            v |= Fld::into_underlying::<Bits>(*fld) << shift;
            shift *= Fld::BITS;
        }
        v
    }
}

mod const_op;
mod dynref;
mod seg2;
mod tag_union;

mod sealed {
    use crate::{
        bitfld::{BitsCast, BitsEnumReLayout, ReLayoutBitsEnum},
        int::{BasicInt, BasicUInt, PrimaryInt},
        nbool::NBool,
        repr::{Uncheckable, Unchecked},
    };

    /// Value can be convert bit field value `Fld`, only has two implement:
    /// * `impl<Fld: BitsCast> IntoBitField<Fld> for Fld`
    /// * `impl<Fld: Uncheckable> IntoBitField<Unchecked<Fld>> for Fld where Unchecked<Fld>: BitsCast`
    /// * `impl<Enum: ReLayoutBitsEnum, ...> IntoBitField<BitsEnumReLayout<Enum, ...>> for Enum`
    ///
    /// For examble:
    /// ```ignore
    /// #[bitfld(u2)]
    /// enum Enum1 {
    ///     A = 0,
    ///     B = 1,
    /// }
    /// #[bitfld(u4, tag(u1), payload(u3))]
    /// enum Enum2 {
    ///     A(u2),
    ///     B(u3),
    /// }
    /// #[bitfld(u8)]
    /// struct Struct1 {
    ///     a: bitfld!(bool, 0),
    ///     b: bitfld!(Unchecked<Enum1>, 1),
    ///     c: bitfld!(NBool, 2),
    ///     d: bitfld!(BitsEnumReLayout<Enum2, u4, 3, 0>, 3..=7)
    /// }
    /// Struct1(0).a().with(true); // field `a` must with a `bool` value
    /// Struct1(0).b().with(Unchecked::from(Enum1::A)).b().with(Enum1::B); // field `b` can with a `Unchecked<Enum1>` value, or with a `Enum1` value
    /// Struct1(0).c().with(NBool::new(true)).c().with(true); // field `c` can with a `NBool` value, or with a `bool` value
    /// Struct1(0).d().with(BitsEnumReLayout::new(Enum2::A(...))).d().with(Enum2::A(...)); // field `d` can with a `Enum2` value, or with a `BitsEnumReLayout<Enum2, ...>` value
    /// ```
    ///
    pub trait IntoBitField<Fld: BitsCast>: Copy {
        fn into_bit_field<Bits: PrimaryInt>(src: Self) -> Bits;
    }

    impl<Fld: BitsCast> IntoBitField<Fld> for Fld {
        #[inline(always)]
        fn into_bit_field<Bits: PrimaryInt>(src: Self) -> Bits {
            BitsCast::into_underlying(src)
        }
    }

    impl IntoBitField<NBool> for bool {
        #[inline(always)]
        fn into_bit_field<Bits: PrimaryInt>(src: Self) -> Bits {
            NBool::into_underlying(NBool::new(src))
        }
    }

    impl<Fld: Uncheckable<UncheckedRaw: BasicInt>> IntoBitField<Unchecked<Fld>> for Fld {
        #[inline(always)]
        fn into_bit_field<Bits: PrimaryInt>(src: Self) -> Bits {
            Uncheckable::raw_value(src).cast_as()
        }
    }

    impl<Enum: ReLayoutBitsEnum, Raw: BasicUInt, const TAG_START: u32, const PAYLOAD_START: u32>
        IntoBitField<BitsEnumReLayout<Enum, Raw, TAG_START, PAYLOAD_START>> for Enum
    {
        fn into_bit_field<Bits: PrimaryInt>(src: Self) -> Bits {
            BitsEnumReLayout::<Enum, Raw, TAG_START, PAYLOAD_START>::new(src).into_bits().cast_as()
        }
    }
}