ark-api 0.16.0

use super::{ffi, ComponentTrait, Entity, Layer, Transform, WorldData};

/// A `WorldEntityQuery` can be used to retrieve entities in the world
/// using various requirements.
pub struct WorldEntityQuery {
    data: WorldData,
}

impl WorldEntityQuery {
    /// Creates a new `WorldEntityQuery` using the ffi struct.
    pub fn new(query: &ffi::WorldEntityQuery) -> Self {
        Self {
            data: WorldData::create_struct(ffi::CreateDataType::WorldEntityQuery, query),
        }
    }

    // TODO: retrieve should perform query so you can reuse this object.
    /// Gets all entities this query resulted in.
    pub fn entities(&self) -> Vec<Entity> {
        let handle = self.data.get_data_handle();
        handle.read_vec()
    }
}

/// Represents an entity layer filter that can be used with entity queries or raycasts.
#[derive(Copy, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "with_serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "with_speedy", derive(speedy::Writable, speedy::Readable))]
#[repr(transparent)]
pub struct EntityLayerMask(u64);

impl EntityLayerMask {
    /// Returns a new mask that includes only the given layer.
    pub const fn layer(layer: usize) -> Self {
        Self(Layer::mask_from_index(layer))
    }

    /// Returns a new mask that includes the union of the layers of this mask and the `other` mask.
    pub const fn with(self, other: EntityLayerMask) -> Self {
        Self(self.0 | other.0)
    }

    /// Returns a new mask that includes the layers that this mask includes except the ones in the `other` mask.
    pub const fn except(self, other: EntityLayerMask) -> Self {
        Self(self.0 & !other.0)
    }

    /// Returns a new mask that also includes the given layer.
    pub const fn with_layer(self, layer: usize) -> Self {
        Self(self.0 | Layer::mask_from_index(layer))
    }

    /// Returns a new mask that excludes the given layer.
    /// If the layer was already excluded from the mask this has no effect.
    pub const fn except_layer(self, layer: usize) -> Self {
        Self(self.0 & !Layer::mask_from_index(layer))
    }

    /// Create from a raw mask value.
    pub const fn from_value(value: u64) -> Self {
        Self(value)
    }

    /// Returns the raw mask value.
    /// Bits set to 1 indicate that the layer should be included.
    pub const fn value(&self) -> u64 {
        self.0
    }

    /// All layers.
    #[inline]
    pub const fn everything() -> Self {
        Self(!0)
    }

    /// No layers.
    #[inline]
    pub const fn nothing() -> Self {
        Self(0)
    }
}

impl std::ops::BitOr for EntityLayerMask {
    type Output = Self;
    /// Returns a new mask that includes the union of the layers of this mask and the `other` mask.
    ///
    /// Identical to `EntityLayerMask::with`.
    fn bitor(self, rhs: Self) -> Self::Output {
        Self(self.0 | rhs.0)
    }
}

impl std::ops::BitOrAssign<Self> for EntityLayerMask {
    /// Returns a new mask that includes the union of the layers of this mask and the `other` mask.
    ///
    /// Identical to `EntityLayerMask::with`.
    fn bitor_assign(&mut self, rhs: Self) {
        self.0 |= rhs.0;
    }
}

impl std::fmt::Debug for EntityLayerMask {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match *self {
            x if x == Self::everything() => write!(f, "LayerMask(everything)"),
            x if x == Self::nothing() => write!(f, "LayerMask(nothing)"),
            x if x.0.count_ones() == 1 => write!(f, "LayerMask(layer {})", self.0.trailing_zeros()),
            _ => write!(f, "LayerMask({:#b})", self.0),
        }
    }
}

/// Represents an entity query type that can be used with entity queries
/// to specify what the query should include before filtering.
#[derive(Copy, Clone)]
pub enum EntityQueryType {
    /// Query all entities in the world. Filtering by layer is optional.
    AllEntities,

    /// Query all direct children, excluding the root. Always filters by layer.
    TransformChildren(Entity),

    /// Query all indirect and direct children, excluding the root. Always filters by layer.
    TransformDescendants(Entity),
}

