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

#[derive(Clone, Copy, Debug)]
pub enum TranslationGranule<const VA_BITS: usize> {
    Granule4k = 4 * 1024,
    Granule16k = 16 * 1024,
    Granule64k = 64 * 1024,
}

impl<const VA_BITS: usize> TranslationGranule<VA_BITS> {
    pub const fn initial_lookup_level(&self) -> isize {
        match self {
            TranslationGranule::Granule4k => match VA_BITS {
                #[cfg(feature = "feat_ttst")]
                16..=21 => 3,
                #[cfg(feature = "feat_ttst")]
                22..=24 => 2,
                25..=30 => 2,
                31..=39 => 1,
                40..=48 => 0,
                #[cfg(feature = "feat_lpa2")]
                49..=52 => -1,
                _ => panic!("Invalid VA_BITS"),
            },
            TranslationGranule::Granule16k => match VA_BITS {
                #[cfg(feature = "feat_ttst")]
                16..=24 => 3,
                25 => 3,
                26..=36 => 2,
                37..=47 => 1,
                48 => 0,
                #[cfg(feature = "feat_lpa2")]
                49..=52 => 0,
                _ => panic!("Invalid VA_BITS"),
            },
            TranslationGranule::Granule64k => match VA_BITS {
                #[cfg(feature = "feat_ttst")]
                17..=24 => 3,
                25..=29 => 3,
                30..=42 => 2,
                43..=48 => 1,
                #[cfg(feature = "feat_lva")]
                49..=52 => 1,
                _ => panic!("Invalid VA_BITS"),
            },
        }
    }

    pub const fn entry_count_at_level(&self, level: isize) -> usize {
        let total_bits_at_level = self.total_bits_at_level(level);
        let bits_per_level = match self {
            TranslationGranule::Granule4k => 9,
            TranslationGranule::Granule16k => 11,
            TranslationGranule::Granule64k => 13,
        };

        let bits = if total_bits_at_level + bits_per_level < VA_BITS {
            total_bits_at_level + bits_per_level
        } else {
            VA_BITS
        } - total_bits_at_level;

        1 << bits
    }

    pub const fn table_size<D>(&self, level: isize) -> usize {
        self.entry_count_at_level(level) * core::mem::size_of::<D>()
    }

    pub const fn table_alignment<D>(&self, level: isize) -> usize {
        // D8.2.5.2
        // For the VMSAv8-64 translation system, if the translation table has fewer than eight
        // entries and an OA size greater than 48 bits is used, then the table is aligned to 64
        // bytes. Otherwise, the translation table is aligned to the size of that translation table.

        let alignment = self.table_size::<D>(level);

        if VA_BITS > 48 && alignment < 64 {
            64
        } else {
            alignment
        }
    }

    pub const fn block_size_at_level(&self, level: isize) -> usize {
        1 << self.total_bits_at_level(level)
    }

