space-crush/space-crush-common/src/systems/fleet_owned_update.rs

use std::collections::HashMap;

use shrev::ReaderId;
use specs::{prelude::*, world::Index, Entities, ReadStorage, System, World, WriteStorage};

use crate::{
    components::{Fleet, FleetOwned, Ship},
    misc::ComponentEvent,
};

pub struct FleetOwnedUpdate {
    fleet_ids: BitSet,
    fleet_owned_ids: BitSet,
    old_fleet_ids: HashMap<Index, Entity>,
    fleet_owned_event_id: ReaderId<ComponentEvent<FleetOwned>>,
}

impl FleetOwnedUpdate {
    pub fn new(world: &mut World) -> Self {
        let fleet_ids = BitSet::new();
        let fleet_owned_ids = BitSet::new();
        let old_fleet_ids = HashMap::new();

        let fleet_owned_event_id = unsafe {
            WriteStorage::<FleetOwned>::setup(world);

            let mut fleet_owned = world.system_data::<WriteStorage<FleetOwned>>();
            fleet_owned
                .unprotected_storage_mut()
                .channel_mut()
                .register_reader()
        };

        Self {
            fleet_ids,
            fleet_owned_ids,
            old_fleet_ids,
            fleet_owned_event_id,
        }
    }
}

#[derive(SystemData)]
pub struct FleetOwnedUpdateData<'a> {
    entities: Entities<'a>,
    ships: ReadStorage<'a, Ship>,
    fleets: WriteStorage<'a, Fleet>,
    fleet_owned: ReadStorage<'a, FleetOwned>,
}

impl<'a> System<'a> for FleetOwnedUpdate {
    type SystemData = FleetOwnedUpdateData<'a>;

    fn run(&mut self, data: Self::SystemData) {
        let FleetOwnedUpdateData {
            entities,
            ships,
            mut fleets,
            fleet_owned,
        } = data;

        self.fleet_ids.clear();
        self.fleet_owned_ids.clear();
        self.old_fleet_ids.clear();

        /* handle events */
        let events = fleet_owned
            .unprotected_storage()
            .channel()
            .read(&mut self.fleet_owned_event_id);
        for event in events {
            match event {
                ComponentEvent::Inserted(id, fleet_owned) => {
                    self.fleet_ids.add(fleet_owned.owner().id());
                    self.fleet_owned_ids.add(*id);
                }
                ComponentEvent::Modified(id, fleet_owned) => {
                    self.fleet_ids.add(fleet_owned.owner().id());
                    self.fleet_owned_ids.add(*id);
                    *self
                        .old_fleet_ids
                        .entry(*id)
                        .or_insert_with(|| fleet_owned.owner()) = fleet_owned.owner();
                }
                ComponentEvent::Removed(id, fleet_owned) => {
                    self.fleet_ids.add(fleet_owned.owner().id());
                    self.fleet_owned_ids.add(*id);
                    *self
                        .old_fleet_ids
                        .entry(*id)
                        .or_insert_with(|| fleet_owned.owner()) = fleet_owned.owner();
                }
            }
        }

        /* find new fleet ids */
        for (fleet_owned, _) in (&fleet_owned, &self.fleet_owned_ids).join() {
            self.fleet_ids.add(fleet_owned.owner().id());
        }

        /* update fleets */
        for (fleet_id, fleet, _) in (&entities, &mut fleets, &self.fleet_ids).join() {
            let data = (&entities, &ships, &fleet_owned, &self.fleet_owned_ids);

            for (ship_id, ship, fleet_owned, _) in data.join() {
                let new_match = fleet_id == fleet_owned.owner();
                let old_match = match self.old_fleet_ids.get(&ship_id.id()) {
                    Some(old_fleet_id) => fleet_id == *old_fleet_id,
                    None => false,
                };

                if old_match && !new_match {
                    let count = fleet.count_mut();
                    let count = &mut count[ship.type_()];
                    *count = count.saturating_sub(1);

                    fleet.owned_mut().remove(ship_id.id());
                } else if !old_match && new_match {
                    let count = fleet.count_mut();
                    let count = &mut count[ship.type_()];
                    *count += 1;

                    fleet.owned_mut().add(ship_id.id());
                }
            }
        }
    }
}