Use train position history to reserve signal blocks

tests/test_game.py

@@ -828,3 +828,93 @@ def test_if_a_train_is_destroyed_the_signals_become_green(self, game: Game):
         train.destroy()
 
         assert game.signal_controller._signal_blocks[0].reserved_by is None
+
+
+class TestReserveAndUnreserveRail:
+    def test_train_reserves_block_when_moving(self, game: Game):
+        """Sends a train right and asserts that it stops reserving the first block
+        and begins to reserve the other instead."""
+        create_objects(
+            game,
+            """
+            . M . F .
+
+            .-S-.-S-.
+            """,
+        )
+        train = game._create_train(*game.grid.stations.values())
+        game.on_update(1 / 60)
+
+        assert game.signal_controller._signal_blocks[0].reserved_by == id(train)
+
+    def test_a_train_without_wagons_reserves_both_blocks_when_leaving(self, game: Game):
+        """Sends a train right and asserts that eventually it reserves both blocks"""
+        create_objects(
+            game,
+            """
+            . M . F .
+
+            .-Sh.-S-.
+            """,
+        )
+        train = game._create_train(*game.grid.stations.values(), wagon_count=0)
+        game.on_update(1 / 60)
+        left_signal_block = game.signal_controller._signal_block_from_position[
+            Vec2(30, 0)
+        ]
+        right_signal_block = game.signal_controller._signal_block_from_position[
+            Vec2(120, 0)
+        ]
+        assert left_signal_block.reserved_by == id(train)
+        while not right_signal_block.reserved_by:
+            game.on_update(1 / 60)
+
+        assert left_signal_block.reserved_by == id(train)
+        assert right_signal_block.reserved_by == id(train)
+
+    def test_a_train_without_wagons_eventually_unreserves_first_block(self, game: Game):
+        """Sends a train right and asserts that eventually it reserves both blocks"""
+        create_objects(
+            game,
+            """
+            . M . F .
+
+            .-Sh.-S-.
+            """,
+        )
+        train = game._create_train(*game.grid.stations.values(), wagon_count=0)
+        game.on_update(1 / 60)
+        left_signal_block = game.signal_controller._signal_block_from_position[
+            Vec2(30, 0)
+        ]
+        right_signal_block = game.signal_controller._signal_block_from_position[
+            Vec2(120, 0)
+        ]
+        while left_signal_block.reserved_by == id(train):
+            game.on_update(1 / 60)
+
+        assert right_signal_block.reserved_by == id(train)
+
+    def test_a_train_with_one_wagon_can_reserve_two_signal_blocks(self, game: Game):
+        """Sends a train right and asserts that it eventually reserves both blocks"""
+        create_objects(
+            game,
+            """
+            . M . F .
+
+            .-Sh.-S-.
+            """,
+        )
+        train = game._create_train(*game.grid.stations.values(), wagon_count=1)
+
+        game.on_update(1 / 60)
+        assert game.signal_controller._signal_blocks[0].reserved_by == id(train)
+
+        while not game.signal_controller._signal_blocks[1].reserved_by:
+            game.on_update(1 / 60)
+
+        assert (
+            game.signal_controller._signal_blocks[1].reserved_by
+            == game.signal_controller._signal_blocks[0].reserved_by
+            == id(train)
+        )

tests/test_model.py

@@ -57,7 +57,9 @@ def train(player, mock_grid: Grid):
     station1 = Station(Vec2(0, 0))
     station2 = Station(Vec2(30, 0))
     mock_grid.create_rail([Rail(0, 0, 30, 0)])
-    return Train(player, station1, station2, mock_grid, SignalController())
+    return Train(
+        player, station1, station2, mock_grid, SignalController(), wagon_count=3
+    )
 
 
 class TestTrain:

trainfinity2/game.py

@@ -207,9 +207,16 @@ def create_signals_at_grid_position(self, x, y) -> list[Signal]:
             signal.add_observer(self.drawer, ChangeEvent)
         return signals
 
-    def _create_train(self, station1: Station, station2: Station):
+    def _create_train(
+        self, station1: Station, station2: Station, *, wagon_count: int = 3
+    ):
         train = Train(
-            self.player, station1, station2, self.grid, self.signal_controller
+            self.player,
+            station1,
+            station2,
+            self.grid,
+            self.signal_controller,
+            wagon_count=wagon_count,
         )
         self.trains.append(train)
         train.add_observer(self, DestroyEvent)

trainfinity2/grid.py

