Restart on game over

src/main.rs

@@ -30,11 +30,23 @@ struct SnakeHead {
 }
 struct Materials {
     head_material: Handle<ColorMaterial>,
+    segment_material: Handle<ColorMaterial>,
     food_material: Handle<ColorMaterial>,
 }
 
 struct SnakeMoveTimer(Timer);
 
+struct GameOverEvent;
+struct GrowthEvent;
+
+#[derive(Default)]
+struct LastTailPosition(Option<Position>);
+
+struct SnakeSegment;
+
+#[derive(Default)]
+struct SnakeSegments(Vec<Entity>);
+
 struct Food;
 
 struct FoodSpawnTimer(Timer);
@@ -66,31 +78,71 @@ fn setup(commands: &mut Commands, mut materials: ResMut<Assets<ColorMaterial>>)
     commands.spawn(Camera2dBundle::default());
     commands.insert_resource(Materials {
         head_material: materials.add(Color::rgb(0.7, 0.7, 0.7).into()),
+        segment_material: materials.add(Color::rgb(0.3, 0.3, 0.3).into()),
         food_material: materials.add(Color::rgb(1.0, 0.0, 1.0).into()),
     });
 }
 
-fn spawn_snake(commands: &mut Commands, materials: Res<Materials>) {
+fn spawn_snake(
+    commands: &mut Commands,
+    materials: Res<Materials>,
+    mut segments: ResMut<SnakeSegments>,
+) {
+    segments.0 = vec![
+        commands
+            .spawn(SpriteBundle {
+                material: materials.head_material.clone(),
+                sprite: Sprite::new(Vec2::new(10.0, 10.0)),
+                ..Default::default()
+            })
+            .with(SnakeHead {
+                direction: Direction::Up,
+            })
+            .with(SnakeSegment)
+            .with(Position { x: 3, y: 3 })
+            .with(Size::square(0.8))
+            .current_entity()
+            .unwrap(),
+        spawn_segment(
+            commands,
+            &materials.segment_material,
+            Position { x: 3, y: 2 },
+        ),
+    ];
+}
+
+fn spawn_segment(
+    commands: &mut Commands,
+    material: &Handle<ColorMaterial>,
+    position: Position,
+) -> Entity {
     commands
         .spawn(SpriteBundle {
-            material: materials.head_material.clone(),
+            material: material.clone(),
-            sprite: Sprite::new(Vec2::new(10.0, 10.0)),
             ..Default::default()
         })
-        .with(SnakeHead {
+        .with(SnakeSegment)
-            direction: Direction::Up,
+        .with(position)
-        })
+        .with(Size::square(0.65))
-        .with(Position { x: 3, y: 3 })
+        .current_entity()
-        .with(Size::square(0.8));
+        .unwrap()
 }
 
 fn snake_movement(
     keyboard_input: Res<Input<KeyCode>>,
     snake_timer: ResMut<SnakeMoveTimer>,
+    mut game_over_events: ResMut<Events<GameOverEvent>>,
+    mut last_tail_position: ResMut<LastTailPosition>,
+    segments: ResMut<SnakeSegments>,
     mut heads: Query<(Entity, &mut SnakeHead)>,
     mut positions: Query<&mut Position>,
 ) {
     if let Some((head_entity, mut head)) = heads.iter_mut().next() {
+        let segment_positions = segments
+            .0
+            .iter()
+            .map(|e| *positions.get_mut(*e).unwrap())
+            .collect::<Vec<Position>>();
         let mut head_pos = positions.get_mut(head_entity).unwrap();
         let dir: Direction = if keyboard_input.pressed(KeyCode::Left) {
             Direction::Left
@@ -123,6 +175,82 @@ fn snake_movement(
                 head_pos.y -= 1;
             }
         };
+
+        // Check if we've hit a wall
+        if head_pos.x < 0
+            || head_pos.y < 0
+            || head_pos.x as u32 >= ARENA_WIDTH
+            || head_pos.y as u32 >= ARENA_HEIGHT
+        {
+            game_over_events.send(GameOverEvent);
+        }
+
+        // Check if we've hit our tail
+        if segment_positions.contains(&head_pos) {
+            game_over_events.send(GameOverEvent);
+        }
+
+        segment_positions
+            .iter()
+            .zip(segments.0.iter().skip(1))
+            .for_each(|(pos, segment)| {
+                *positions.get_mut(*segment).unwrap() = *pos;
+            });
+        last_tail_position.0 = Some(*segment_positions.last().unwrap());
+    }
+}
+
+fn game_over(
+    commands: &mut Commands,
+    mut reader: Local<EventReader<GameOverEvent>>,
+    game_over_events: Res<Events<GameOverEvent>>,
+    materials: Res<Materials>,
+    segments_res: ResMut<SnakeSegments>,
+    food: Query<Entity, With<Food>>,
+    segments: Query<Entity, With<SnakeSegment>>,
+) {
+    if reader.iter(&game_over_events).next().is_some() {
+        for ent in food.iter().chain(segments.iter()) {
+            commands.despawn(ent);
+        }
+        spawn_snake(commands, materials, segments_res);
+    }
+}
+
+fn snake_eating(
+    commands: &mut Commands,
+    snake_timer: ResMut<SnakeMoveTimer>,
+    mut growth_events: ResMut<Events<GrowthEvent>>,
+    food_positions: Query<(Entity, &Position), With<Food>>,
+    head_positions: Query<&Position, With<SnakeHead>>,
+) {
+    if !snake_timer.0.finished() {
+        return;
+    }
+    for head_pos in head_positions.iter() {
+        for (ent, food_pos) in food_positions.iter() {
+            if food_pos == head_pos {
+                commands.despawn(ent);
+                growth_events.send(GrowthEvent);
+            }
+        }
+    }
+}
+
+fn snake_growth(
+    commands: &mut Commands,
+    last_tail_position: Res<LastTailPosition>,
+    growth_events: Res<Events<GrowthEvent>>,
+    mut segments: ResMut<SnakeSegments>,
+    mut growth_reader: Local<EventReader<GrowthEvent>>,
+    materials: Res<Materials>,
+) {
+    if growth_reader.iter(&growth_events).next().is_some() {
+        segments.0.push(spawn_segment(
+            commands,
+            &materials.segment_material,
+            last_tail_position.0.unwrap(),
+        ));
     }
 }
 
@@ -186,10 +314,11 @@ fn main() {
             height: 500.0,
             ..Default::default()
         })
-        .add_resource(SnakeMoveTimer(Timer::new(
+        .add_resource(SnakeMoveTimer(Timer::new(Duration::from_millis(150), true)))
-            Duration::from_millis(150),
+        .add_resource(SnakeSegments::default())
-            true,
+        .add_resource(LastTailPosition::default())
-        )))
+        .add_event::<GrowthEvent>()
+        .add_event::<GameOverEvent>()
         .add_startup_system(setup.system())
         .add_startup_stage("game_setup", SystemStage::single(spawn_snake.system()))
         .add_system(snake_movement.system())
@@ -197,6 +326,9 @@ fn main() {
         .add_system(size_scaling.system())
         .add_system(food_spawner.system())
         .add_system(snake_timer.system())
+        .add_system(snake_eating.system())
+        .add_system(snake_growth.system())
+        .add_system(game_over.system())
         .add_plugins(DefaultPlugins)
         .run();
 }