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
170
171
172
173
174
175
176
177
178
179
180
181
182
183
#![doc = include_str!("../README.md")]
use std::cmp::Reverse;
use std::collections::HashMap;
use std::fmt::Debug;
use std::hash::Hash;
use std::marker::PhantomData;

use bevy::render::RenderApp;
use bevy::sprite::{extract_sprites, queue_sprites, ExtractedSprite, SpriteSystem};
use bevy::{prelude::*, render::Extract, sprite::ExtractedSprites};
use ordered_float::OrderedFloat;
#[cfg(feature = "parallel_y_sort")]
use rayon::slice::ParallelSliceMut;

/// This plugin will modify the z-coordinates of the extracted sprites stored
/// in Bevy's [`ExtractedSprites`] so that they're rendered in the proper
/// order. See the crate documentation for how to use it.
///
/// In general you should only instantiate this plugin with a single type you
/// use throughout your program.
///
/// By default your sprites will also be y-sorted. If you don't need this,
/// replace the [`SpriteLayerOptions`] like so:
///
/// ```
/// # use bevy::prelude::*;
/// # use extol_sprite_layer::SpriteLayerOptions;
/// # let mut app = App::new();
/// app.insert_resource(SpriteLayerOptions { y_sort: false });
/// ```
pub struct SpriteLayerPlugin<Layer> {
    phantom: PhantomData<Layer>,
}

impl<Layer> Default for SpriteLayerPlugin<Layer> {
    fn default() -> Self {
        Self {
            phantom: Default::default(),
        }
    }
}

impl<Layer: LayerIndex> Plugin for SpriteLayerPlugin<Layer> {
    fn build(&self, app: &mut App) {
        app.init_resource::<SpriteLayerOptions>();
        let render_app = app.sub_app_mut(RenderApp);
        render_app.add_systems(
            ExtractSchedule,
            update_sprite_z_coordinates::<Layer>
                .in_set(SpriteSystem::ExtractSprites)
                .in_set(SpriteLayerSet)
                .after(extract_sprites)
                .before(queue_sprites),
        );
    }
}

/// Configure how the sprite layer
#[derive(Debug, Resource, Reflect)]
pub struct SpriteLayerOptions {
    pub y_sort: bool,
}

impl Default for SpriteLayerOptions {
    fn default() -> Self {
        Self { y_sort: true }
    }
}

/// Set for all systems related to [`SpriteLayerPlugin`]. This is run in the
/// render app's [`ExtractSchedule`], *not* the main app.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, SystemSet)]
pub struct SpriteLayerSet;

/// Trait for the type you use to indicate your sprites' layers. Add this as a
/// component to any entity you want to treat as a sprite. Note that this does
/// *not* propagate.
pub trait LayerIndex: Eq + Hash + Component + Clone + Debug {
    /// The actual numeric z-value that the layer index corresponds to.  Note
    /// that the z-value for an entity can be any value in the range
    /// `layer.as_z_coordinate() <= z < layer.as_z_coordinate() + 1.0`, and the
    /// exact values are an implementation detail!
    ///
    /// With the default Bevy camera settings, your return values from this
    /// function should be between 0 and 999.0, since the camera is at z =
    /// 1000.0. Prefer smaller z-values since that gives more precision.
    fn as_z_coordinate(&self) -> f32;
}

/// Update the z-coordinates of the transform of every sprite with a
/// `LayerIndex` component so that they're rendered in the proper layer with
/// y-sorting.
#[allow(clippy::type_complexity)]
fn update_sprite_z_coordinates<Layer: LayerIndex>(
    mut extracted_sprites: ResMut<ExtractedSprites>,
    options: Extract<Res<SpriteLayerOptions>>,
    transform_query: Extract<Query<(Entity, &GlobalTransform), With<Layer>>>,
    layer_query: Extract<Query<&Layer>>,
) {
    if options.y_sort {
        let z_index_map = map_z_indices(transform_query, layer_query);
        for (entity, sprite) in extracted_sprites.sprites.iter_mut() {
            if let Some(z) = z_index_map.get(entity) {
                set_sprite_coordinate(sprite, *z);
            }
        }
    } else {
        for (entity, sprite) in extracted_sprites.sprites.iter_mut() {
            if let Ok(layer) = layer_query.get(*entity) {
                set_sprite_coordinate(sprite, layer.as_z_coordinate());
            }
        }
    }
}

/// Sets the z-coordinate of the sprite's transform.
fn set_sprite_coordinate(sprite: &mut ExtractedSprite, z: f32) {
    if sprite.transform.translation().z != 0.0 {
        // not currently disableable, but I'm open if you file an issue :)
        warn!(
            "Entity {:?} has a LabelLayer *and* a nonzero z-coordinate {}; this is probably not what you want!",
            sprite.original_entity,
            sprite.transform.translation().z
        );
    }
    // hacky hacky; I can't find a way to directly mutate the GlobalTransform.
    let mut affine = sprite.transform.affine();
    affine.translation.z = z;
    sprite.transform = GlobalTransform::from(affine);
}

/// Used to sort the entities within a sprite layer.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
struct ZIndexSortKey(Reverse<OrderedFloat<f32>>);

impl ZIndexSortKey {
    // This is reversed because bevy uses +y pointing upwards, which is the
    // opposite of what you generally want.
    fn new(transform: &GlobalTransform) -> Self {
        Self(Reverse(OrderedFloat(transform.translation().y)))
    }
}

/// Determines the z-value to use for each entity. The z-value is set to
/// `layer.as_z_coordinate() + offset`, where `offset` is calculated so that
/// entities with a higher y-coordinate have a higher offset.
#[allow(clippy::type_complexity)]
fn map_z_indices<Layer: LayerIndex>(
    transform_query: Extract<Query<(Entity, &GlobalTransform), With<Layer>>>,
    layer_query: Extract<Query<&Layer>>,
) -> HashMap<Entity, f32> {
    // We y-sort everything because this avoids the overhead of grouping
    // entities by their layer. Using sort_by_cached_key to make the vec's
    // elements smaller doesn't seem to help here.
    let mut all_entities: Vec<(ZIndexSortKey, Entity)> = transform_query
        .iter()
        .map(|(entity, transform)| (ZIndexSortKey::new(transform), entity))
        .collect();

    // most of the expense is here.
    #[cfg(feature = "parallel_y_sort")]
    all_entities.par_sort_unstable();
    #[cfg(not(feature = "parallel_y_sort"))]
    all_entities.sort_unstable();

    let scale_factor = 1.0 / all_entities.len() as f32;
    all_entities
        .into_iter()
        .enumerate()
        .map(|(i, (_, entity))| {
            (
                entity,
                // NOTE: it's possible that the scale factor will be small
                // enough relative to the z coordinate that these are equal for
                // consecutive values. This occurs when z-coordinate *
                // len(all_entities) > 2^23 (floats have 24 bits of
                // precision). Even with a z-coordinate of 1000, this requires
                // over 8000 entities to hit, which I think is fine.
                layer_query.get(entity).unwrap().as_z_coordinate() + i as f32 * scale_factor,
            )
        })
        .collect()
}