@@ -138,13 +138,13 @@ def _create_in_random_unoccupied_location(
     def _create_factories(self):
         self._create_factory_in_random_unoccupied_location()
 
-    def snap_to(self, x, y) -> Vec2:
+    def snap_to(self, x: float, y: float) -> Vec2:
         return Vec2(self.snap_to_x(x), self.snap_to_y(y))
 
-    def snap_to_x(self, x) -> int:
+    def snap_to_x(self, x: float) -> int:
         return math.floor(x / GRID_BOX_SIZE) * GRID_BOX_SIZE
 
-    def snap_to_y(self, y) -> int:
+    def snap_to_y(self, y: float) -> int:
         return math.floor(y / GRID_BOX_SIZE) * GRID_BOX_SIZE
 
     def find_route_between_stations(

trainfinity2/signal_controller.py

@@ -1,4 +1,6 @@
+from collections import defaultdict
 from dataclasses import dataclass
+from typing import Iterable
 
 from pyglet.math import Vec2
 
@@ -55,7 +57,9 @@ def __init__(
         # self._signal_blocks: list[SignalBlock] = []
         self._signal_block_from_position: dict[Vec2, SignalBlock] = {}
         self._signals: list[Signal] = []
-        self._reserved_position_from_reserver_id: dict[int, Vec2] = {}
+        self._reserved_positions_from_reserver_id: dict[int, set[Vec2]] = defaultdict(
+            set
+        )
 
     def __repr__(self) -> str:
         return (
@@ -147,23 +151,16 @@ def create_signal_blocks(
     def reserver(self, position: Vec2) -> int | None:
         return self._signal_block_from_position[position].reserved_by
 
-    def reserve(self, reserver_id: int, position: Vec2):
-        """Called by trains when they enter a new rail. The correct blocks
-        are then reserved and unreserved."""
-        self._reserved_position_from_reserver_id[reserver_id] = position
-        self._update_signal_block_reservations()
-
-    def unreserve(self, reserver_id: int):
-        """Called by a train when it is destroyed."""
-        if reserver_id in self._reserved_position_from_reserver_id:
-            del self._reserved_position_from_reserver_id[reserver_id]
+    def reserve(self, reserver_id: int, positions: Iterable[Vec2]):
+        """Called by trains when they enter a new rail, or when they are destroyed."""
+        self._reserved_positions_from_reserver_id[reserver_id] = set(positions)
         self._update_signal_block_reservations()
 
     def _update_signal_block_reservations(self):
         for signal_block in self._signal_blocks:
             signal_block.reserved_by = None
-            for id_, position in self._reserved_position_from_reserver_id.items():
-                if position in signal_block.positions:
+            for id_, positions in self._reserved_positions_from_reserver_id.items():
+                if positions & signal_block.positions:
                     signal_block.reserved_by = id_
         self._update_signals()
 

trainfinity2/train.py

@@ -6,6 +6,7 @@
 from typing import Sequence
 from pyglet.math import Vec2
 
+
 from .constants import GRID_BOX_SIZE
 from .grid import Grid
 from .model import Player, Rail, Station
@@ -76,6 +77,7 @@ class Train(Subject):
     second_station: Station
     grid: Grid
     signal_controller: SignalController
+    wagon_count: int
     x: float = field(init=False)
     y: float = field(init=False)
     target_x: float = field(init=False)
@@ -85,7 +87,7 @@ class Train(Subject):
     selected = False
     wait_timer: float = 0.0
     angle: float = 0
-    speed = 0.0  # Pixels per second
+    speed: float = 0.0  # Pixels per second
 
     MAX_SPEED = 120.0  # 60.0  # Pixels per second
     ACCELERATION = 40.0  # Pixels per second squared
@@ -98,15 +100,13 @@ def __post_init__(self):
         self.target_y = self.y
         self._target_station = self.first_station
         self._rails_on_route: list[Rail] | None = []
-        self._position_history: deque[Vec2] = deque(
-            maxlen=4
-        )  # TODO: needs to be more than number of wagons
+        # The position history needs to be approximately as long as the train,
+        # since it is used for reserving positions. As long as one wagon is
+        # approximately as long as a block, this will do.
+        self._position_history: deque[Vec2] = deque(maxlen=self.wagon_count + 1)
         self._previous_targets_y = []
         # TODO: wagons are now created on top of train
-        self.wagons = []
-        self.wagons.append(Wagon(self.x, self.y))
-        self.wagons.append(Wagon(self.x, self.y))
-        self.wagons.append(Wagon(self.x, self.y))
+        self.wagons = [Wagon(self.x, self.y) for _ in range(self.wagon_count)]
 
     @property
     def rails_on_route(self):
@@ -163,9 +163,7 @@ def is_at(self, x, y):
         )
 
     def destroy(self):
-        train_and_wagon_ids = {id(obj) for obj in [self] + self.wagons}
-        for id_ in train_and_wagon_ids:
-            self.signal_controller.unreserve(id_)
+        self.signal_controller.reserve(id(self), set())
         self.notify(DestroyEvent())
 
     def is_colliding_with(self, train):
@@ -175,11 +173,7 @@ def is_colliding_with(self, train):
         )
 
     def _can_reserve_position(self, position: Vec2) -> bool:
-        train_and_wagon_ids = {id(obj) for obj in [self] + self.wagons}
-        return (
-            not self.signal_controller.reserver(position)
-            or self.signal_controller.reserver(position) in train_and_wagon_ids
-        )
+        return self.signal_controller.reserver(position) in {id(self), None}
 
     def _stop_at_target_station(self):
         # Check factories before mines, or a the iron will
@@ -249,16 +243,9 @@ def _on_reached_target(self):
         self._position_history.appendleft(Vec2(self.target_x, self.target_y))
         self._update_current_rail_and_target_xy(next_rail, self.target_x, self.target_y)
 
-        self.signal_controller.reserve(id(self), next_position)
-        for wagon in self.wagons:
-            # Kind of dangerous, an implementation change could
-            # mean that another train gets the chance to
-            # reserve the block before the wagon gets the chance
-            # to reserve it again
-            self.signal_controller.unreserve(id(wagon))
-            self.signal_controller.reserve(
-                id(wagon), self.grid.snap_to(wagon.x, wagon.y)
-            )
+        self.signal_controller.reserve(
+            id(self), [*self._position_history, next_position]
+        )
 
     def _update_current_rail_and_target_xy(self, next_rail: Rail, x, y):
         self.current_rail = next_rail