PkmnLib_rs/src/dynamic_data/flow/choice_queue.rs

use crate::dynamic_data::choices::TurnChoice;
use crate::dynamic_data::script_handling::ScriptSource;
use crate::dynamic_data::Pokemon;
use crate::{script_hook, PkmnError, VecExt};
use anyhow::Result;
use anyhow_ext::anyhow;
use parking_lot::lock_api::MappedRwLockReadGuard;
use parking_lot::{RawRwLock, RwLock, RwLockReadGuard};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

/// The ChoiceQueue is used to run choices one by one.
///
/// It functions internally by holding a vector of choices, and passing ownership of the turn choice
/// to the turn executor one by one, replacing it with empty spots at the start. It holds several
/// helper functions to change the turn order while doing the execution. This is needed, as several
/// moves in Pokemon actively mess with this order.
#[derive(Debug)]
pub struct ChoiceQueue {
    /// Our storage of turn choices. Starts out completely filled, then slowly empties as turns get
    /// executed.
    queue: RwLock<Vec<Option<Arc<TurnChoice>>>>,
    /// The current index of the turn we need to execute next.
    current: AtomicUsize,
}

impl ChoiceQueue {
    /// Initializes a ChoiceQueue, and sort the choices.
    pub(crate) fn new(mut queue: Vec<Option<Arc<TurnChoice>>>) -> Self {
        queue.sort_unstable_by(|a, b| b.cmp(a));
        Self {
            queue: RwLock::new(queue),
            current: AtomicUsize::new(0),
        }
    }

    /// Dequeues the next turn choice to be executed. This gives ownership to the callee, and replaces
    /// our own reference to the turn choice with an empty spot. It also increments the current position
    /// by one.
    pub fn dequeue(&self) -> Result<Option<Arc<TurnChoice>>> {
        let mut write_lock = self.queue.write();
        if self.current.load(Ordering::Relaxed) >= write_lock.len() {
            return Ok(None);
        }
        let c = write_lock
            .get_mut(self.current.load(Ordering::Relaxed))
            .ok_or(anyhow!("Unable to get current turn choice"))?
            .take();
        self.current.fetch_add(1, Ordering::Relaxed);
        Ok(c)
    }

