// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! Manages MMIO range partitioning between VTLs.

use super::PartitionInfo;
use super::dt::DtError;
use memory_range::MemoryRange;

/// The start address of MMIO high range.
const MMIO_HIGH_RANGE_START: u64 = 1 << 32;

impl PartitionInfo {
    /// Select the mmio range that VTL2 should use from looking at VTL0 mmio
    /// ranges.
    ///
    /// VTL2 MMIO is partitioned such that:
    /// - All MMIO low range is assigned to VTL0.
    /// - VTL2_MMIO_HIGH_RANGE_SIZE bytes from the end of the high range is
    ///   assigned to VTL2.
    /// - The remaining high range is assigned to VTL0.
    ///
    /// Assumes input ranges are non-overlapping and in increasing address
    /// order.
    ///
    /// On success, returns the mmio that VTL2 should use.
    ///
    /// Returns an error if the input VTL0 MMIO range is invalid or if the VTL2
    /// allocation amount was not satisfied due to a lack of high MMIO assigned
    /// to VTL0.
    pub fn select_vtl2_mmio_range(&self, vtl2_size: u64) -> Result<MemoryRange, DtError> {
        // Iterate over the list of MMIO ranges in reverse address order so that
        // the VTL2 range is carved out from the end.
        for range in self.vmbus_vtl0.mmio.iter().rev() {
            // Do not select low MMIO ranges for VTL2.
            if range.start() < MMIO_HIGH_RANGE_START {
                continue;
            }

            // Compute the length of the VTL2 subrange. If there is not enough
            // mmio, give up.
            if range.len() < vtl2_size {
                return Err(DtError::NotEnoughMmio);
            }

            let vtl2_range_start = range.end() - vtl2_size;

            return Ok(MemoryRange::new(vtl2_range_start..range.end()));
        }

        Err(DtError::NotEnoughMmio)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use arrayvec::ArrayVec;
    use memory_range::subtract_ranges;

    /// The size (in bytes) of MMIO high range assigned to VTL2.
    const VTL2_MMIO_HIGH_RANGE_SIZE: u64 = 128 * (1 << 20);

    // Tests that MMIO range is partitioned correctly between VTL0 and VTL2
    // for a variety of input ranges.
    #[test]
    fn mmio_range_partitioned_correctly() {
        #[derive(Debug)]
        struct TestCase {
            // Input
            mmio: ArrayVec<MemoryRange, 2>,
            // Expected output
            succeeds: bool,
            vtl0_range: ArrayVec<MemoryRange, 2>,
            vtl2_range: MemoryRange,
        }

        let testcases = vec![
            TestCase {
                // No MMIO range is provided, fails.
                mmio: ArrayVec::new(),
                succeeds: false,
                vtl0_range: ArrayVec::new(),
                vtl2_range: MemoryRange::EMPTY,
            },
            TestCase {
                // Only low mmio, fails.
                mmio: ArrayVec::from([
                    MemoryRange::new(0x3000_0000..0x4000_0000),
                    MemoryRange::new(0x4000_0000..0x5000_0000),
                ]),
                succeeds: false,
                vtl0_range: ArrayVec::new(),
                vtl2_range: MemoryRange::EMPTY,
            },
            TestCase {
                // MMIO high range is less than what VTL2 requested.
                mmio: ArrayVec::from([
                    MemoryRange::new(0x3000_0000..0x4000_0000),
                    MemoryRange::new(
                        MMIO_HIGH_RANGE_START
                            ..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE / 2),
                    ),
                ]),
                succeeds: false,
                vtl0_range: ArrayVec::new(),
                vtl2_range: MemoryRange::EMPTY,
            },
            TestCase {
                // MMIO high range is just enough for VTL2.
                // Low range should be assigned to VTL0.
                // High range should be assigned to VTL2.
                mmio: ArrayVec::from([
                    MemoryRange::new(0x3000_0000..0x4000_0000),
                    MemoryRange::new(
                        MMIO_HIGH_RANGE_START..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE),
                    ),
                ]),
                succeeds: true,
                vtl0_range: [MemoryRange::new(0x3000_0000..0x4000_0000)]
                    .into_iter()
                    .collect(),
                vtl2_range: MemoryRange::new(
                    MMIO_HIGH_RANGE_START..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE),
                ),
            },
            TestCase {
                // MMIO high range is more than what VTL2 requested.
                // VTL2 should be assigned SIZE from the end of the high range.
                // VTL0 should be assigned the remaining high range.
                mmio: ArrayVec::from([
                    MemoryRange::new(0x3000_0000..0x4000_0000),
                    MemoryRange::new(
                        MMIO_HIGH_RANGE_START
                            ..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE * 4),
                    ),
                ]),
                succeeds: true,
                vtl0_range: ArrayVec::from([
                    MemoryRange::new(0x3000_0000..0x4000_0000),
                    MemoryRange::new(
                        MMIO_HIGH_RANGE_START
                            ..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE * 3),
                    ),
                ]),
                vtl2_range: MemoryRange::new(
                    MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE * 3
                        ..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE * 4),
                ),
            },
            TestCase {
                // Multiple MMIO high ranges are provided.
                // VTL2 should be assigned SIZE from the very end of the high range.
                // VTL0 should be assigned all the remaining high ranges.
                mmio: ArrayVec::from([
                    MemoryRange::new(
                        MMIO_HIGH_RANGE_START
                            ..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE * 2),
                    ),
                    MemoryRange::new(
                        (MMIO_HIGH_RANGE_START * 2)
                            ..(MMIO_HIGH_RANGE_START * 2 + VTL2_MMIO_HIGH_RANGE_SIZE * 2),
                    ),
                ]),
                succeeds: true,
                vtl0_range: ArrayVec::from([
                    MemoryRange::new(
                        MMIO_HIGH_RANGE_START
                            ..(MMIO_HIGH_RANGE_START + VTL2_MMIO_HIGH_RANGE_SIZE * 2),
                    ),
                    MemoryRange::new(
                        (MMIO_HIGH_RANGE_START * 2)
                            ..(MMIO_HIGH_RANGE_START * 2 + VTL2_MMIO_HIGH_RANGE_SIZE),
                    ),
                ]),
                vtl2_range: MemoryRange::new(
                    (MMIO_HIGH_RANGE_START * 2 + VTL2_MMIO_HIGH_RANGE_SIZE)
                        ..(MMIO_HIGH_RANGE_START * 2 + VTL2_MMIO_HIGH_RANGE_SIZE * 2),
                ),
            },
        ];

        // Run all test cases.
        for (i, tc) in testcases.iter().enumerate() {
            let mut vtl2_info = PartitionInfo::new();
            vtl2_info.vmbus_vtl0.mmio.clone_from(&tc.mmio);

            let result = vtl2_info.select_vtl2_mmio_range(VTL2_MMIO_HIGH_RANGE_SIZE);

            assert_eq!(
                tc.succeeds,
                result.is_ok(),
                "test case #{i}: unexpected result"
            );

            if tc.succeeds {
                let vtl2_mmio = result.unwrap();
                let vtl0_mmio =
                    subtract_ranges(tc.mmio.iter().cloned(), [vtl2_mmio]).collect::<Vec<_>>();

                assert_eq!(
                    tc.vtl0_range.as_slice(),
                    vtl0_mmio.as_slice(),
                    "test case #{i}: vtl0 was assigned an unexpected mmio range"
                );
                assert_eq!(
                    tc.vtl2_range, vtl2_mmio,
                    "test case #{i}: vtl1 was assigned an unexpected mmio range"
                );
            }
        }
    }
}
