mmtk 0.3.2

use super::*;
use crate::util::constants::BYTES_IN_PAGE;
use crate::util::heap::layout::vm_layout_constants::BYTES_IN_CHUNK;
use crate::util::memory;
use crate::util::{constants, Address};
use std::io::Result;
use std::sync::atomic::{AtomicU16, AtomicU32, AtomicU8, AtomicUsize, Ordering};

#[derive(Clone, Copy)]
pub enum SideMetadataScope {
    Global,
    PolicySpecific,
}

impl SideMetadataScope {
    pub const fn is_global(&self) -> bool {
        matches!(self, SideMetadataScope::Global)
    }
}

/// This struct stores the specification of a side metadata bit-set.
/// It is used as an input to the (inline) functions provided by the side metadata module.
///
/// Each plan or policy which uses a metadata bit-set, needs to create an instance of this struct.
///
/// For performance reasons, objects of this struct should be constants.
#[derive(Clone, Copy)]
pub struct SideMetadataSpec {
    pub scope: SideMetadataScope,
    pub offset: usize,
    pub log_num_of_bits: usize,
    pub log_min_obj_size: usize,
}

// ** NOTE: **
//  Regardless of the number of bits in a metadata unit, we always represent its content as a word.

/// Tries to map the required metadata space and returns `true` is successful.
///
/// # Arguments
///
/// * `start` - The starting address of the source data.
///
/// * `size` - The size of the source data (in bytes).
///
/// * `global_metadata_spec_vec` - A vector of SideMetadataSpec objects containing all global side metadata.
///
/// * `local_metadata_spec_vec` - A vector of SideMetadataSpec objects containing all local side metadata.
///
pub fn try_map_metadata_space(
    start: Address,
    size: usize,
    global_metadata_spec_vec: &[SideMetadataSpec],
    local_metadata_spec_vec: &[SideMetadataSpec],
) -> Result<()> {
    debug_assert!(start.is_aligned_to(BYTES_IN_PAGE));
    debug_assert!(size % BYTES_IN_PAGE == 0);

    for spec in global_metadata_spec_vec {
        let res = try_mmap_contiguous_metadata_space(start, size, spec, false);
        if res.is_err() {
            return res;
        }
    }

    #[cfg(target_pointer_width = "32")]
    let mut lsize: usize = 0;

    for spec in local_metadata_spec_vec {
        // For local side metadata, we always have to reserve address space for all
        // local metadata required by all policies in MMTk to be able to calculate a constant offset for each local metadata at compile-time
        // (it's like assigning an ID to each policy).
        // As the plan is chosen at run-time, we will never know which subset of policies will be used during run-time.
        // We can't afford this much address space in 32-bits.
        // So, we switch to the chunk-based approach for this specific case.
        //
        // The global metadata is different in that for each plan, we can calculate its constant base addresses at compile-time.
        // Using the chunk-based approach will need the same address space size as the current not-chunked approach.
        #[cfg(target_pointer_width = "64")]
        {
            let res = try_mmap_contiguous_metadata_space(start, size, spec, false);
            if res.is_err() {
                return res;
            }
        }
        #[cfg(target_pointer_width = "32")]
        {
            lsize += meta_bytes_per_chunk(spec.log_min_obj_size, spec.log_num_of_bits);
        }
    }

    #[cfg(target_pointer_width = "32")]
    if lsize > 0 {
        let max = BYTES_IN_CHUNK >> LOG_LOCAL_SIDE_METADATA_WORST_CASE_RATIO;
        debug_assert!(
            lsize <= max,
            "local side metadata per chunk (0x{:x}) must be less than (0x{:x})",
            lsize,
            max
        );
        return try_map_per_chunk_metadata_space(start, size, lsize, false);
    }

    Ok(())
}

