// SPDX-FileCopyrightText: Copyright 2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Memory descriptor

use packed_struct::prelude::*;

use core::cell::UnsafeCell;
use core::ptr;

use crate::MemoryAttributesIndex;

use super::address::PhysicalAddress;
use super::TranslationGranule;

/// Memory shareability
#[derive(PrimitiveEnum_u8, Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum Shareability {
    #[default]
    NonShareable = 0b00,
    Outer = 0b10,
    Inner = 0b11,
}

/// Data access permission
#[allow(non_camel_case_types)]
#[derive(PrimitiveEnum_u8, Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum DataAccessPermissions {
    #[default]
    ReadWrite_None = 0b00,
    ReadWrite_ReadWrite = 0b01,
    ReadOnly_None = 0b10,
    ReadOnly_ReadOnly = 0b11,
}

/// Memory attributes
#[derive(PackedStruct, Clone, Debug, PartialEq, Eq, Default)]
#[packed_struct(size_bytes = "8", bit_numbering = "lsb0")]
pub struct Attributes {
    #[packed_field(bits = "54")]
    pub uxn: bool,
    #[packed_field(bits = "53")]
    pub pxn: bool,
    #[packed_field(bits = "52")]
    pub contiguous: bool,
    #[packed_field(bits = "11")]
    pub not_global: bool,
    #[packed_field(bits = "10")]
    pub access_flag: bool,
    #[packed_field(bits = "9..=8", ty = "enum")]
    pub shareability: Shareability,
    #[packed_field(bits = "7..=6", ty = "enum")]
    pub data_access_permissions: DataAccessPermissions,
    #[packed_field(bits = "5")]
    pub non_secure: bool,
    #[packed_field(bits = "4..=2", ty = "enum")]
    pub mem_attr_index: MemoryAttributesIndex,
}

impl From<Attributes> for u64 {
    fn from(attributes: Attributes) -> Self {
        u64::from_be_bytes(attributes.pack().unwrap())
    }
}

impl From<u64> for Attributes {
    fn from(bits: u64) -> Self {
        Self::unpack(&bits.to_be_bytes()).unwrap()
    }
}

/// Next level attributes
#[derive(PackedStruct, Clone, Debug, PartialEq, Eq, Default)]
#[packed_struct(size_bytes = "8", bit_numbering = "lsb0")]
pub struct NextLevelAttributes {
    #[packed_field(bits = "63")]
    ns_table: bool,
    #[packed_field(bits = "62..=61")]
    ap_table: Integer<u8, packed_bits::Bits<2>>,
    #[packed_field(bits = "60")]
    xn_table: bool,
    #[packed_field(bits = "59")]
    pxn_table: bool,
}

impl From<NextLevelAttributes> for u64 {
    fn from(attributes: NextLevelAttributes) -> Self {
        u64::from_be_bytes(attributes.pack().unwrap())
    }
}

impl From<u64> for NextLevelAttributes {
    fn from(bits: u64) -> Self {
        Self::unpack(&bits.to_be_bytes()).unwrap()
    }
}

/// Memory descriptor type
#[derive(PartialEq, Eq, Debug)]
pub enum DescriptorType {
    Invalid,
    Block,
    Table,
}

/// Memory descriptor of a memory translation table
#[repr(C)]
#[derive(Default)]
pub struct Descriptor {
    cell: UnsafeCell<u64>,
}

impl Descriptor {
    const ATTR_MASK: u64 = 0xfff8_0000_0000_0ffc;
    const DESCRIPTOR_TYPE_MASK: u64 = 0b11;
    const INVALID_DESCRIPTOR_VALUE: u64 = 0x0;
    const NEXT_ATTR_MASK: u64 = 0xf800_0000_0000_0000;
    const TABLE_BIT: u64 = 0b10;
    const TA_MASK: u64 = 0x0000_ffff_ffff_f000;
    const VALID_BIT: u64 = 0b01;

