diff --git a/examples/spawn_despawn_tilemap.rs b/examples/spawn_despawn_tilemap.rs index f8b924fa..2a9fb320 100644 --- a/examples/spawn_despawn_tilemap.rs +++ b/examples/spawn_despawn_tilemap.rs @@ -84,11 +84,8 @@ fn despawn_map(mut commands: Commands, mut maps: Query<(Entity, &mut TileStorage }; commands.entity(tilemap_entity).despawn_recursive(); - - for maybe_entity in tile_storage.iter_mut() { - if let Some(entity) = maybe_entity.take() { - commands.entity(entity).despawn(); - } + for entity in tile_storage.drain() { + commands.entity(entity).despawn(); } } diff --git a/src/tiles/storage.rs b/src/tiles/storage.rs index 15a71928..4988c854 100644 --- a/src/tiles/storage.rs +++ b/src/tiles/storage.rs @@ -86,20 +86,37 @@ impl TileStorage { self.tiles.iter_mut() } - /// Remove entity at the given tile position, if there was one, leaving `None` in its place. + /// Removes any stored `Entity` at the given tile position, leaving `None` in its place and + /// returning the `Entity`. /// /// Panics if the given `tile_pos` doesn't lie within the extents of the underlying tile map. - pub fn remove(&mut self, tile_pos: &TilePos) { - self.tiles[tile_pos.to_index(&self.size)].take(); + pub fn remove(&mut self, tile_pos: &TilePos) -> Option { + self.tiles[tile_pos.to_index(&self.size)].take() } - /// Remove any stored entity at the given tile position, if the given `tile_pos` does lie within - /// the extents of the underlying map. + /// Remove any stored `Entity` at the given tile position, leaving `None` in its place and + /// returning the `Entity`. /// - /// Otherwise, nothing is done. - pub fn checked_remove(&mut self, tile_pos: &TilePos) { - if tile_pos.within_map_bounds(&self.size) { - self.tiles[tile_pos.to_index(&self.size)].take(); - } + /// Checks that the given `tile_pos` lies within the extents of the underlying map. + pub fn checked_remove(&mut self, tile_pos: &TilePos) -> Option { + self.tiles.get_mut(tile_pos.to_index(&self.size))?.take() + } + + /// Removes all stored `Entity`s, leaving `None` in their place and + /// returning them in an iterator. + /// + /// Example: + /// ``` + /// # use bevy::prelude::Commands; + /// # use bevy_ecs_tilemap::prelude::{TilemapSize, TileStorage}; + /// # fn example(mut commands: Commands) { + /// # let mut storage = TileStorage::empty(TilemapSize { x: 16, y: 16 }); + /// for entity in storage.drain() { + /// commands.entity(entity).despawn(); + /// } + /// # } + /// ``` + pub fn drain(&mut self) -> impl Iterator + use<'_> { + self.tiles.iter_mut().filter_map(|opt| opt.take()) } }