    /// This reads what the next choice to execute will be, without modifying state.
    pub fn peek(&self) -> Result<Option<MappedRwLockReadGuard<'_, RawRwLock, Arc<TurnChoice>>>> {
        let read_lock = self.queue.read();
        if self.current.load(Ordering::Relaxed) >= read_lock.len() {
            Ok(None)
        } else {
            let v = RwLockReadGuard::try_map(read_lock, |a| match a.get(self.current.load(Ordering::Relaxed)) {
                Some(Some(v)) => Some(v),
                _ => None,
            });
            match v {
                Ok(v) => Ok(Some(v)),
                Err(_) => Err(anyhow!("Could not map choice")),
            }
        }
    }

    /// Check if we have any choices remaining.
    pub fn has_next(&self) -> bool {
        self.current.load(Ordering::Relaxed) < self.queue.read().len()
    }

    /// This resorts the yet to be executed choices. This can be useful for dealing with situations
    /// such as Pokemon changing forms just after the very start of a turn, when turn order has
    /// technically already been decided.
    pub fn resort(&mut self) -> Result<()> {
        let len = self.queue.read().len();
        let mut write_lock = self.queue.write();
        for index in self.current.load(Ordering::Relaxed)..len {
            let choice = &mut write_lock.get_mut_res(index)?;
            if let Some(choice) = choice {
                let mut speed = choice.user().boosted_stats().speed();
                script_hook!(change_speed, (*choice), choice, &mut speed);
                choice.set_speed(speed)
            }
        }

        write_lock
            .get_mut(self.current.load(Ordering::Relaxed)..len)
            .ok_or(PkmnError::IndexOutOfBounds {
                index: self.current.load(Ordering::Relaxed),
                len,
            })?
            .sort_unstable_by(|a, b| b.cmp(a));
        Ok(())
    }

    /// This moves the choice of a specific Pokemon up to the next choice to be executed.
    pub fn move_pokemon_choice_next(&self, pokemon: &Pokemon) -> Result<bool> {
        let mut queue_lock = self.queue.write();
        let mut desired_index = None;
        // Find the index for the choice we want to move up.
        for index in self.current.load(Ordering::Relaxed)..queue_lock.len() {
            if let Some(Some(choice)) = &queue_lock.get(index) {
                if pokemon.eq(choice.user()) {
                    desired_index = Some(index);
                    break;
                }
            }
        }
        let result = match desired_index {
            Some(desired_index) => {
                // If the choice we want to move up is already the next choice, just return.
                if desired_index == self.current.load(Ordering::Relaxed) {
                    return Ok(true);
                }

                // Take the choice we want to move forward out of it's place.
                let choice = queue_lock
                    .get_mut_res(desired_index)?
                    .take()
                    .ok_or(anyhow!("Choice was already taken"))?;
                // Iterate backwards from the spot before the choice we want to move up, push them all back
                // by 1 spot.
                let current = self.current.load(Ordering::Relaxed);
                for index in (current..desired_index).rev() {
                    queue_lock.swap(index, index + 1);
                }
                // Place the choice that needs to be next in the next to be executed position.
                let _ = queue_lock
                    .get_mut_res(self.current.load(Ordering::Relaxed))?
                    .insert(choice);
                true
            }
            None => false,
        };
        Ok(result)
    }

    /// Internal helper function to be easily able to iterate over the yet to be executed choices.
    pub(crate) fn get_queue(&self) -> Result<MappedRwLockReadGuard<'_, RawRwLock, [Option<Arc<TurnChoice>>]>> {
        let read_lock = self.queue.read();
        match RwLockReadGuard::try_map(read_lock, |a| {
            a.get(self.current.load(Ordering::Relaxed)..self.queue.read().len())
        }) {
            Ok(v) => Ok(v),
            Err(_) => Err(PkmnError::IndexOutOfBounds {
                index: self.current.load(Ordering::Relaxed),
                len: self.queue.read().len(),
            }
            .into()),
        }
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
#[allow(clippy::indexing_slicing)]
mod tests {
    use super::*;
    use crate::defines::LevelInt;
    use crate::dynamic_data::{DynamicLibrary, PassChoice};
    use crate::static_data::{AbilityIndex, Gender};
    use std::sync::Arc;
    #[test]
    fn create_empty_queue() {
        let queue = ChoiceQueue::new(Vec::new());
        assert!(!queue.has_next());
        assert!(queue.peek().unwrap().is_none());
    }

    #[test]
    fn dequeue_from_empty_queue() {
        let queue = ChoiceQueue::new(Vec::new());
        assert!(queue.dequeue().unwrap().is_none());
    }

    fn get_user(level: LevelInt) -> Pokemon {
        let lib = Arc::new(crate::dynamic_data::libraries::dynamic_library::test::build());
        let species = lib.static_data().species().get(&"foo".into()).unwrap();
        let form = species.get_form(&"default".into()).unwrap();

        Pokemon::new(
            lib,
            species,
            &form,
            AbilityIndex {
                hidden: false,
                index: 0,
            },
            level,
            0,
            Gender::Male,
            0,
            &"test_nature".into(),
        )
        .unwrap()
    }

    #[test]
    fn create_queue_with_single_item() {
        let user = get_user(10);

        let queue = ChoiceQueue::new(vec![Some(Arc::new(TurnChoice::Pass(PassChoice::new(user))))]);
        assert!(queue.has_next());
        assert!(queue.peek().unwrap().is_some());
        assert_eq!(7, queue.peek().unwrap().unwrap().speed());
    }

    #[test]
    fn dequeue_from_queue_with_single_item() {
        let user = get_user(10);

        let queue = ChoiceQueue::new(vec![Some(Arc::new(TurnChoice::Pass(PassChoice::new(user))))]);
        assert!(queue.has_next());
        assert_eq!(7, queue.dequeue().unwrap().unwrap().speed());
        assert!(!queue.has_next());
        assert!(queue.peek().unwrap().is_none());
    }

