pumpkin-core 0.5.0

use std::ops::Range;

use crate::basic_types::PredicateId;
use crate::containers::HashMap;
use crate::containers::StorageKey;
use crate::propagators::nogoods::NogoodId;

/// An arena allocator for storing nogoods.
///
/// The idea is to avoid double indirection by storing one large structure with [`PredicateId`]s.
///
/// Currently, deleting nogoods is not supported.
#[derive(Clone, Default, Debug)]
pub(crate) struct ArenaAllocator {
    /// A list of [`PredicateId`]s representing the nogoods.
    ///
    /// If there is a [`NogoodId`] with value `i`, then the [`PredicateId`] at position `i` will
    /// contain the length `x` of the nogood and the [`PredicateId`] at position `i + 1` will
    /// contain the last-traversed watcher index. The next `i + 2 + x` elements are then the nogood
    /// pointed to by the [`NogoodId`] with value `i`.
    pub(crate) nogoods: Vec<PredicateId>,
    /// Maps each [`NogoodId`] to an index; this is to prevent unnecessary allocations for other
    /// structures such as the [`NogoodInfo`] which use direct hashing for storing information
    /// about nogoods.
    pub(crate) nogood_id_to_index: HashMap<NogoodId, NogoodIndex>,
    /// The current index for the next [`NogoodId`] which is entered; see
    /// [`ArenaAllocator::nogood_id_to_index`].
    current_index: u32,
    /// The number of elements (i.e., [`PredicateId`]s), that are created when the arena is
    /// initialised.
    ///
    /// Note that it is lazily initialised, so that this memory is only allocated the first time
    /// that a nogood is added to the arena.
    initial_capacity: usize,
}

/// The index offset which determines how many elements to skip before the actual nogood predicates
/// begin.
///
/// See [`ArenaAllocator::nogoods`] for more information.
const OFFSET: usize = 2;

#[derive(Clone, Copy, Debug, Hash)]
pub(crate) struct NogoodIndex(u32);

impl StorageKey for NogoodIndex {
    fn index(&self) -> usize {
        self.0 as usize
    }

    fn create_from_index(index: usize) -> Self {
        NogoodIndex(index as u32)
    }
}

impl ArenaAllocator {
    pub(crate) fn new(capacity: usize) -> Self {
        Self {
            nogoods: Vec::default(),
            nogood_id_to_index: HashMap::default(),
            current_index: 0,
            initial_capacity: capacity,
        }
    }

    /// Inserts the nogood consisting of [`PredicateId`]s and returns its corresponding
    /// [`NogoodId`].
    pub(crate) fn insert(&mut self, nogood: Vec<PredicateId>) -> NogoodId {
        if self.nogoods.is_empty() {
            self.nogoods.reserve_exact(self.initial_capacity);
        }

        let nogood_id = NogoodId::create_from_index(self.nogoods.len());

        // We store the NogoodId with its index.
        let _ = self
            .nogood_id_to_index
            .insert(nogood_id, NogoodIndex(self.current_index));
        self.current_index += 1;

        // We push a PredicateId which stores the length of the nogood
        self.nogoods
            .push(PredicateId::create_from_index(nogood.len()));
        // We also push a PredicateId which stores the last-traversed watcher (defaults to the
        // first non-watcher element)
        self.nogoods.push(PredicateId::create_from_index(2));
        self.nogoods.extend(nogood);

        nogood_id
    }

    /// Returns the index of the provided [`NogoodId`].
    ///
    /// In other words, if the nogood with ID [`NogoodId`] was the `n`th nogood to be inserted then
    /// this method will return `n`.
    pub(crate) fn get_nogood_index(&self, nogood_id: &NogoodId) -> NogoodIndex {
        *self
            .nogood_id_to_index
            .get(nogood_id)
            .expect("Expected nogood predicate to exist")
    }

    /// Returns a list of all the present [`NogoodId`]s.
    pub(crate) fn nogoods_ids(&self) -> impl Iterator<Item = NogoodId> + '_ {
        NogoodIdIterator {
            nogoods: &self.nogoods,
            current_index: 0,
        }
    }

    /// Returns the length of the nogood corresponding to the provided [`NogoodId`].
    fn len_of_nogood(&self, nogood_id: NogoodId) -> usize {
        self.nogoods[nogood_id.index()].index()
    }

    /// Calculates the range of the nogood spanned by the nogood with ID [`NogoodId`].
    ///
    /// Does not include any of the information [`PredicateId`]s.
    fn calculate_range_of_nogood(&self, nogood_id: NogoodId) -> Range<usize> {
        let len = self.len_of_nogood(nogood_id);
        nogood_id.index() + OFFSET..nogood_id.index() + OFFSET + len
    }

    /// Calculates the range of the nogood spanned by the nogood with ID [`NogoodId`].
    ///
    /// Includes the [`PredicateId`] storing the last-traversed watcher index as the first element.
    #[allow(unused, reason = "Currently inlined due to borrow issues")]
    pub(crate) fn calculate_range_of_nogood_including_last_traversed(
        &self,
        nogood_id: NogoodId,
    ) -> Range<usize> {
        let len = self.len_of_nogood(nogood_id);
        nogood_id.index() + OFFSET - 1..nogood_id.index() + OFFSET + len
    }

    /// Returns the nogood pointed to by [`NogoodId`].
    pub(crate) fn get_nogood(&self, nogood_id: NogoodId) -> &[PredicateId] {
        let nogood_range = self.calculate_range_of_nogood(nogood_id);

        &self.nogoods[nogood_range]
    }

    /// Returns a mutable reference to the nogood pointed to by [`NogoodId`].
    #[allow(unused, reason = "Standard API")]
    pub(crate) fn get_nogood_mut(&mut self, nogood_id: NogoodId) -> &mut [PredicateId] {
        let nogood_range = self.calculate_range_of_nogood(nogood_id);

        &mut self.nogoods[nogood_range]
    }

    /// Returns a tuple consisting of a mutable reference to the index of the last-traversed watcher
    /// and a mutable reference to the nogood pointed to by [`NogoodId`].
    #[allow(unused, reason = "Currently inlined due to borrow issues")]
    pub(crate) fn get_nogood_mut_with_last_traversed(
        &mut self,
        nogood_id: NogoodId,
    ) -> (&mut u32, &mut [PredicateId]) {
        let nogood_range = self.calculate_range_of_nogood_including_last_traversed(nogood_id);

        self.nogoods[nogood_range]
            .split_first_mut()
            .map(|(last_traversed, nogood)| (&mut last_traversed.id, nogood))
            .expect("Expected nogood to be at least of length two")
    }
}

pub(crate) struct NogoodIdIterator<'a> {
    nogoods: &'a Vec<PredicateId>,
    current_index: usize,
}

impl Iterator for NogoodIdIterator<'_> {
    type Item = NogoodId;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_index >= self.nogoods.len() {
            return None;
        }
        let id = NogoodId::create_from_index(self.current_index);
        self.current_index += self.nogoods[self.current_index].id as usize + OFFSET;

        Some(id)
    }
}