/// Tries to map the required metadata address range, without reserving swap-space/physical memory for it.
/// This will make sure the address range is exclusive to the caller.
///
/// NOTE: Accessing addresses in this range will produce a segmentation fault if swap-space is not mapped using the `try_map_metadata_space` function.
pub fn try_map_metadata_address_range(
    start: Address,
    size: usize,
    global_metadata_spec_vec: &[SideMetadataSpec],
    local_metadata_spec_vec: &[SideMetadataSpec],
) -> Result<()> {
    info!(
        "try_map_metadata_address_range({}, 0x{:x}, {}, {})",
        start,
        size,
        global_metadata_spec_vec.len(),
        local_metadata_spec_vec.len()
    );
    debug_assert!(start.is_aligned_to(BYTES_IN_CHUNK));
    debug_assert!(size % BYTES_IN_CHUNK == 0);

    for spec in global_metadata_spec_vec {
        let res = try_mmap_contiguous_metadata_space(start, size, spec, true);
        if res.is_err() {
            return res;
        }
    }

    #[cfg(target_pointer_width = "32")]
    let mut lsize: usize = 0;

    for spec in local_metadata_spec_vec {
        #[cfg(target_pointer_width = "64")]
        {
            let res = try_mmap_contiguous_metadata_space(start, size, spec, true);
            if res.is_err() {
                return res;
            }
        }
        #[cfg(target_pointer_width = "32")]
        {
            lsize += meta_bytes_per_chunk(spec.log_min_obj_size, spec.log_num_of_bits);
        }
    }

    #[cfg(target_pointer_width = "32")]
    if lsize > 0 {
        let max = BYTES_IN_CHUNK >> LOG_LOCAL_SIDE_METADATA_WORST_CASE_RATIO;
        debug_assert!(
            lsize <= max,
            "local side metadata per chunk (0x{:x}) must be less than (0x{:x})",
            lsize,
            max
        );
        return try_map_per_chunk_metadata_space(start, size, lsize, true);
    }

    Ok(())
}

/// Unmap the corresponding metadata space or panic.
///
/// Note-1: This function is only used for test and debug right now.
///
/// Note-2: This function uses munmap() which works at page granularity.
///     If the corresponding metadata space's size is not a multiple of page size,
///     the actual unmapped space will be bigger than what you specify.
pub fn ensure_unmap_metadata_space(
    start: Address,
    size: usize,
    global_metadata_spec_vec: &[SideMetadataSpec],
    local_metadata_spec_vec: &[SideMetadataSpec],
) {
    trace!("ensure_unmap_metadata_space({}, 0x{:x})", start, size);
    debug_assert!(start.is_aligned_to(BYTES_IN_PAGE));
    debug_assert!(size % BYTES_IN_PAGE == 0);

    for spec in global_metadata_spec_vec {
        ensure_munmap_contiguos_metadata_space(start, size, spec);
    }

    for spec in local_metadata_spec_vec {
        #[cfg(target_pointer_width = "64")]
        {
            ensure_munmap_contiguos_metadata_space(start, size, spec);
        }
        #[cfg(target_pointer_width = "32")]
        {
            ensure_munmap_chunked_metadata_space(start, size, spec);
        }
    }
}

// Used only for debugging
// Panics in the required metadata for data_addr is not mapped
pub fn ensure_metadata_is_mapped(metadata_spec: SideMetadataSpec, data_addr: Address) {
    let meta_start = address_to_meta_address(metadata_spec, data_addr).align_down(BYTES_IN_PAGE);

    debug!(
        "ensure_metadata_is_mapped({}).meta_start({})",
        data_addr, meta_start
    );

    assert!(memory::check_is_mmapped(meta_start, BYTES_IN_PAGE).is_ok())
}