    pub const fn total_bits_at_level(&self, level: isize) -> usize {
        assert!(self.initial_lookup_level() <= level && level <= 3);

        match self {
            TranslationGranule::Granule4k => match level {
                -1 => 48,
                0 => 39,
                1 => 30,
                2 => 21,
                3 => 12,
                _ => panic!("Invalid translation level"),
            },
            TranslationGranule::Granule16k => match level {
                0 => 47,
                1 => 36,
                2 => 25,
                3 => 14,
                _ => panic!("Invalid translation level"),
            },
            TranslationGranule::Granule64k => match level {
                1 => 42,
                2 => 29,
                3 => 16,
                _ => panic!("Invalid translation level"),
            },
        }
    }
}

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

    macro_rules! validate_initial_lookup_level {
        ( $granule:expr, $tnsz_min:literal, $tnsz_max:literal, $level:literal ) => {
            assert_eq!(
                $level,
                TranslationGranule::<{ 64 - $tnsz_min }>::initial_lookup_level(&$granule)
            );
            assert_eq!(
                $level,
                TranslationGranule::<{ 64 - $tnsz_max }>::initial_lookup_level(&$granule)
            );
        };
    }

    #[test]
    fn test_initial_lookup_level() {
        // Table D8-16 4KB granule, determining stage 1 initial lookup level
        #[cfg(feature = "feat_lpa2")]
        validate_initial_lookup_level!(TranslationGranule::Granule4k, 12, 15, -1);
        validate_initial_lookup_level!(TranslationGranule::Granule4k, 16, 24, 0);
        validate_initial_lookup_level!(TranslationGranule::Granule4k, 25, 33, 1);
        validate_initial_lookup_level!(TranslationGranule::Granule4k, 34, 39, 2);
        #[cfg(feature = "feat_ttst")]
        validate_initial_lookup_level!(TranslationGranule::Granule4k, 40, 42, 2);
        #[cfg(feature = "feat_ttst")]
        validate_initial_lookup_level!(TranslationGranule::Granule4k, 43, 48, 3);

        //// Table D8-26 16KB granule, determining stage 1 initial lookup level
        #[cfg(feature = "feat_lpa2")]
        validate_initial_lookup_level!(TranslationGranule::Granule16k, 12, 15, 0);
        validate_initial_lookup_level!(TranslationGranule::Granule16k, 16, 16, 0);
        validate_initial_lookup_level!(TranslationGranule::Granule16k, 17, 27, 1);
        validate_initial_lookup_level!(TranslationGranule::Granule16k, 28, 38, 2);
        validate_initial_lookup_level!(TranslationGranule::Granule16k, 39, 39, 3);
        #[cfg(feature = "feat_ttst")]
        validate_initial_lookup_level!(TranslationGranule::Granule16k, 40, 48, 3);
        //
        //// D8.2.10.1 VMSAv8-64 Stage 1 address translation using the 64KB translation granule
        #[cfg(feature = "feat_lva")]
        validate_initial_lookup_level!(TranslationGranule::Granule64k, 12, 15, 1);
        validate_initial_lookup_level!(TranslationGranule::Granule64k, 16, 21, 1);
        validate_initial_lookup_level!(TranslationGranule::Granule64k, 22, 34, 2);
        validate_initial_lookup_level!(TranslationGranule::Granule64k, 35, 39, 3);
        #[cfg(feature = "feat_ttst")]
        validate_initial_lookup_level!(TranslationGranule::Granule64k, 40, 47, 3);
    }

    #[test]
    fn test_entry_count_at_level() {
        // Table D8-14 4KB granule translation table properties at each lookup level
        #[cfg(feature = "feat_lpa2")]
        {
            let granule = TranslationGranule::<52>::Granule4k;
            assert_eq!(16, granule.entry_count_at_level(-1));
            assert_eq!(512, granule.entry_count_at_level(0));
            assert_eq!(512, granule.entry_count_at_level(1));
            assert_eq!(512, granule.entry_count_at_level(2));
            assert_eq!(512, granule.entry_count_at_level(3));
        }
        #[cfg(not(feature = "feat_lpa2"))]
        {
            let granule = TranslationGranule::<48>::Granule4k;
            assert_eq!(512, granule.entry_count_at_level(0));
            assert_eq!(512, granule.entry_count_at_level(1));
            assert_eq!(512, granule.entry_count_at_level(2));
            assert_eq!(512, granule.entry_count_at_level(3));
        }

        let granule = TranslationGranule::<42>::Granule4k;
        assert_eq!(8, granule.entry_count_at_level(0));
        assert_eq!(512, granule.entry_count_at_level(1));
        assert_eq!(512, granule.entry_count_at_level(2));
        assert_eq!(512, granule.entry_count_at_level(3));

        let granule = TranslationGranule::<32>::Granule4k;
        assert_eq!(4, granule.entry_count_at_level(1));
        assert_eq!(512, granule.entry_count_at_level(2));
        assert_eq!(512, granule.entry_count_at_level(3));

        let granule = TranslationGranule::<26>::Granule4k;
        assert_eq!(32, granule.entry_count_at_level(2));
        assert_eq!(512, granule.entry_count_at_level(3));

        // Table D8-24 16KB granule translation table properties at each lookup level
        #[cfg(feature = "feat_lpa2")]
        {
            let granule = TranslationGranule::<52>::Granule16k;
            assert_eq!(32, granule.entry_count_at_level(0));
            assert_eq!(2048, granule.entry_count_at_level(1));
            assert_eq!(2048, granule.entry_count_at_level(2));
            assert_eq!(2048, granule.entry_count_at_level(3));
        }
        #[cfg(not(feature = "feat_lpa2"))]
        {
            let granule = TranslationGranule::<48>::Granule16k;
            assert_eq!(2, granule.entry_count_at_level(0));
            assert_eq!(2048, granule.entry_count_at_level(1));
            assert_eq!(2048, granule.entry_count_at_level(2));
            assert_eq!(2048, granule.entry_count_at_level(3));
        }

        let granule = TranslationGranule::<40>::Granule16k;
        assert_eq!(16, granule.entry_count_at_level(1));
        assert_eq!(2048, granule.entry_count_at_level(2));
        assert_eq!(2048, granule.entry_count_at_level(3));

        let granule = TranslationGranule::<30>::Granule16k;
        assert_eq!(32, granule.entry_count_at_level(2));
        assert_eq!(2048, granule.entry_count_at_level(3));

        // Table D8-33 64KB granule translation table properties at each lookup level
        #[cfg(feature = "feat_lva")]
        {
            let granule = TranslationGranule::<52>::Granule64k;
            assert_eq!(1024, granule.entry_count_at_level(1));
        }

        let granule = TranslationGranule::<48>::Granule64k;
        #[cfg(not(feature = "feat_lva"))]
        {
            assert_eq!(64, granule.entry_count_at_level(1));
        }

        assert_eq!(8192, granule.entry_count_at_level(2));
        assert_eq!(8192, granule.entry_count_at_level(3));

        let granule = TranslationGranule::<40>::Granule64k;
        assert_eq!(2048, granule.entry_count_at_level(2));
        assert_eq!(8192, granule.entry_count_at_level(3));
    }

    #[test]
    fn test_granule_size_at_level() {
        // Table D8-15 4KB granule block descriptor properties at each lookup level
        #[cfg(feature = "feat_lpa2")]
        {
            let granule = TranslationGranule::<52>::Granule4k;
            assert_eq!(256 << 40, granule.block_size_at_level(-1));
        }

        let granule = TranslationGranule::<48>::Granule4k;
        assert_eq!(512 << 30, granule.block_size_at_level(0));
        assert_eq!(1 << 30, granule.block_size_at_level(1));
        assert_eq!(2 << 20, granule.block_size_at_level(2));
        assert_eq!(4 << 10, granule.block_size_at_level(3));

        // Table D8-25 16KB granule block descriptor properties at each lookup level
        let granule = TranslationGranule::<48>::Granule16k;
        assert_eq!(128 << 40, granule.block_size_at_level(0));
        assert_eq!(64 << 30, granule.block_size_at_level(1));
        assert_eq!(32 << 20, granule.block_size_at_level(2));
        assert_eq!(16 << 10, granule.block_size_at_level(3));

        // Table D8-34 64KB granule block descriptor properties at each lookup level
        let granule = TranslationGranule::<48>::Granule64k;
        assert_eq!(4 << 40, granule.block_size_at_level(1));
        assert_eq!(512 << 20, granule.block_size_at_level(2));
        assert_eq!(64 << 10, granule.block_size_at_level(3));
    }
}