fixup! editoast: filter train schedules for LMR

Signed-off-by: Youness CHRIFI ALAOUI <[email protected]>

editoast/src/models/timetable.rs

@@ -5,9 +5,12 @@ use diesel::sql_query;
 use diesel::sql_types::Array;
 use diesel::sql_types::BigInt;
 use diesel_async::RunQueryDsl;
+use futures_util::stream::TryStreamExt;
 use std::ops::DerefMut;
 
 use crate::error::Result;
+use crate::models::prelude::*;
+use crate::models::train_schedule::TrainSchedule;
 use crate::models::Identifiable;
 use crate::models::{DeleteStatic, Retrieve};
 use crate::Exists;
@@ -56,6 +59,29 @@ impl Timetable {
             .await
             .map_err(Into::into)
     }
+
+    pub async fn schedules_before_date(
+        self,
+        conn: &mut DbConnection,
+        time: DateTime<Utc>,
+    ) -> Result<Vec<TrainSchedule>> {
+        use diesel::prelude::*;
+        use diesel_async::RunQueryDsl;
+        use editoast_models::tables::train_schedule::dsl;
+
+        let train_schedules = dsl::train_schedule
+            .filter(dsl::start_time.le(time))
+            .filter(dsl::timetable_id.eq(self.id))
+            .load_stream::<Row<TrainSchedule>>(conn.write().await.deref_mut())
+            .await?
+            .map_ok(|ts| ts.into())
+            .try_collect::<Vec<TrainSchedule>>()
+            .await;
+        match train_schedules {
+            Ok(train_schedules) => Ok(train_schedules),
+            Err(err) => Err(err.into()),
+        }
+    }
 }
 
 #[async_trait::async_trait]

editoast/src/models/train_schedule.rs

@@ -1,10 +1,6 @@
-use crate::error::Result;
 use chrono::DateTime;
 use chrono::Utc;
-use diesel::prelude::*;
-use diesel_async::RunQueryDsl;
 use editoast_derive::Model;
-use editoast_models::DbConnection;
 use editoast_schemas::train_schedule::Comfort;
 use editoast_schemas::train_schedule::Distribution;
 use editoast_schemas::train_schedule::Margins;
@@ -13,12 +9,8 @@ use editoast_schemas::train_schedule::PowerRestrictionItem;
 use editoast_schemas::train_schedule::ScheduleItem;
 use editoast_schemas::train_schedule::TrainScheduleBase;
 use editoast_schemas::train_schedule::TrainScheduleOptions;
-use futures_util::stream::TryStreamExt;
-use std::ops::DerefMut;
 
-use super::Model as _;
-use crate::diesel::ExpressionMethods;
-use crate::Row;
+use crate::models::prelude::*;
 
 #[derive(Debug, Default, Clone, Model)]
 #[model(table = editoast_models::tables::train_schedule)]
@@ -82,23 +74,3 @@ impl From<TrainScheduleBase> for TrainScheduleChangeset {
             .options(options)
     }
 }
-
-pub async fn filter_train_by_start_time_and_timetable(
-    conn: &mut DbConnection,
-    time: DateTime<Utc>,
-    timetable_id: i64,
-) -> Result<Vec<TrainSchedule>> {
-    use editoast_models::tables::train_schedule::dsl;
-    let train_schedules = dsl::train_schedule
-        .filter(dsl::start_time.le(time))
-        .filter(dsl::timetable_id.eq(timetable_id))
-        .load_stream::<Row<TrainSchedule>>(conn.write().await.deref_mut())
-        .await?
-        .map_ok(|ts| ts.into())
-        .try_collect::<Vec<TrainSchedule>>()
-        .await;
-    match train_schedules {
-        Ok(train_schedules) => Ok(train_schedules),
-        Err(err) => Err(err.into()),
-    }
-}

editoast/src/views/timetable/stdcm.rs

@@ -1,7 +1,6 @@
 mod failure_handler;
 mod request;
 
-use crate::models::train_schedule::filter_train_by_start_time_and_timetable;
 use axum::extract::Json;
 use axum::extract::Path;
 use axum::extract::Query;
@@ -148,7 +147,7 @@ async fn stdcm(
     let timetable_id = id;
     let infra_id = query.infra;
 
-    // 1. Retrieve Timetable / Infra / Trains / Simulation / Rolling Stock
+    // 1.  Infra / Timetable / Trains / Simulation / Rolling Stock
 
     let infra = Infra::retrieve_or_fail(&mut conn, infra_id, || StdcmError::InfraNotFound {
         infra_id,
@@ -189,25 +188,64 @@ async fn stdcm(
         StdcmError::TimetableNotFound { timetable_id }
     })
     .await?;
-    // Filter train
-    let mut train_schedules =
-        filter_train_by_start_time_and_timetable(&mut conn, latest_simulation_end, timetable.id)
-            .await?;
+
+    // Filter trains
+    // The goal is to filter out as many trains as possible whose schedules overlap
+    // with the LMR train being searched for.
+    // The diagram below shows an LMR train inserted into a timetable.
+
+    // '?': unscheduled arrival times.
+    // '|': scheduled arrival times.
+    // tA: earliest_departure_time
+    // tB: latest_simulation_end
+    //
+    //                           tA                     tB
+    //       LMR Train           |----------------------|
+    // Train 1            |--------------|
+    // Train 2      |------------|
+    //                                        |----------?   Train 3
+    // Train 4        |-------?
+    //                         Train 5  |---------?
+    //                                       |----------?   Train 6
+
+    // Step 1 (SQL Filter):
+    // Trains that depart after the latest arrival time of the LMR train are excluded.
+    // In this example, Train 3 and Train 6 are filtered out.
+
+    // It's not easy to write an SQL query to filter trains when the train departure time < latest_simulation_ended
+    // because there are two cases : when the train departure time > tA (Step 2) and the train departure time < tA (Step 3).
+
+    // Step 2 (Rust filter) :
+    // If the train departure time > LMR train departure (tA), the train is kept (e.g., train_5)
+    // Step 3 (Rust filter) :
+    // For trains departing before the LMR train departure (tA):
+
+    // If the train's arrival time is unscheduled (?), the train is kept (e.g., Train 4 and Train 5).
+    // If the train's arrival time is scheduled (|), the train is kept only if its arrival time is after the LMR train's earliest departure time.
+    // Train 1 is kept and train 2 is filtered out.
+
+    // Step 1
+    let mut train_schedules = timetable
+        .schedules_before_date(&mut conn, latest_simulation_end)
+        .await?;
 
     train_schedules.retain(|train_schedule| {
-        if train_schedule.start_time >= earliest_departure_time {
-            true
-        } else {
-            if let Some(last_item) = train_schedule.schedule.last() {
-                if let Some(last_path_item) = train_schedule.path.last() {
-                    last_item.at == last_path_item.id
-                } else {
-                    false
-                }
-            } else {
-                false
-            }
-        }
+        // Step 2 and 3
+        train_schedule.start_time >= earliest_departure_time
+            || train_schedule
+                .schedule
+                .last()
+                .and_then(|last_schedule_item| {
+                    train_schedule.path.last().and_then(|last_path_item| {
+                        (last_schedule_item.at == last_path_item.id).then_some(last_schedule_item)
+                    })
+                })
+                .and_then(|last_schedule_item| {
+                    last_schedule_item.arrival.clone().map(|arrival| {
+                        train_schedule.start_time + *arrival > earliest_departure_time
+                    })
+                })
+                .unwrap_or(true)
     });
 
     // 3. Get scheduled train requirements