#[inline(always)]
pub fn load_atomic(metadata_spec: SideMetadataSpec, data_addr: Address) -> usize {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log <= 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;
        let byte_val = unsafe { meta_addr.atomic_load::<AtomicU8>(Ordering::SeqCst) };

        ((byte_val & mask) as usize) >> lshift
    } else if bits_num_log == 4 {
        unsafe { meta_addr.atomic_load::<AtomicU16>(Ordering::SeqCst) as usize }
    } else if bits_num_log == 5 {
        unsafe { meta_addr.atomic_load::<AtomicU32>(Ordering::SeqCst) as usize }
    } else if bits_num_log == 6 {
        unsafe { meta_addr.atomic_load::<AtomicUsize>(Ordering::SeqCst) }
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

pub fn store_atomic(metadata_spec: SideMetadataSpec, data_addr: Address, metadata: usize) {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log < 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;

        let mut old_val = unsafe { meta_addr.load::<u8>() };
        let mut new_val = (old_val & !mask) | ((metadata as u8) << lshift);

        while unsafe {
            meta_addr
                .compare_exchange::<AtomicU8>(old_val, new_val, Ordering::SeqCst, Ordering::SeqCst)
                .is_err()
        } {
            old_val = unsafe { meta_addr.load::<u8>() };
            new_val = (old_val & !mask) | ((metadata as u8) << lshift);
        }
    } else if bits_num_log == 3 {
        unsafe { meta_addr.atomic_store::<AtomicU8>(metadata as u8, Ordering::SeqCst) };
    } else if bits_num_log == 4 {
        unsafe { meta_addr.atomic_store::<AtomicU16>(metadata as u16, Ordering::SeqCst) };
    } else if bits_num_log == 5 {
        unsafe { meta_addr.atomic_store::<AtomicU32>(metadata as u32, Ordering::SeqCst) };
    } else if bits_num_log == 6 {
        unsafe { meta_addr.atomic_store::<AtomicUsize>(metadata as usize, Ordering::SeqCst) }
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

pub fn compare_exchange_atomic(
    metadata_spec: SideMetadataSpec,
    data_addr: Address,
    old_metadata: usize,
    new_metadata: usize,
) -> bool {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log < 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;

        let real_old_byte = unsafe { meta_addr.atomic_load::<AtomicU8>(Ordering::SeqCst) };
        let expected_old_byte = (real_old_byte & !mask) | ((old_metadata as u8) << lshift);
        let expected_new_byte = (expected_old_byte & !mask) | ((new_metadata as u8) << lshift);

        unsafe {
            meta_addr
                .compare_exchange::<AtomicU8>(
                    expected_old_byte,
                    expected_new_byte,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                )
                .is_ok()
        }
    } else if bits_num_log == 3 {
        unsafe {
            meta_addr
                .compare_exchange::<AtomicU8>(
                    old_metadata as u8,
                    new_metadata as u8,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                )
                .is_ok()
        }
    } else if bits_num_log == 4 {
        unsafe {
            meta_addr
                .compare_exchange::<AtomicU16>(
                    old_metadata as u16,
                    new_metadata as u16,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                )
                .is_ok()
        }
    } else if bits_num_log == 5 {
        unsafe {
            meta_addr
                .compare_exchange::<AtomicU32>(
                    old_metadata as u32,
                    new_metadata as u32,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                )
                .is_ok()
        }
    } else if bits_num_log == 6 {
        unsafe {
            meta_addr
                .compare_exchange::<AtomicUsize>(
                    old_metadata,
                    new_metadata,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                )
                .is_ok()
        }
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

// same as Rust atomics, this wraps around on overflow
pub fn fetch_add_atomic(metadata_spec: SideMetadataSpec, data_addr: Address, val: usize) -> usize {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log < 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;

        let mut old_val = unsafe { meta_addr.load::<u8>() };
        let mut new_sub_val = (((old_val & mask) >> lshift) + (val as u8)) & (mask >> lshift);
        let mut new_val = (old_val & !mask) | (new_sub_val << lshift);

        while unsafe {
            meta_addr
                .compare_exchange::<AtomicU8>(old_val, new_val, Ordering::SeqCst, Ordering::SeqCst)
                .is_err()
        } {
            old_val = unsafe { meta_addr.load::<u8>() };
            new_sub_val = (((old_val & mask) >> lshift) + (val as u8)) & (mask >> lshift);
            new_val = (old_val & !mask) | (new_sub_val << lshift);
        }

        (old_val & mask) as usize
    } else if bits_num_log == 3 {
        unsafe {
            (&*meta_addr.to_ptr::<AtomicU8>()).fetch_add(val as u8, Ordering::SeqCst) as usize
        }
    } else if bits_num_log == 4 {
        unsafe {
            (&*meta_addr.to_ptr::<AtomicU16>()).fetch_add(val as u16, Ordering::SeqCst) as usize
        }
    } else if bits_num_log == 5 {
        unsafe {
            (&*meta_addr.to_ptr::<AtomicU32>()).fetch_add(val as u32, Ordering::SeqCst) as usize
        }
    } else if bits_num_log == 6 {
        unsafe { (&*meta_addr.to_ptr::<AtomicUsize>()).fetch_add(val, Ordering::SeqCst) }
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

// same as Rust atomics, this wraps around on overflow
pub fn fetch_sub_atomic(metadata_spec: SideMetadataSpec, data_addr: Address, val: usize) -> usize {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log < 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;

        let mut old_val = unsafe { meta_addr.load::<u8>() };
        let mut new_sub_val = (((old_val & mask) >> lshift) - (val as u8)) & (mask >> lshift);
        let mut new_val = (old_val & !mask) | (new_sub_val << lshift);

        while unsafe {
            meta_addr
                .compare_exchange::<AtomicU8>(old_val, new_val, Ordering::SeqCst, Ordering::SeqCst)
                .is_err()
        } {
            old_val = unsafe { meta_addr.load::<u8>() };
            new_sub_val = (((old_val & mask) >> lshift) - (val as u8)) & (mask >> lshift);
            new_val = (old_val & !mask) | (new_sub_val << lshift);
        }

        (old_val & mask) as usize
    } else if bits_num_log == 3 {
        unsafe {
            (&*meta_addr.to_ptr::<AtomicU8>()).fetch_sub(val as u8, Ordering::SeqCst) as usize
        }
    } else if bits_num_log == 4 {
        unsafe {
            (&*meta_addr.to_ptr::<AtomicU16>()).fetch_sub(val as u16, Ordering::SeqCst) as usize
        }
    } else if bits_num_log == 5 {
        unsafe {
            (&*meta_addr.to_ptr::<AtomicU32>()).fetch_sub(val as u32, Ordering::SeqCst) as usize
        }
    } else if bits_num_log == 6 {
        unsafe { (&*meta_addr.to_ptr::<AtomicUsize>()).fetch_sub(val, Ordering::SeqCst) }
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

/// Non-atomic load of metadata.
///
/// # Safety
///
/// This is unsafe because:
///
/// 1. Concurrent access to this operation is undefined behaviour.
/// 2. Interleaving Non-atomic and atomic operations is undefined behaviour.
///
pub unsafe fn load(metadata_spec: SideMetadataSpec, data_addr: Address) -> usize {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log <= 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;
        let byte_val = meta_addr.load::<u8>();

        ((byte_val & mask) as usize) >> lshift
    } else if bits_num_log == 4 {
        meta_addr.load::<u16>() as usize
    } else if bits_num_log == 5 {
        meta_addr.load::<u32>() as usize
    } else if bits_num_log == 6 {
        meta_addr.load::<usize>() as usize
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

/// Non-atomic store of metadata.
///
/// # Safety
///
/// This is unsafe because:
///
/// 1. Concurrent access to this operation is undefined behaviour.
/// 2. Interleaving Non-atomic and atomic operations is undefined behaviour.
///
pub unsafe fn store(metadata_spec: SideMetadataSpec, data_addr: Address, metadata: usize) {
    let meta_addr = address_to_meta_address(metadata_spec, data_addr);
    if cfg!(debug_assertions) {
        ensure_metadata_is_mapped(metadata_spec, data_addr);
    }

    let bits_num_log = metadata_spec.log_num_of_bits;

    if bits_num_log < 3 {
        let lshift = meta_byte_lshift(metadata_spec, data_addr);
        let mask = meta_byte_mask(metadata_spec) << lshift;

        let old_val = meta_addr.load::<u8>();
        let new_val = (old_val & !mask) | ((metadata as u8) << lshift);

        meta_addr.store::<u8>(new_val);
    } else if bits_num_log == 3 {
        meta_addr.store::<u8>(metadata as u8);
    } else if bits_num_log == 4 {
        meta_addr.store::<u16>(metadata as u16);
    } else if bits_num_log == 5 {
        meta_addr.store::<u32>(metadata as u32);
    } else if bits_num_log == 6 {
        meta_addr.store::<usize>(metadata as usize);
    } else {
        unreachable!(
            "side metadata > {}-bits is not supported!",
            constants::BITS_IN_WORD
        );
    }
}

/// Bulk-zero a specific metadata for a chunk.
///
/// # Arguments
///
/// * `metadata_spec` - The specification of the target side metadata.
///
/// * `chunk_start` - The starting address of the chunk whose metadata is being zeroed.
///
pub fn bzero_metadata(metadata_spec: SideMetadataSpec, start: Address, size: usize) {
    debug_assert!(
        start.is_aligned_to(BYTES_IN_PAGE) && meta_byte_lshift(metadata_spec, start) == 0
    );

    let meta_start = address_to_meta_address(metadata_spec, start);
    if cfg!(target_pointer_width = "64") || metadata_spec.scope.is_global() {
        memory::zero(
            meta_start,
            address_to_meta_address(metadata_spec, start + size) - meta_start,
        );
    }
    #[cfg(target_pointer_width = "32")]
    if !metadata_spec.scope.is_global() {
        // per chunk policy-specific metadata for 32-bits targets
        let chunk_num = ((start + size).align_down(BYTES_IN_CHUNK)
            - start.align_down(BYTES_IN_CHUNK))
            / BYTES_IN_CHUNK;
        if chunk_num == 0 {
            memory::zero(
                meta_start,
                address_to_meta_address(metadata_spec, start + size) - meta_start,
            );
        } else {
            let second_data_chunk = start.align_up(BYTES_IN_CHUNK);
            // bzero the first sub-chunk
            memory::zero(
                meta_start,
                address_to_meta_address(metadata_spec, second_data_chunk) - meta_start,
            );
            let last_data_chunk = (start + size).align_down(BYTES_IN_CHUNK);
            let last_meta_chunk = address_to_meta_address(metadata_spec, last_data_chunk);
            // bzero the last sub-chunk
            memory::zero(
                last_meta_chunk,
                address_to_meta_address(metadata_spec, start + size) - last_meta_chunk,
            );
            let mut next_data_chunk = second_data_chunk;
            // bzero all chunks in the middle
            while next_data_chunk != last_data_chunk {
                memory::zero(
                    address_to_meta_address(metadata_spec, next_data_chunk),
                    meta_bytes_per_chunk(
                        metadata_spec.log_min_obj_size,
                        metadata_spec.log_num_of_bits,
                    ),
                );
                next_data_chunk += BYTES_IN_CHUNK;
            }
        }
    }
}