iter_combinations/
iter_combinations.rs

1//! Shows how to iterate over combinations of query results.
2
3use bevy::{color::palettes::css::ORANGE_RED, math::FloatPow, prelude::*};
4use rand::{Rng, SeedableRng};
5use rand_chacha::ChaCha8Rng;
6
7fn main() {
8    App::new()
9        .add_plugins(DefaultPlugins)
10        .insert_resource(ClearColor(Color::BLACK))
11        .add_systems(Startup, generate_bodies)
12        .add_systems(FixedUpdate, (interact_bodies, integrate))
13        .add_systems(Update, look_at_star)
14        .run();
15}
16
17const GRAVITY_CONSTANT: f32 = 0.001;
18const NUM_BODIES: usize = 100;
19
20#[derive(Component, Default)]
21struct Mass(f32);
22#[derive(Component, Default)]
23struct Acceleration(Vec3);
24#[derive(Component, Default)]
25struct LastPos(Vec3);
26#[derive(Component)]
27struct Star;
28
29#[derive(Bundle, Default)]
30struct BodyBundle {
31    mesh: Mesh3d,
32    material: MeshMaterial3d<StandardMaterial>,
33    mass: Mass,
34    last_pos: LastPos,
35    acceleration: Acceleration,
36}
37
38fn generate_bodies(
39    time: Res<Time<Fixed>>,
40    mut commands: Commands,
41    mut meshes: ResMut<Assets<Mesh>>,
42    mut materials: ResMut<Assets<StandardMaterial>>,
43) {
44    let mesh = meshes.add(Sphere::new(1.0).mesh().ico(3).unwrap());
45
46    let color_range = 0.5..1.0;
47    let vel_range = -0.5..0.5;
48
49    // We're seeding the PRNG here to make this example deterministic for testing purposes.
50    // This isn't strictly required in practical use unless you need your app to be deterministic.
51    let mut rng = ChaCha8Rng::seed_from_u64(19878367467713);
52    for _ in 0..NUM_BODIES {
53        let radius: f32 = rng.random_range(0.1..0.7);
54        let mass_value = FloatPow::cubed(radius) * 10.;
55
56        let position = Vec3::new(
57            rng.random_range(-1.0..1.0),
58            rng.random_range(-1.0..1.0),
59            rng.random_range(-1.0..1.0),
60        )
61        .normalize()
62            * ops::cbrt(rng.random_range(0.2f32..1.0))
63            * 15.;
64
65        commands.spawn((
66            BodyBundle {
67                mesh: Mesh3d(mesh.clone()),
68                material: MeshMaterial3d(materials.add(Color::srgb(
69                    rng.random_range(color_range.clone()),
70                    rng.random_range(color_range.clone()),
71                    rng.random_range(color_range.clone()),
72                ))),
73                mass: Mass(mass_value),
74                acceleration: Acceleration(Vec3::ZERO),
75                last_pos: LastPos(
76                    position
77                        - Vec3::new(
78                            rng.random_range(vel_range.clone()),
79                            rng.random_range(vel_range.clone()),
80                            rng.random_range(vel_range.clone()),
81                        ) * time.timestep().as_secs_f32(),
82                ),
83            },
84            Transform {
85                translation: position,
86                scale: Vec3::splat(radius),
87                ..default()
88            },
89        ));
90    }
91
92    // add bigger "star" body in the center
93    let star_radius = 1.;
94    commands
95        .spawn((
96            BodyBundle {
97                mesh: Mesh3d(meshes.add(Sphere::new(1.0).mesh().ico(5).unwrap())),
98                material: MeshMaterial3d(materials.add(StandardMaterial {
99                    base_color: ORANGE_RED.into(),
100                    emissive: LinearRgba::from(ORANGE_RED) * 2.,
101                    ..default()
102                })),
103
104                mass: Mass(500.0),
105                ..default()
106            },
107            Transform::from_scale(Vec3::splat(star_radius)),
108            Star,
109        ))
110        .with_child(PointLight {
111            color: Color::WHITE,
112            range: 100.0,
113            radius: star_radius,
114            ..default()
115        });
116    commands.spawn((
117        Camera3d::default(),
118        Transform::from_xyz(0.0, 10.5, -30.0).looking_at(Vec3::ZERO, Vec3::Y),
119    ));
120}
121
122fn interact_bodies(mut query: Query<(&Mass, &GlobalTransform, &mut Acceleration)>) {
123    let mut iter = query.iter_combinations_mut();
124    while let Some([(Mass(m1), transform1, mut acc1), (Mass(m2), transform2, mut acc2)]) =
125        iter.fetch_next()
126    {
127        let delta = transform2.translation() - transform1.translation();
128        let distance_sq: f32 = delta.length_squared();
129
130        let f = GRAVITY_CONSTANT / distance_sq;
131        let force_unit_mass = delta * f;
132        acc1.0 += force_unit_mass * *m2;
133        acc2.0 -= force_unit_mass * *m1;
134    }
135}
136
137fn integrate(time: Res<Time>, mut query: Query<(&mut Acceleration, &mut Transform, &mut LastPos)>) {
138    let dt_sq = time.delta_secs() * time.delta_secs();
139    for (mut acceleration, mut transform, mut last_pos) in &mut query {
140        // verlet integration
141        // x(t+dt) = 2x(t) - x(t-dt) + a(t)dt^2 + O(dt^4)
142
143        let new_pos = transform.translation * 2.0 - last_pos.0 + acceleration.0 * dt_sq;
144        acceleration.0 = Vec3::ZERO;
145        last_pos.0 = transform.translation;
146        transform.translation = new_pos;
147    }
148}
149
150fn look_at_star(
151    mut camera: Single<&mut Transform, (With<Camera>, Without<Star>)>,
152    star: Single<&Transform, With<Star>>,
153) {
154    let new_rotation = camera
155        .looking_at(star.translation, Vec3::Y)
156        .rotation
157        .lerp(camera.rotation, 0.1);
158    camera.rotation = new_rotation;
159}
iter_combinations/iter_combinations.rs

iter_combinations/
iter_combinations.rs
