swamp-abi-memtypes 0.0.1

use std::fmt::{Display, Formatter};
use std::ops::{Add, Div, Sub};

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum MemoryAlignment {
    // Do not change the order.
    U8,
    U16,
    U32,
    U64,
}

impl MemoryAlignment {
    #[must_use]
    const fn rank(&self) -> usize {
        match self {
            Self::U8 => 1,
            Self::U16 => 2,
            Self::U32 => 3,
            Self::U64 => 4,
        }
    }
    #[must_use]
    pub const fn greater_than(&self, other: Self) -> bool {
        self.rank() > other.rank()
    }
}

impl From<MemoryAlignment> for usize {
    fn from(val: MemoryAlignment) -> Self {
        match val {
            MemoryAlignment::U8 => 1,
            MemoryAlignment::U16 => 2,
            MemoryAlignment::U32 => 4,
            MemoryAlignment::U64 => 8,
        }
    }
}

impl From<MemoryAlignment> for u8 {
    fn from(val: MemoryAlignment) -> Self {
        match val {
            MemoryAlignment::U8 => 1,
            MemoryAlignment::U16 => 2,
            MemoryAlignment::U32 => 4,
            MemoryAlignment::U64 => 8,
        }
    }
}

impl TryInto<MemoryAlignment> for usize {
    type Error = ();

    fn try_into(self) -> Result<MemoryAlignment, Self::Error> {
        let converted = match self {
            1 => MemoryAlignment::U8,
            2 => MemoryAlignment::U16,
            4 => MemoryAlignment::U32,
            8 => MemoryAlignment::U64,

            _ => return Err(()),
        };
        Ok(converted)
    }
}


#[derive(Debug, Copy, Clone, PartialOrd, Ord, Eq, PartialEq)]
pub struct MemorySize(pub u32);

impl Display for MemorySize {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let description = human_memsize::human_size(self.0.into());
        write!(f, "{description}")
    }
}



#[derive(Debug, Copy, Eq, PartialEq, Hash, Clone, Ord, PartialOrd)]
pub struct MemoryOffset(pub u32);

impl MemoryOffset {
    #[must_use]
    pub const fn to_size(&self) -> MemorySize {
        MemorySize(self.0)
    }
}

impl Display for MemoryOffset {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "+{:X}", self.0)
    }
}

impl MemoryOffset {
    pub fn space(&mut self, memory_size: MemorySize, alignment: MemoryAlignment) -> Self {
        let start = align(self.0 as usize, alignment.into()) as u32;
        self.0 = start + memory_size.0;
        Self(start)
    }
}

impl Add<MemorySize> for MemoryOffset {
    type Output = Self;

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

impl From<MemoryAlignment> for MemoryOffset {
    fn from(val: MemoryAlignment) -> Self {
        let octets: usize = val.into();
        Self(octets as u32)
    }
}

impl From<MemorySize> for usize {
    fn from(val: MemorySize) -> Self {
        val.0 as Self
    }
}


#[must_use]
pub fn align_to(addr: MemoryOffset, alignment: MemoryAlignment) -> MemoryOffset {
    MemoryOffset(align(addr.0 as usize, alignment.into()) as u32)
}

/// # Arguments
/// * `offset` - The offset after the last field (end of layout).
/// * `base_offset` - The starting offset of the struct/tuple/union.
/// * `max_alignment` - The maximum alignment required by any field.
///
/// # Returns
/// The total size, rounded up to `max_alignment`.
/// # Notes
/// The total size of a struct is always rounded up to a multiple of its alignment.
/// It might be strange in that it "wastes" memory for the potential parent struct
/// to place items of lower memory alignment. (reuse tail padding).
/// It simplifies things as well with code generation and similar, that a struct
/// is always the same size and doesn't have to rely on where the struct is contained.
/// It also ensures that arrays of the struct are correctly aligned according to the ABI,
/// and matches the behavior of C, C++, and Rust.
/// Note: The tail padding at the end of a struct is not reused for subsequent fields
/// in a parent struct-this is required for safe and predictable layout
#[must_use]
pub fn adjust_size_to_alignment(
    unaligned_size: MemorySize,
    max_alignment: MemoryAlignment,
) -> MemorySize {
    align_to(MemoryOffset(unaligned_size.0), max_alignment).to_size()
}


#[derive(Copy, Clone)]
pub struct CountU32(pub u32);



impl Div<Self> for MemorySize {
    type Output = CountU32;

    fn div(self, rhs: Self) -> Self::Output {
        assert!(rhs.0 > 0, "Division by zero in MemorySize");
        assert!(
            self.0 > 0,
            "Numerator must be positive in MemorySize division"
        );
        assert_eq!(
            self.0 % rhs.0,
            0,
            "MemorySize division must be exact and positive"
        );

        CountU32(self.0 / rhs.0)
    }
}

impl Add<Self> for MemoryOffset {
    type Output = Self;

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

impl Sub<Self> for MemoryOffset {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self {
        assert!(rhs.0 <= self.0);
        Self(self.0 - rhs.0)
    }
}

impl MemoryOffset {
    #[must_use]
    pub const fn as_size(&self) -> MemorySize {
        MemorySize(self.0)
    }
}

pub const SAFE_ALIGNMENT: usize = 8;

#[must_use]
pub fn align(addr: usize, alignment: usize) -> usize {
    debug_assert!(
        alignment.is_power_of_two(),
        "alignment must be a power of two"
    );
    (addr + alignment - 1) & !(alignment - 1)
}


impl MemoryOffset {
    #[must_use]
    pub fn add(&self, size: MemorySize, alignment: MemoryAlignment) -> Self {
        let new_start = align(self.0 as usize, alignment.into()) as u32;
        Self(new_start + size.0)
    }
}