    /// Query descriptor type
    pub fn get_descriptor_type(&self, level: isize) -> DescriptorType {
        assert!(level <= 3);

        let desc_type_bits = self.get() & Self::DESCRIPTOR_TYPE_MASK;
        if desc_type_bits & Self::VALID_BIT != 0 {
            if level == 3 {
                assert_eq!(Self::TABLE_BIT, desc_type_bits & Self::TABLE_BIT);
                DescriptorType::Block
            } else if desc_type_bits & Self::TABLE_BIT != 0 {
                DescriptorType::Table
            } else {
                DescriptorType::Block
            }
        } else {
            DescriptorType::Invalid
        }
    }

    // Invalid descriptor functions

    /// Check if it is a valid descriptor
    pub fn is_valid(&self) -> bool {
        self.get() & Self::VALID_BIT != 0
    }

    // Block descriptor functions

    /// Set block descriptor
    pub fn set_block_descriptor<const VA_BITS: usize>(
        &mut self,
        granule: TranslationGranule<VA_BITS>,
        level: isize,
        output_address: PhysicalAddress,
        attributes: Attributes,
    ) {
        let attr: u64 = attributes.into();

        assert!(level <= 3);
        assert!(self.get_descriptor_type(level) != DescriptorType::Table);
        assert_eq!(0, attr & !Self::ATTR_MASK);

        let oa_granule_mask = !(granule.block_size_at_level(level) - 1);

        // Figure D8-14 VMSAv8-64 Block descriptor formats
        let oa_bits = match granule {
            // 4KB and 16KB granules, 52-bit OA
            #[cfg(feature = "feat_lpa2")]
            TranslationGranule::Granule4k | TranslationGranule::Granule16k => {
                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_f000;
                assert_eq!(0, output_address.0 & !oa_mask);

                let address = output_address.0 & oa_mask;

                // OA[49:n] remains in place, OA[51:50] is mapped to [9:8]
                let lsbs = address & 0x0003_ffff_ffff_f000;
                let msbs = ((address >> 50) & 0x3) << 8;

                lsbs | msbs
            }

            // 64KB granule, 52-bit OA
            #[cfg(feature = "feat_lpa")]
            TranslationGranule::Granule64k => {
                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_0000;
                assert_eq!(0, output_address.0 & !oa_mask);

                let address = output_address.0 & oa_mask;

                // OA[47:n] remains in place, OA[51:48] is mapped to [15:12]
                let lsbs = address & 0x0000_ffff_ffff_0000;
                let msbs = ((address >> 48) & 0xf) << 12;

                lsbs | msbs
            }

            // 4KB, 16KB, and 64KB granules, 48-bit OA
            #[cfg(not(all(feature = "feat_lpa", feature = "feat_lpa2")))]
            _ => {
                let oa_mask = oa_granule_mask & 0x0000_ffff_ffff_f000;
                assert_eq!(0, output_address.0 & !oa_mask);
                output_address.0 & oa_mask
            }
        };

        let table_bit = if level < 3 { 0 } else { Self::TABLE_BIT };

        self.set(Self::VALID_BIT | table_bit | oa_bits as u64 | attr);
    }

    /// Get output address from the block descriptor
    pub fn get_block_output_address<const VA_BITS: usize>(
        &self,
        granule: TranslationGranule<VA_BITS>,
        level: isize,
    ) -> PhysicalAddress {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        let oa_granule_mask = !(granule.block_size_at_level(level) - 1);
        let descriptor_value = self.get();

        // Figure D8-14 VMSAv8-64 Block descriptor formats
        let pa = match granule {
            // 4KB and 16KB granules, 52-bit OA
            #[cfg(feature = "feat_lpa2")]
            TranslationGranule::Granule4k | TranslationGranule::Granule16k => {
                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_f000;

                // OA[49:n] remains in place, OA[51:50] is mapped from [9:8]
                let lsbs = descriptor_value & 0x0003_ffff_ffff_f000;
                let msbs = ((descriptor_value >> 8) & 0x3) << 50;

                (lsbs | msbs) as usize & oa_mask
            }

            // 64KB granule, 52-bit OA
            #[cfg(feature = "feat_lpa")]
            TranslationGranule::Granule64k => {
                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_0000;

                // OA[47:n] remains in place, OA[51:48] is mapped from [15:12]
                let lsbs = descriptor_value & 0x0003_ffff_ffff_0000;
                let msbs = ((descriptor_value >> 12) & 0xf) << 48;

                (lsbs | msbs) as usize & oa_mask
            }

            // 4KB, 16KB, and 64KB granules, 48-bit OA
            #[cfg(not(all(feature = "feat_lpa", feature = "feat_lpa2")))]
            _ => {
                let oa_mask = oa_granule_mask & 0x0000_ffff_ffff_f000;
                descriptor_value as usize & oa_mask
            }
        };

        PhysicalAddress(pa)
    }

