Swarm (#36)

* make azalea-pathfinder dir

* start writing d* lite impl

* more work on d* lite

* work more on implementing d* lite

* full d* lite impl

* updated edges

* add next() function

* add NoPathError

* why does dstar lite not work

* fix d* lite implementation

* make the test actually check the coords

* replace while loop with if statement

* fix clippy complaints

* make W only have to be PartialOrd

* fix PartialOrd issues

* implement mtd* lite

* add a test to mtd* lite

* remove normal d* lite

* make heuristic only take in one arg

* add `success` function

* Update README.md

* evil black magic to make .entity not need dimension

* start adding moves

* slightly improve the vec3/position situation

new macro that implements all the useful functions

* moves stuff

* make it compile

* update deps in az-pathfinder

* make it compile again

* more pathfinding stuff

* add Bot::look_at

* replace EntityMut and EntityRef with just Entity

* block pos pathfinding stuff

* rename movedirection to walkdirection

* execute path every tick

* advance path

* change az-pf version

* make azalea_client keep plugin state

* fix Plugins::get

* why does it think there is air

* start debugging incorrect air

* update some From methods to use rem_euclid

* start adding swarm

* fix deadlock

i still don't understand why it was happening but the solution was to keep the Client::player lock for shorter so it didn't overlap with the Client::dimension lock

* make lookat actually work probably

* fix going too fast

* Update main.rs

* make a thing immutable

* direction_looking_at

* fix rotations

* import swarm in an example

* fix stuff from merge

* remove azalea_pathfinder import

* delete azalea-pathfinder crate

already in azalea::pathfinder module

* swarms

* start working on shared dimensions

* Shared worlds work

* start adding Swarm::add_account

* add_account works

* change "client" to "bot" in some places

* Fix issues from merge

* Update world.rs

* add SwarmEvent::Disconnect(Account)

* almost add SwarmEvent::Chat and new plugin system

it panics rn

* make plugins have to provide the State associated type

* improve comments

* make fn build slightly cleaner

* fix SwarmEvent::Chat

* change a println in bot/main.rs

* Client::shutdown -> disconnect

* polish

fix clippy warnings + improve some docs a bit

* fix shared worlds*

*there's a bug that entities and bots will have their positions exaggerated because the relative movement packet is applied for every entity once per bot

* i am being trolled by rust

for some reason some stuff is really slow for literally no reason and it makes no sense i am going insane

* make world an RwLock again

* remove debug messages

* fix skipping event ticks

unfortunately now sending events is `.send().await?` instead of just `.send()`

* fix deadlock + warnings

* turns out my floor_mod impl was wrong

and i32::rem_euclid has the correct behavior LOL

* still errors with lots of bots

* make swarm iter & fix new chunks not loading

* improve docs

* start fixing tests

* fix all the tests

except the examples i don't know how to exclude them from the tests

* improve docs some more

1 files changed, 263 insertions, 59 deletions

diff --git a/azalea-world/src/entity_storage.rs b/azalea-world/src/entity_storage.rs
index 02d7d55a..c8c58a75 100755
--- a/azalea-world/src/entity_storage.rs
+++ b/azalea-world/src/entity_storage.rs
@@ -2,101 +2,229 @@ use crate::entity::EntityData;
 use azalea_core::ChunkPos;
 use log::warn;
 use nohash_hasher::{IntMap, IntSet};
-use std::collections::HashMap;
+use parking_lot::RwLock;
+use std::{
+    collections::HashMap,
+    sync::{Arc, Weak},
+};
 use uuid::Uuid;
 
-#[derive(Debug)]
-pub struct EntityStorage {
-    data_by_id: IntMap<u32, EntityData>,
-    id_by_chunk: HashMap<ChunkPos, IntSet<u32>>,
+// How entity updates are processed (to avoid issues with shared worlds)
+// - each bot contains a map of { entity id: updates received }
+// - the shared world also contains a canonical "true" updates received for each entity
+// - when a client loads an entity, its "updates received" is set to the same as the global "updates received"
+// - when the shared world sees an entity for the first time, the "updates received" is set to 1.
+// - clients can force the shared "updates received" to 0 to make it so certain entities (i.e. other bots in our swarm) don't get confused and updated by other bots
+// - when a client gets an update to an entity, we check if our "updates received" is the same as the shared world's "updates received":
+//      if it is, then process the update and increment the client's and shared world's "updates received"
+//      if not, then we simply increment our local "updates received" and do nothing else
+
+/// Store a map of entities by ID. To get an iterator over all entities, use
+/// `storage.shared.read().entities` [`WeakEntityStorage::entities`].
+///
+/// This is meant to be used with shared worlds.
+#[derive(Debug, Default)]
+pub struct PartialEntityStorage {
+    pub shared: Arc<RwLock<WeakEntityStorage>>,
+
+    /// The entity id of the player that owns this struct.
+    pub owner_entity_id: u32,
+    pub updates_received: IntMap<u32, u32>,
+    /// Strong references to the entities we have loaded.
+    data_by_id: IntMap<u32, Arc<EntityData>>,
+}
+
+/// Weakly store entities in a world. If the entities aren't being referenced
+/// by anything else (like an [`PartialEntityStorage`]), they'll be forgotten.
+#[derive(Debug, Default)]
+pub struct WeakEntityStorage {
+    data_by_id: IntMap<u32, Weak<EntityData>>,
+    /// An index of all the entity ids we know are in a chunk
+    ids_by_chunk: HashMap<ChunkPos, IntSet<u32>>,
+    /// An index of entity ids by their UUIDs
     id_by_uuid: HashMap<Uuid, u32>,
+
+    pub updates_received: IntMap<u32, u32>,
 }
 
-impl EntityStorage {
-    pub fn new() -> Self {
+impl PartialEntityStorage {
+    pub fn new(shared: Arc<RwLock<WeakEntityStorage>>, owner_entity_id: u32) -> Self {
+        shared.write().updates_received.insert(owner_entity_id, 0);
         Self {
+            shared,
+            owner_entity_id,
+            updates_received: IntMap::default(),
             data_by_id: IntMap::default(),
-            id_by_chunk: HashMap::default(),
-            id_by_uuid: HashMap::default(),
         }
     }
 
     /// Add an entity to the storage.
     #[inline]
     pub fn insert(&mut self, id: u32, entity: EntityData) {
-        self.id_by_chunk
+        // if the entity is already in the shared world, we don't need to do anything
+        if self.shared.read().data_by_id.contains_key(&id) {
+            return;
+        }
+
+        // add the entity to the "indexes"
+        let mut shared = self.shared.write();
+        shared
+            .ids_by_chunk
             .entry(ChunkPos::from(entity.pos()))
             .or_default()
             .insert(id);
-        self.id_by_uuid.insert(entity.uuid, id);
+        shared.id_by_uuid.insert(entity.uuid, id);
+
+        // now store the actual entity data
+        let entity = Arc::new(entity);
+        shared.data_by_id.insert(id, Arc::downgrade(&entity));
         self.data_by_id.insert(id, entity);
+        // set our updates_received to the shared updates_received, unless it's
+        // not there in which case set both to 1
+        if let Some(&shared_updates_received) = shared.updates_received.get(&id) {
+            // 0 means we're never tracking updates for this entity
+            if shared_updates_received != 0 || id == self.owner_entity_id {
+                self.updates_received.insert(id, 1);
+            }
+        } else {
+            shared.updates_received.insert(id, 1);
+            self.updates_received.insert(id, 1);
+        }
     }
 
-    /// Remove an entity from the storage by its id.
+    /// Remove an entity from this storage by its id. It will only be removed
+    /// from the shared storage if there are no other references to it.
     #[inline]
     pub fn remove_by_id(&mut self, id: u32) {
         if let Some(entity) = self.data_by_id.remove(&id) {
-            let entity_chunk = ChunkPos::from(entity.pos());
-            let entity_uuid = entity.uuid;
-            if self.id_by_chunk.remove(&entity_chunk).is_none() {
-                warn!("Tried to remove entity with id {id} from chunk {entity_chunk:?} but it was not found.");
-            }
-            if self.id_by_uuid.remove(&entity_uuid).is_none() {
-                warn!("Tried to remove entity with id {id} from uuid {entity_uuid:?} but it was not found.");
-            }
+            let chunk = ChunkPos::from(entity.pos());
+            let uuid = entity.uuid;
+            self.updates_received.remove(&id);
+            drop(entity);
+            // maybe remove it from the storage
+            self.shared.write().remove_entity_if_unused(id, uuid, chunk);
         } else {
             warn!("Tried to remove entity with id {id} but it was not found.")
         }
     }
 
-    /// Check if there is an entity that exists with the given id.
+    /// Whether the entity with the given id is being loaded by this storage.
+    /// If you want to check whether the entity is in the shared storage, use
+    /// [`WeakEntityStorage::contains_id`].
     #[inline]
-    pub fn contains_id(&self, id: &u32) -> bool {
+    pub fn limited_contains_id(&self, id: &u32) -> bool {
         self.data_by_id.contains_key(id)
     }
 
-    /// Get a reference to an entity by its id.
+    /// Whether the entity with the given id is in the shared storage (i.e.
+    /// it's possible we don't see the entity but something else in the shared
+    /// storage does). To check whether the entity is being loaded by this
+    /// storage, use [`PartialEntityStorage::limited_contains_id`].
     #[inline]
-    pub fn get_by_id(&self, id: u32) -> Option<&EntityData> {
+    pub fn contains_id(&self, id: &u32) -> bool {
+        self.shared.read().data_by_id.contains_key(id)
+    }
+
+    /// Get a reference to an entity by its id, if it's being loaded by this storage.
+    #[inline]
+    pub fn limited_get_by_id(&self, id: u32) -> Option<&Arc<EntityData>> {
         self.data_by_id.get(&id)
     }
 
-    /// Get a mutable reference to an entity by its id.
+    /// Get a mutable reference to an entity by its id, if it's being loaded by
+    /// this storage.
     #[inline]
-    pub fn get_mut_by_id(&mut self, id: u32) -> Option<&mut EntityData> {
+    pub fn limited_get_mut_by_id(&mut self, id: u32) -> Option<&mut Arc<EntityData>> {
         self.data_by_id.get_mut(&id)
     }
 
-    /// Get a reference to an entity by its uuid.
+    /// Returns whether we're allowed to update this entity (to prevent two clients in
+    /// a shared world updating it twice), and acknowleges that we WILL update
+    /// it if it's true. Don't call this unless you actually got an entity
+    /// update that all other clients within render distance will get too.
+    pub fn maybe_update(&mut self, id: u32) -> bool {
+        let this_client_updates_received = self.updates_received.get(&id).copied();
+        let shared_updates_received = self.shared.read().updates_received.get(&id).copied();
+
+        let can_update = this_client_updates_received == shared_updates_received;
+        if can_update {
+            let new_updates_received = this_client_updates_received.unwrap_or(0) + 1;
+            self.updates_received.insert(id, new_updates_received);
+            self.shared
+                .write()
+                .updates_received
+                .insert(id, new_updates_received);
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Get an entity in the shared storage by its id, if it exists.
     #[inline]
-    pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<&EntityData> {
-        self.id_by_uuid
+    pub fn get_by_id(&self, id: u32) -> Option<Arc<EntityData>> {
+        self.shared
+            .read()
+            .data_by_id
+            .get(&id)
+            .and_then(|e| e.upgrade())
+    }
+
+    /// Get a reference to an entity by its UUID, if it's being loaded by this
+    /// storage.
+    #[inline]
+    pub fn limited_get_by_uuid(&self, uuid: &Uuid) -> Option<&Arc<EntityData>> {
+        self.shared
+            .read()
+            .id_by_uuid
             .get(uuid)
             .and_then(|id| self.data_by_id.get(id))
     }
 
-    /// Get a mutable reference to an entity by its uuid.
+    /// Get a mutable reference to an entity by its UUID, if it's being loaded
+    /// by this storage.
     #[inline]
-    pub fn get_mut_by_uuid(&mut self, uuid: &Uuid) -> Option<&mut EntityData> {
-        self.id_by_uuid
+    pub fn limited_get_mut_by_uuid(&mut self, uuid: &Uuid) -> Option<&mut Arc<EntityData>> {
+        self.shared
+            .read()
+            .id_by_uuid
             .get(uuid)
             .and_then(|id| self.data_by_id.get_mut(id))
     }
 
-    /// Clear all entities in a chunk.
+    /// Get an entity in the shared storage by its UUID, if it exists.
+    #[inline]
+    pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<Arc<EntityData>> {
+        self.shared.read().id_by_uuid.get(uuid).and_then(|id| {
+            self.shared
+                .read()
+                .data_by_id
+                .get(id)
+                .and_then(|e| e.upgrade())
+        })
+    }
+
+    /// Clear all entities in a chunk. This will not clear them from the
+    /// shared storage, unless there are no other references to them.
     pub fn clear_chunk(&mut self, chunk: &ChunkPos) {
-        if let Some(entities) = self.id_by_chunk.remove(chunk) {
-            for entity_id in entities {
-                if let Some(entity) = self.data_by_id.remove(&entity_id) {
-                    self.id_by_uuid.remove(&entity.uuid);
-                } else {
-                    warn!("While clearing chunk {chunk:?}, found an entity that isn't in by_id {entity_id}.");
+        if let Some(entities) = self.shared.read().ids_by_chunk.get(chunk) {
+            for id in entities.iter() {
+                if let Some(entity) = self.data_by_id.remove(id) {
+                    let uuid = entity.uuid;
+                    drop(entity);
+                    // maybe remove it from the storage
+                    self.shared
+                        .write()
+                        .remove_entity_if_unused(*id, uuid, *chunk);
                 }
             }
+            //     for entity_id in entities {
+            //         self.remove_by_id(entity_id);
+            //     }
         }
     }
 
-    /// Updates an entity from its old chunk.
+    /// Move an entity from its old chunk to a new chunk.
     #[inline]
     pub fn update_entity_chunk(
         &mut self,
@@ -104,36 +232,40 @@ impl EntityStorage {
         old_chunk: &ChunkPos,
         new_chunk: &ChunkPos,
     ) {
-        if let Some(entities) = self.id_by_chunk.get_mut(old_chunk) {
+        if let Some(entities) = self.shared.write().ids_by_chunk.get_mut(old_chunk) {
             entities.remove(&entity_id);
         }
-        self.id_by_chunk
+        self.shared
+            .write()
+            .ids_by_chunk
             .entry(*new_chunk)
             .or_default()
             .insert(entity_id);
     }
 
-    /// Get an iterator over all entities.
-    #[inline]
-    pub fn entities(&self) -> std::collections::hash_map::Values<'_, u32, EntityData> {
-        self.data_by_id.values()
-    }
-
-    pub fn find_one_entity<F>(&self, mut f: F) -> Option<&EntityData>
+    pub fn find_one_entity<F>(&self, mut f: F) -> Option<Arc<EntityData>>
     where
-        F: FnMut(&EntityData) -> bool,
+        F: FnMut(&Arc<EntityData>) -> bool,
     {
-        self.entities().find(|&entity| f(entity))
+        for entity in self.shared.read().entities() {
+            if let Some(entity) = entity.upgrade() {
+                if f(&entity) {
+                    return Some(entity);
+                }
+            }
+        }
+        None
     }
 
-    pub fn find_one_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<&EntityData>
+    pub fn find_one_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<Arc<EntityData>>
     where
         F: FnMut(&EntityData) -> bool,
     {
-        if let Some(entities) = self.id_by_chunk.get(chunk) {
+        let shared = self.shared.read();
+        if let Some(entities) = shared.ids_by_chunk.get(chunk) {
             for entity_id in entities {
-                if let Some(entity) = self.data_by_id.get(entity_id) {
-                    if f(entity) {
+                if let Some(entity) = shared.data_by_id.get(entity_id).and_then(|e| e.upgrade()) {
+                    if f(&entity) {
                         return Some(entity);
                     }
                 }
@@ -143,9 +275,81 @@ impl EntityStorage {
     }
 }
 
-impl Default for EntityStorage {
-    fn default() -> Self {
-        Self::new()
+impl WeakEntityStorage {
+    pub fn new() -> Self {
+        Self {
+            data_by_id: IntMap::default(),
+            ids_by_chunk: HashMap::default(),
+            id_by_uuid: HashMap::default(),
+            updates_received: IntMap::default(),
+        }
+    }
+
+    /// Remove an entity from the storage if it has no strong references left.
+    /// Returns whether the entity was removed.
+    pub fn remove_entity_if_unused(&mut self, id: u32, uuid: Uuid, chunk: ChunkPos) -> bool {
+        if self.data_by_id.get(&id).and_then(|e| e.upgrade()).is_some() {
+            // if we could get the entity, that means there are still strong
+            // references to it
+            false
+        } else {
+            if self.ids_by_chunk.remove(&chunk).is_none() {
+                warn!("Tried to remove entity with id {id} from chunk {chunk:?} but it was not found.");
+            }
+            if self.id_by_uuid.remove(&uuid).is_none() {
+                warn!(
+                    "Tried to remove entity with id {id} from uuid {uuid:?} but it was not found."
+                );
+            }
+            if self.updates_received.remove(&id).is_none() {
+                // if this happens it means we weren't tracking the updates_received for the client (bad)
+                warn!(
+                    "Tried to remove entity with id {id} from updates_received but it was not found."
+                );
+            }
+            true
+        }
+    }
+
+    /// Remove a chunk from the storage if the entities in it have no strong
+    /// references left.
+    pub fn remove_chunk_if_unused(&mut self, chunk: &ChunkPos) {
+        if let Some(entities) = self.ids_by_chunk.get(chunk) {
+            if entities.is_empty() {
+                self.ids_by_chunk.remove(chunk);
+            }
+        }
+    }
+
+    /// Get an iterator over all entities in the shared storage. The iterator
+    /// is over `Weak<EntityData>`s, so you'll have to manually try to upgrade.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// let mut storage = EntityStorage::new();
+    ///     storage.insert(
+    ///     0,
+    ///     Arc::new(EntityData::new(
+    ///         uuid,
+    ///         Vec3::default(),
+    ///     EntityMetadata::Player(metadata::Player::default()),
+    /// )),
+    /// );
+    /// for entity in storage.shared.read().entities() {
+    ///     if let Some(entity) = entity.upgrade() {
+    ///         println!("Entity: {:?}", entity);
+    ///     }
+    /// }
+    /// ```
+    pub fn entities(&self) -> std::collections::hash_map::Values<'_, u32, Weak<EntityData>> {
+        self.data_by_id.values()
+    }
+
+    /// Whether the entity with the given id is in the shared storage.
+    #[inline]
+    pub fn contains_id(&self, id: &u32) -> bool {
+        self.data_by_id.contains_key(id)
     }
 }
 
@@ -158,7 +362,7 @@ mod tests {
 
     #[test]
     fn test_store_entity() {
-        let mut storage = EntityStorage::new();
+        let mut storage = PartialEntityStorage::default();
         assert!(storage.get_by_id(0).is_none());
 
         let uuid = Uuid::from_u128(100);