    #[test]
    fn create_queue_with_two_items_with_equal_order() {
        let user1 = get_user(10);
        let user2 = get_user(10);

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1)))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2)))),
        ]);
        assert!(queue.has_next());
        assert!(queue.peek().unwrap().is_some());
        assert_eq!(7, queue.peek().unwrap().unwrap().speed());
    }

    #[test]
    fn create_queue_with_two_items_get_queue() {
        let user1 = get_user(10);
        let user2 = get_user(5);

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1.clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2)))),
        ]);
        let inner_queue = queue.get_queue().unwrap();
        assert_eq!(inner_queue[0].as_ref().unwrap().user().clone(), user1);
    }

    #[test]
    fn create_queue_with_two_items_in_wrong_order_sorts_correctly() {
        let user1 = get_user(5);
        let user2 = get_user(100);

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1)))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2)))),
        ]);
        assert_eq!(25, queue.dequeue().unwrap().unwrap().speed());
        assert_eq!(6, queue.dequeue().unwrap().unwrap().speed());
    }

    #[test]
    fn resort_with_two_choices() {
        let user1 = get_user(50);
        let user2 = get_user(1);

        let mut queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1.clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2.clone())))),
        ]);

        user2.change_level_by(60).unwrap();
        assert_eq!(user1, queue.peek().unwrap().unwrap().user().clone(),);
        queue.resort().unwrap();
        assert_eq!(user2, queue.peek().unwrap().unwrap().user().clone(),);
    }

    #[test]
    fn move_pokemon_choice_first_with_two_choices() {
        let user1 = get_user(100);
        let user2 = get_user(1);

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1.clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2.clone())))),
        ]);
        assert_eq!(user1, queue.peek().unwrap().unwrap().user().clone(),);
        assert!(queue.move_pokemon_choice_next(&user2).unwrap());

        assert_eq!(user2, queue.dequeue().unwrap().unwrap().user().clone(),);
        assert_eq!(user1, queue.dequeue().unwrap().unwrap().user().clone(),);
    }

    #[test]
    fn move_pokemon_choice_first_when_choice_has_already_been() {
        let user1 = get_user(10);
        let user2 = get_user(100);

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1.clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2.clone())))),
        ]);
        assert_eq!(user2, queue.dequeue().unwrap().unwrap().user().clone(),);
        assert!(!queue.move_pokemon_choice_next(&user2).unwrap());

        assert_eq!(user1, queue.dequeue().unwrap().unwrap().user().clone(),);
        assert!(queue.peek().unwrap().is_none())
    }

    #[test]
    fn move_pokemon_choice_first_when_choice_is_next() {
        let user1 = get_user(100);
        let user2 = get_user(10);

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user1.clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(user2)))),
        ]);
        assert_eq!(user1, queue.peek().unwrap().unwrap().user().clone(),);
        assert!(queue.move_pokemon_choice_next(&user1).unwrap());
        assert_eq!(user1, queue.peek().unwrap().unwrap().user().clone(),);
    }

    #[test]
    fn move_pokemon_choice_first_with_seven_choices() {
        let users = [
            get_user(100),
            get_user(90),
            get_user(80),
            get_user(70),
            get_user(60),
            get_user(50),
            get_user(40),
        ];

        let queue = ChoiceQueue::new(vec![
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[0].clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[1].clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[2].clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[3].clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[4].clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[5].clone())))),
            Some(Arc::new(TurnChoice::Pass(PassChoice::new(users[6].clone())))),
        ]);
        assert_eq!(users[0], queue.peek().unwrap().unwrap().user().clone(),);
        assert!(queue.move_pokemon_choice_next(&users[4]).unwrap());

        assert_eq!(users[4], queue.dequeue().unwrap().unwrap().user().clone(),);
        for index in 0..4 {
            assert_eq!(users[index], queue.dequeue().unwrap().unwrap().user().clone(),);
        }
        for index in 5..7 {
            assert_eq!(users[index], queue.dequeue().unwrap().unwrap().user().clone(),);
        }
    }
}