    /// Set the attributes of the block descriptor
    pub fn set_block_attributes(&mut self, level: isize, attributes: Attributes) {
        assert!(level <= 3);
        let attr: u64 = attributes.into();
        assert_eq!(0, attr & !Self::ATTR_MASK);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        self.modify(|d| (d & !Self::ATTR_MASK) | attr);
    }

    /// Get the attributes of the block descriptor
    pub fn get_block_attributes(&self, level: isize) -> Attributes {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        Attributes::from(self.get() & Self::ATTR_MASK)
    }

    /// Set block descriptor to invalid
    pub fn set_block_descriptor_to_invalid(&mut self, level: isize) {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        self.set(Self::INVALID_DESCRIPTOR_VALUE)
    }

    /// Set block or invalid descriptor to invalid
    pub fn set_block_or_invalid_descriptor_to_invalid(&mut self, level: isize) {
        assert!(level <= 3);
        assert!(DescriptorType::Table != self.get_descriptor_type(level));

        self.set(Self::INVALID_DESCRIPTOR_VALUE)
    }

    /// Set table descriptor
    pub fn set_table_descriptor(
        &mut self,
        level: isize,
        next_level_table: PhysicalAddress,
        next_level_attributes: Option<NextLevelAttributes>,
    ) {
        assert!(level <= 2);
        assert!(self.get_descriptor_type(level) != DescriptorType::Table);

        let table_addr = next_level_table.0 as u64;
        assert_eq!(0, table_addr & !Self::TA_MASK);

        let mut raw_desc_value = Self::VALID_BIT | Self::TABLE_BIT | table_addr;

        if let Some(next_attr) = next_level_attributes {
            let next_attr_bits: u64 = next_attr.into();
            assert_eq!(0, next_attr_bits & !Self::NEXT_ATTR_MASK);
            raw_desc_value |= next_attr_bits;
        }

        self.set(raw_desc_value);
    }

    /// Get next level table
    pub fn get_next_level_table(&self, level: isize) -> PhysicalAddress {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        PhysicalAddress((self.get() & Self::TA_MASK) as usize)
    }

    /// Get next level attributes
    pub fn get_next_level_attributes(&self, level: isize) -> NextLevelAttributes {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        NextLevelAttributes::from(self.get() & Self::NEXT_ATTR_MASK)
    }

    /// Set table descriptor to invalid
    pub fn set_table_descriptor_to_invalid(&mut self, level: isize) -> PhysicalAddress {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let pa = PhysicalAddress((self.get() & Self::TA_MASK) as usize);
        self.set(Self::INVALID_DESCRIPTOR_VALUE);
        pa
    }

    /// Get raw descriptor value
    fn get(&self) -> u64 {
        unsafe { ptr::read_volatile(self.cell.get()) }
    }

    /// Set raw descriptor value
    fn set(&mut self, value: u64) {
        // Safety: The cell should point to a valid address and the assembly code is just a data
        // barrier.
        unsafe {
            ptr::write_volatile(self.cell.get(), value);
            #[cfg(target_arch = "aarch64")]
            core::arch::asm!("dsb nsh");
        }
    }

