// 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::kernel_space::KernelSpace;
use crate::MemoryAttributesIndex;

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

/// Data access permission
#[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;
    pub const GRANULE_SIZES: [usize; 4] = [0, 0x4000_0000, 0x0020_0000, 0x0000_1000];
    const INVALID_DESCRIPTOR_VALUE: u64 = 0x0;
    const NEXT_ATTR_MASK: u64 = 0xf800_0000_0000_0000;
    const OA_MASK: u64 = 0x0000_ffff_ffff_f000;
    const TABLE_BIT: u64 = 0b10;
    const TABLE_ENTRY_COUNT: usize = 512;
    const TA_MASK: u64 = 0x0000_ffff_ffff_f000;
    const VALID_BIT: u64 = 0b01;

    /// Query descriptor type
    pub fn get_descriptor_type(&self, level: usize) -> 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(
        &mut self,
        level: usize,
        output_address: usize,
        attributes: Attributes,
    ) {
        let attr: u64 = attributes.into();

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

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

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

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

        (self.get() & Self::OA_MASK) as usize
    }

    /// Set the attributes of the block descriptor
    pub fn set_block_attributes(&mut self, level: usize, 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: usize) -> 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: usize) {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        self.set(Self::INVALID_DESCRIPTOR_VALUE)
    }

    /// Set table descriptor
    ///
    /// **Unsafe**: The caller has to ensure that the passed next level table has the same life as
    /// the descriptor.
    pub unsafe fn set_table_descriptor(
        &mut self,
        level: usize,
        next_level_table: &mut [Descriptor],
        next_level_attributes: Option<NextLevelAttributes>,
    ) {
        assert!(level <= 2);
        assert_eq!(Self::TABLE_ENTRY_COUNT, next_level_table.len());
        assert!(self.get_descriptor_type(level) != DescriptorType::Table);

        let table_addr = KernelSpace::kernel_to_pa(next_level_table.as_ptr() 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
    ///
    /// **Unsafe**: The returned next level table is based on the address read from the descriptor.
    /// The caller has to ensure that no other references are being used of the table.
    pub unsafe fn get_next_level_table(&self, level: usize) -> &[Descriptor] {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let table_address =
            KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *const Descriptor;
        core::slice::from_raw_parts(table_address, Self::TABLE_ENTRY_COUNT)
    }

    /// Get mutable next level table
    ///
    /// **Unsafe**: The returned next level table is based on the address read from the descriptor.
    /// The caller has to ensure that no other references are being used of the table.
    pub unsafe fn get_next_level_table_mut(&mut self, level: usize) -> &mut [Descriptor] {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let table_address =
            KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *mut Descriptor;
        core::slice::from_raw_parts_mut(table_address, Self::TABLE_ENTRY_COUNT)
    }

    /// Get next level attributes
    pub fn get_next_level_attributes(&self, level: usize) -> 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
    ///
    /// **Unsafe:** The returned descriptor reference must be released by the caller, i.e. release
    /// to `PagePool`
    pub unsafe fn set_table_descriptor_to_invalid(&mut self, level: usize) -> &mut [Descriptor] {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let table_address =
            KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *mut Descriptor;
        self.set(Self::INVALID_DESCRIPTOR_VALUE);
        core::slice::from_raw_parts_mut(table_address, Self::TABLE_ENTRY_COUNT)
    }

    /// 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) {
        unsafe { ptr::write_volatile(self.cell.get(), value) }
    }

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

    /// Get output address mask
    fn get_oa_mask(level: usize) -> usize {
        Self::OA_MASK as usize & !(Self::GRANULE_SIZES[level] - 1)
    }
}

#[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(1, 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(1, 1 << 63, Attributes::default());
}

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

    descriptor.set_block_descriptor(
        1,
        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(
        3,
        0x0000000f_fffff000,
        Attributes {
            uxn: true,
            ..Default::default()
        },
    );
    assert_eq!(0x0040000f_fffff003, descriptor.get());

    assert_eq!(0x0000000f_fffff000, descriptor.get_block_output_address(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]
#[should_panic]
fn test_descriptor_level3_to_table() {
    let mut next_level_table = [Descriptor {
        cell: UnsafeCell::new(0),
    }];
    let mut descriptor = Descriptor {
        cell: UnsafeCell::new(0),
    };

    unsafe {
        descriptor.set_table_descriptor(3, &mut next_level_table, None);
    }
}

#[test]
fn test_descriptor_block_to_table() {
    let next_level_table =
        unsafe { core::slice::from_raw_parts_mut(0x1000 as *mut Descriptor, 512) };
    let mut descriptor = Descriptor {
        cell: UnsafeCell::new(1),
    };

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

#[test]
#[should_panic]
fn test_descriptor_table_invalid_count() {
    let next_level_table =
        unsafe { core::slice::from_raw_parts_mut(0x800 as *mut Descriptor, 511) };
    let mut descriptor = Descriptor {
        cell: UnsafeCell::new(0),
    };

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

#[test]
#[should_panic]
fn test_descriptor_table_non_aligned() {
    let next_level_table =
        unsafe { core::slice::from_raw_parts_mut(0x800 as *mut Descriptor, 512) };
    let mut descriptor = Descriptor {
        cell: UnsafeCell::new(0),
    };

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

#[test]
fn test_descriptor_table() {
    let next_level_table =
        unsafe { core::slice::from_raw_parts_mut(0x0000_000c_ba98_7000 as *mut Descriptor, 512) };
    let mut descriptor = Descriptor {
        cell: UnsafeCell::new(0),
    };

    unsafe {
        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 =
        unsafe { core::slice::from_raw_parts_mut(NEXT_LEVEL_ADDR as *mut Descriptor, 512) };
    let mut descriptor = Descriptor {
        cell: UnsafeCell::new(0),
    };

    unsafe {
        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!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
        descriptor.get_next_level_table(0).as_ptr() as u64
    });
}

#[test]
fn test_descriptor_table_get_next_level_table_mut() {
    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!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
        descriptor.get_next_level_table_mut(0).as_ptr() as *mut Descriptor 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!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
        descriptor.set_table_descriptor_to_invalid(0).as_ptr() as *mut Descriptor 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());
}