/// Use a `WorldEntityQueryBuilder` to build a `WorldEntityQuery` with an ergonomic interface.
#[derive(Copy, Clone)]
pub struct WorldEntityQueryBuilder {
    query: ffi::WorldEntityQuery,
    layer_filter: Option<EntityLayerMask>,
}

impl Default for WorldEntityQueryBuilder {
    fn default() -> Self {
        Self::new(EntityQueryType::AllEntities)
    }
}

impl WorldEntityQueryBuilder {
    /// Builds a new `WorldEntityQueryBuilder` from a `EntityQueryType`
    pub fn new(query_type: EntityQueryType) -> Self {
        let (query_type, root_entity) = match query_type {
            EntityQueryType::AllEntities => {
                (ffi::WorldEntityQueryType::AllEntities, Entity::invalid())
            }
            EntityQueryType::TransformChildren(entity) => {
                (ffi::WorldEntityQueryType::TransformChildren, entity)
            }
            EntityQueryType::TransformDescendants(entity) => {
                (ffi::WorldEntityQueryType::TransformDescendants, entity)
            }
        };

        Self {
            query: ffi::WorldEntityQuery {
                layer_mask: !0u64,
                num_include_tags: 0,
                num_exclude_tags: 0,
                include_tags: Default::default(),
                exclude_tags: Default::default(),

                query_type,
                root_entity: root_entity.as_ffi(),
                _pad0: Default::default(),
                _pad1: Default::default(),
            },
            layer_filter: None,
        }
    }

    /// Sets a layer filter on the query to be built.
    ///
    /// Note that not calling this means to return everything regardless of layer. This cannot
    /// currently be expressed with an `EntityLayerMask` as `everything` will exclude entities
    /// that belong to no layer.
    pub fn with_layer_filter(&mut self, layer_filter: EntityLayerMask) -> &mut Self {
        self.layer_filter = Some(layer_filter);
        self
    }

    /// Add a tag to be included in the query
    pub fn with_tag(&mut self, tag: u64) -> &mut Self {
        if (self.query.num_include_tags as usize) >= self.query.include_tags.len() {
            log::warn!(
                "WorldEntityQueryBuilder: Trying to include more tags than allowed. Max {}",
                self.query.include_tags.len()
            );
            return self;
        }

        self.query.include_tags[self.query.num_include_tags as usize] = tag;
        self.query.num_include_tags += 1;
        self
    }

    /// Exclude a tag from the query
    pub fn without_tag(&mut self, tag: u64) -> &mut Self {
        if (self.query.num_exclude_tags as usize) >= self.query.exclude_tags.len() {
            log::warn!(
                "WorldEntityQueryBuilder: Trying to exclude more tags than allowed. Max {}",
                self.query.exclude_tags.len()
            );
            return self;
        }

        self.query.exclude_tags[self.query.num_exclude_tags as usize] = tag;
        self.query.num_exclude_tags += 1;
        self
    }

    /// Build the `WorldEntityQuery`
    pub fn build(&mut self) -> WorldEntityQuery {
        if let Some(layer_mask) = self.layer_filter {
            self.query.layer_mask = layer_mask.value();
        } else if self.query.query_type == ffi::WorldEntityQueryType::AllEntities {
            self.query.query_type = ffi::WorldEntityQueryType::AllEntitiesNoLayerFilter;
        } else if self.query.query_type == ffi::WorldEntityQueryType::TransformChildren {
            self.query.query_type = ffi::WorldEntityQueryType::TransformChildrenNoLayerFilter;
        } else if self.query.query_type == ffi::WorldEntityQueryType::TransformDescendants {
            self.query.query_type = ffi::WorldEntityQueryType::TransformDescendantsNoLayerFilter;
        } else {
            self.query.layer_mask = !0u64;
        }
        WorldEntityQuery::new(&self.query)
    }
}

impl Transform {
    /// Retrieves all the direct transform children of this entity
    pub fn children(&self) -> Vec<Entity> {
        WorldEntityQueryBuilder::new(EntityQueryType::TransformChildren(self.entity()))
            .build()
            .entities()
    }

    /// Retrieves all the transform descendants of this entity.
    pub fn descendants(&self) -> Vec<Entity> {
        WorldEntityQueryBuilder::new(EntityQueryType::TransformDescendants(self.entity()))
            .build()
            .entities()
    }
}