    /// Modify raw descriptor value
    fn modify<F>(&mut self, f: F)
    where
        F: Fn(u64) -> u64,
    {
        self.set(f(self.get()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_attributes() {
        let attributes = Attributes::default();
        assert_eq!(0u64, attributes.into());

        let attributes = Attributes {
            uxn: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 54, attributes.into());

        let attributes = Attributes {
            pxn: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 53, attributes.into());

        let attributes = Attributes {
            contiguous: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 52, attributes.into());

        let attributes = Attributes {
            not_global: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 11, attributes.into());

        let attributes = Attributes {
            access_flag: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 10, attributes.into());

        let attributes = Attributes {
            non_secure: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 5, attributes.into());

        let attributes = Attributes {
            mem_attr_index: MemoryAttributesIndex::Normal_IWBWA_OWBWA,
            ..Default::default()
        };
        assert_eq!(
            (MemoryAttributesIndex::Normal_IWBWA_OWBWA as u64) << 2,
            attributes.into()
        );

        let attributes: Attributes = 0.into();
        assert!(!attributes.uxn);
        assert!(!attributes.pxn);
        assert!(!attributes.contiguous);
        assert!(!attributes.not_global);
        assert!(!attributes.access_flag);
        assert_eq!(Shareability::NonShareable, attributes.shareability);
        assert_eq!(
            DataAccessPermissions::ReadWrite_None,
            attributes.data_access_permissions
        );
        assert!(!attributes.non_secure);
        assert_eq!(
            MemoryAttributesIndex::Device_nGnRnE,
            attributes.mem_attr_index
        );
    }

    #[test]
    fn test_next_level_attributes() {
        let next_level_attributes = NextLevelAttributes::default();
        assert_eq!(0u64, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            ns_table: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 63, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            ap_table: 3.into(),
            ..Default::default()
        };
        assert_eq!(3u64 << 61, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            xn_table: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 60, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            pxn_table: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 59, next_level_attributes.into());

        let next_level_attributes: NextLevelAttributes = 0.into();
        assert!(!next_level_attributes.ns_table);
        assert_eq!(0u8, next_level_attributes.ap_table.into());
        assert!(!next_level_attributes.xn_table);
        assert!(!next_level_attributes.pxn_table);

        let next_level_attributes: NextLevelAttributes = u64::MAX.into();
        assert!(next_level_attributes.ns_table);
        assert_eq!(3u8, next_level_attributes.ap_table.into());
        assert!(next_level_attributes.xn_table);
        assert!(next_level_attributes.pxn_table);
    }

    #[test]
    fn test_descriptor_get_type() {
        let descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };
        assert_eq!(DescriptorType::Invalid, descriptor.get_descriptor_type(1));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };
        assert_eq!(DescriptorType::Block, descriptor.get_descriptor_type(1));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };
        assert_eq!(DescriptorType::Table, descriptor.get_descriptor_type(1));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };
        assert_eq!(DescriptorType::Invalid, descriptor.get_descriptor_type(3));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };
        assert_eq!(DescriptorType::Block, descriptor.get_descriptor_type(3));
    }

    #[test]
    fn test_descriptor_is_valid() {
        let descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };
        assert!(!descriptor.is_valid());

        let descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };
        assert!(descriptor.is_valid());
    }

    #[test]
    fn test_descriptor_set_block_to_block_again() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };

        descriptor.set_block_descriptor::<48>(
            TranslationGranule::Granule4k,
            1,
            PhysicalAddress(0),
            Attributes::default(),
        );
        assert_eq!(0x1, descriptor.get());
    }

    #[test]
    #[should_panic]
    fn test_descriptor_set_block_invalid_oa() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor::<48>(
            TranslationGranule::Granule4k,
            1,
            PhysicalAddress(1 << 63),
            Attributes::default(),
        );
    }

    #[test]
    fn test_descriptor_block() {
        let granule: TranslationGranule<48> = TranslationGranule::Granule4k;
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor(
            granule,
            1,
            PhysicalAddress(0x0000000f_c0000000),
            Attributes {
                uxn: true,
                ..Default::default()
            },
        );
        assert_eq!(0x0040000f_c0000001, descriptor.get());

        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor(
            granule,
            3,
            PhysicalAddress(0x0000000f_fffff000),
            Attributes {
                uxn: true,
                ..Default::default()
            },
        );
        assert_eq!(0x0040000f_fffff003, descriptor.get());

        assert_eq!(
            PhysicalAddress(0x0000000f_fffff000),
            descriptor.get_block_output_address(granule, 3)
        );
        assert_eq!(
            Attributes {
                uxn: true,
                ..Default::default()
            },
            descriptor.get_block_attributes(3)
        );

        descriptor.set_block_attributes(
            3,
            Attributes {
                pxn: true,
                ..Default::default()
            },
        );
        assert_eq!(
            Attributes {
                pxn: true,
                ..Default::default()
            },
            descriptor.get_block_attributes(3)
        );
    }

    #[test]
    #[should_panic]
    fn test_descriptor_invalid_block_to_invalid() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor_to_invalid(0);
    }

    #[test]
    fn test_descriptor_block_to_invalid() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };

        descriptor.set_block_descriptor_to_invalid(3);
        assert_eq!(0, descriptor.get());
    }

    #[test]
    fn test_descriptor_block_or_invalid_block_to_invalid() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_or_invalid_descriptor_to_invalid(0);
        assert_eq!(0, descriptor.get());

        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };

        descriptor.set_block_or_invalid_descriptor_to_invalid(3);
        assert_eq!(0, descriptor.get());
    }

    #[test]
    #[should_panic]
    fn test_descriptor_level3_to_table() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_table_descriptor(3, PhysicalAddress(0), None);
    }

    #[test]
    fn test_descriptor_block_to_table() {
        let next_level_table = PhysicalAddress(0x1000);
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };

        descriptor.set_table_descriptor(0, next_level_table, None);
        assert_eq!(0x1003, descriptor.get());
    }

    #[test]
    #[should_panic]
    fn test_descriptor_table_non_aligned() {
        let next_level_table = PhysicalAddress(0x800);
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_table_descriptor(0, next_level_table, None);
    }

    #[test]
    fn test_descriptor_table() {
        let next_level_table = PhysicalAddress(0x0000_000c_ba98_7000);
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_table_descriptor(0, next_level_table, None);
        assert_eq!(0x0000_000c_ba98_7003, descriptor.get());
    }

    #[test]
    fn test_descriptor_table_next_level_attr() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let next_level_table = PhysicalAddress(0x0000_000c_ba98_7000);
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_table_descriptor(
            0,
            next_level_table,
            Some(NextLevelAttributes {
                ns_table: true,
                ..Default::default()
            }),
        );
        assert_eq!(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003, descriptor.get());
    }

    #[test]
    fn test_descriptor_table_get_next_level_table() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(NEXT_LEVEL_ADDR, descriptor.get_next_level_table(0).0 as u64);
    }

    #[test]
    fn test_descriptor_table_get_next_level_attr() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(
            NextLevelAttributes {
                ns_table: true,
                ..Default::default()
            },
            descriptor.get_next_level_attributes(0)
        );
    }

    #[test]
    fn test_descriptor_table_set_to_invalid() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(
            NEXT_LEVEL_ADDR,
            descriptor.set_table_descriptor_to_invalid(0).0 as u64
        );
        assert_eq!(0, descriptor.get());
    }

    #[test]
    fn test_descriptor_raw_interface() {
        let cell_value = 0x01234567_89abcdefu64;
        let cell_new_value = 0x12345678_9abcdef0u64;

        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(cell_value),
        };

        assert_eq!(cell_value, descriptor.get());

        descriptor.set(cell_new_value);
        assert_eq!(cell_new_value, descriptor.get());

        descriptor.modify(|d| d + 1);
        assert_eq!(cell_new_value + 1, descriptor.get());
    }
}