aboutsummaryrefslogtreecommitdiff
path: root/azalea-world/src/entity_storage.rs
blob: 4dd3ec12aafcb2f68b64abe3d2038e8eb1d6e174 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
use crate::entity::EntityData;
use azalea_core::ChunkPos;
use log::warn;
use nohash_hasher::{IntMap, IntSet};
use std::collections::HashMap;
use uuid::Uuid;

#[derive(Debug)]
pub struct EntityStorage {
    data_by_id: IntMap<u32, EntityData>,
    id_by_chunk: HashMap<ChunkPos, IntSet<u32>>,
    id_by_uuid: HashMap<Uuid, u32>,
}

impl EntityStorage {
    pub fn new() -> Self {
        Self {
            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
            .entry(ChunkPos::from(entity.pos()))
            .or_default()
            .insert(id);
        self.id_by_uuid.insert(entity.uuid, id);
        self.data_by_id.insert(id, entity);
    }

    /// Remove an entity from the storage by its id.
    #[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.");
            }
        } 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.
    #[inline]
    pub fn contains_id(&self, id: &u32) -> bool {
        self.data_by_id.contains_key(id)
    }

    /// Get a reference to an entity by its id.
    #[inline]
    pub fn get_by_id(&self, id: u32) -> Option<&EntityData> {
        self.data_by_id.get(&id)
    }

    /// Get a mutable reference to an entity by its id.
    #[inline]
    pub fn get_mut_by_id(&mut self, id: u32) -> Option<&mut EntityData> {
        self.data_by_id.get_mut(&id)
    }

    /// Get a reference to an entity by its uuid.
    #[inline]
    pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<&EntityData> {
        self.id_by_uuid
            .get(uuid)
            .and_then(|id| self.data_by_id.get(id))
    }

    /// Get a mutable reference to an entity by its uuid.
    #[inline]
    pub fn get_mut_by_uuid(&mut self, uuid: &Uuid) -> Option<&mut EntityData> {
        self.id_by_uuid
            .get(uuid)
            .and_then(|id| self.data_by_id.get_mut(id))
    }

    /// Clear all entities in a chunk.
    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}.");
                }
            }
        }
    }

    /// Updates an entity from its old chunk.
    #[inline]
    pub fn update_entity_chunk(
        &mut self,
        entity_id: u32,
        old_chunk: &ChunkPos,
        new_chunk: &ChunkPos,
    ) {
        if let Some(entities) = self.id_by_chunk.get_mut(old_chunk) {
            entities.remove(&entity_id);
        }
        self.id_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>
    where
        F: FnMut(&EntityData) -> bool,
    {
        self.entities().find(|&entity| f(entity))
    }

    pub fn find_one_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<&EntityData>
    where
        F: FnMut(&EntityData) -> bool,
    {
        if let Some(entities) = self.id_by_chunk.get(chunk) {
            for entity_id in entities {
                if let Some(entity) = self.data_by_id.get(entity_id) {
                    if f(entity) {
                        return Some(entity);
                    }
                }
            }
        }
        None
    }
}

impl Default for EntityStorage {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use azalea_core::Vec3;

    #[test]
    fn test_store_entity() {
        let mut storage = EntityStorage::new();
        assert!(storage.get_by_id(0).is_none());

        let uuid = Uuid::from_u128(100);
        storage.insert(0, EntityData::new(uuid, Vec3::default()));
        assert_eq!(storage.get_by_id(0).unwrap().uuid, uuid);

        storage.remove_by_id(0);
        assert!(storage.get_by_id(0).is_none());
    }
}