// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Copyright © 2019 Intel Corporation
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE-BSD-3-Clause file.
//
// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause

use vm_memory::{GuestAddress, GuestUsize};

use crate::address::AddressAllocator;
#[cfg(target_arch = "x86_64")]
use crate::gsi::GsiApic;
use crate::gsi::{GsiAllocator, InterruptAllocError};
use crate::page_size::get_page_size;

/// Manages allocating system resources such as address space and interrupt numbers.
///
/// # Example - Use the `SystemAddress` builder.
///
/// ```rust
/// # #[cfg(target_arch = "x86_64")]
/// # use vm_allocator::{GsiApic, SystemAllocator};
/// # #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
/// # use vm_allocator::SystemAllocator;
/// # use vm_memory::{Address, GuestAddress, GuestUsize};
///   let mut allocator = SystemAllocator::new(
///           GuestAddress(0x1000),
///           0x10000,
///           GuestAddress(0x10000000), 0x10000000,
///           #[cfg(target_arch = "x86_64")] &[GsiApic::new(5, 19)]).unwrap();
///   #[cfg(target_arch = "x86_64")]
///   assert_eq!(allocator.allocate_irq(), Ok(5));
///   #[cfg(target_arch = "aarch64")]
///   assert_eq!(allocator.allocate_irq(), Ok(32));
///   #[cfg(target_arch = "riscv64")]
///   assert_eq!(allocator.allocate_irq(), Ok(0));
///   #[cfg(target_arch = "x86_64")]
///   assert_eq!(allocator.allocate_irq(), Ok(6));
///   #[cfg(target_arch = "aarch64")]
///   assert_eq!(allocator.allocate_irq(), Ok(33));
///   #[cfg(target_arch = "riscv64")]
///   assert_eq!(allocator.allocate_irq(), Ok(1));
///   assert_eq!(allocator.allocate_platform_mmio_addresses(None, 0x1000, Some(0x1000)), Some(GuestAddress(0x1fff_f000)));
///
/// ```
pub struct SystemAllocator {
    io_address_space: AddressAllocator,
    platform_mmio_address_space: AddressAllocator,
    gsi_allocator: GsiAllocator,
}

impl SystemAllocator {
    /// Creates a new `SystemAllocator` for managing addresses and irq numbers.
    /// Can return `None` if `base` + `size` overflows a u64
    ///
    /// * `io_base` - (X86) The starting address of IO memory.
    /// * `io_size` - (X86) The size of IO memory.
    /// * `platform_mmio_base` - The starting address of platform MMIO memory.
    /// * `platform_mmio_size` - The size of platform MMIO memory.
    /// * `apics` - (X86) slice of APIC's.
    ///
    pub fn new(
        io_base: GuestAddress,
        io_size: GuestUsize,
        platform_mmio_base: GuestAddress,
        platform_mmio_size: GuestUsize,
        #[cfg(target_arch = "x86_64")] apics: &[GsiApic],
    ) -> Option<Self> {
        Some(SystemAllocator {
            io_address_space: AddressAllocator::new(io_base, io_size)?,
            platform_mmio_address_space: AddressAllocator::new(
                platform_mmio_base,
                platform_mmio_size,
            )?,
            #[cfg(target_arch = "x86_64")]
            gsi_allocator: GsiAllocator::new(apics),
            #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
            gsi_allocator: GsiAllocator::new(),
        })
    }

    /// Reserves the next available system irq number.
    pub fn allocate_irq(&mut self) -> Result<u32, InterruptAllocError> {
        self.gsi_allocator.allocate_irq()
    }

    /// Reserves the next available GSI.
    pub fn allocate_gsi(&mut self) -> Result<u32, InterruptAllocError> {
        self.gsi_allocator.allocate_gsi()
    }

    /// Frees the given GSI.
    pub fn free_gsi(&mut self, vector: u32) -> Result<(), InterruptAllocError> {
        self.gsi_allocator.free_gsi(vector)
    }

    /// Reserves a section of `size` bytes of IO address space.
    pub fn allocate_io_addresses(
        &mut self,
        address: Option<GuestAddress>,
        size: GuestUsize,
        align_size: Option<GuestUsize>,
    ) -> Option<GuestAddress> {
        self.io_address_space
            .allocate(address, size, Some(align_size.unwrap_or(0x1)))
    }

    /// Reserves a section of `size` bytes of platform MMIO address space.
    pub fn allocate_platform_mmio_addresses(
        &mut self,
        address: Option<GuestAddress>,
        size: GuestUsize,
        align_size: Option<GuestUsize>,
    ) -> Option<GuestAddress> {
        self.platform_mmio_address_space.allocate(
            address,
            size,
            Some(align_size.unwrap_or_else(get_page_size)),
        )
    }

    /// Free an IO address range.
    /// We can only free a range if it matches exactly an already allocated range.
    pub fn free_io_addresses(&mut self, address: GuestAddress, size: GuestUsize) {
        self.io_address_space.free(address, size);
    }

    /// Free a platform MMIO address range.
    /// We can only free a range if it matches exactly an already allocated range.
    pub fn free_platform_mmio_addresses(&mut self, address: GuestAddress, size: GuestUsize) {
        self.platform_mmio_address_space.free(address, size);
    }
}