editoast: move map_to_core_work_schedule to WorkSchedule

Signed-off-by: hamz2a <[email protected]>

Author: hamz2a

editoast/src/models/work_schedules.rs

@@ -1,4 +1,9 @@
+use std::cmp::max;
+
+use chrono::DateTime;
 use chrono::NaiveDateTime;
+use chrono::TimeZone;
+use chrono::Utc;
 use editoast_derive::Model;
 use editoast_schemas::infra::TrackRange;
 use strum::FromRepr;
@@ -7,6 +12,8 @@ use serde::Deserialize;
 use serde::Serialize;
 use utoipa::ToSchema;
 
+use crate::core::stdcm::UndirectedTrackRange;
+
 #[derive(Debug, Clone, Model)]
 #[model(table = editoast_models::tables::work_schedule_group)]
 #[model(gen(ops = crd, batch_ops = c, list))]
@@ -38,3 +45,46 @@ pub struct WorkSchedule {
     pub work_schedule_type: WorkScheduleType,
     pub work_schedule_group_id: i64,
 }
+
+impl WorkSchedule {
+    pub fn map_to_core_work_schedule(
+        &self,
+        start_time: DateTime<Utc>,
+    ) -> crate::core::stdcm::WorkSchedule {
+        crate::core::stdcm::WorkSchedule {
+            start_time: elapsed_since_time_ms(&self.start_date_time, &start_time),
+            end_time: elapsed_since_time_ms(&self.end_date_time, &start_time),
+            track_ranges: self
+                .track_ranges
+                .iter()
+                .map(|track| UndirectedTrackRange {
+                    track_section: track.track.to_string(),
+                    begin: (track.begin * 1000.0) as u64,
+                    end: (track.end * 1000.0) as u64,
+                })
+                .collect(),
+        }
+    }
+
+    pub fn make_stdcm_work_schedule(
+        &self,
+        start_time: DateTime<Utc>,
+        latest_simulation_end: DateTime<Utc>,
+    ) -> Option<crate::core::stdcm::WorkSchedule> {
+        let search_window_duration = (latest_simulation_end - start_time).num_milliseconds() as u64;
+
+        let ws = self.map_to_core_work_schedule(start_time);
+        if ws.end_time > 0 && ws.start_time < search_window_duration {
+            Some(ws)
+        } else {
+            None
+        }
+    }
+}
+
+fn elapsed_since_time_ms(time: &NaiveDateTime, start_time: &DateTime<Utc>) -> u64 {
+    max(
+        0,
+        (Utc.from_utc_datetime(time) - start_time).num_milliseconds(),
+    ) as u64
+}

editoast/src/views/timetable.rs

@@ -2,7 +2,6 @@ pub mod path_not_found_handler;
 pub mod stdcm;
 pub mod stdcm_request_payload;
 
-use std::cmp::max;
 use std::collections::HashMap;
 
 use axum::extract::Json;
@@ -12,10 +11,6 @@ use axum::extract::State;
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::Extension;
-use chrono::DateTime;
-use chrono::NaiveDateTime;
-use chrono::TimeZone;
-use chrono::Utc;
 use derivative::Derivative;
 use editoast_authz::BuiltinRole;
 use editoast_derive::EditoastError;
@@ -31,15 +26,13 @@ use crate::core::conflict_detection::Conflict;
 use crate::core::conflict_detection::ConflictDetectionRequest;
 use crate::core::conflict_detection::TrainRequirements;
 use crate::core::simulation::SimulationResponse;
-use crate::core::stdcm::UndirectedTrackRange;
 use crate::core::AsCoreRequest;
 use crate::error::Result;
 use crate::models::prelude::*;
 use crate::models::timetable::Timetable;
 use crate::models::timetable::TimetableWithTrains;
 use crate::models::train_schedule::TrainSchedule;
 use crate::models::train_schedule::TrainScheduleChangeset;
-use crate::models::work_schedules::WorkSchedule;
 use crate::models::Infra;
 use crate::views::train_schedule::train_simulation_batch;
 use crate::views::train_schedule::TrainScheduleForm;
@@ -359,29 +352,6 @@ async fn conflicts(
     Ok(Json(conflict_detection_response.conflicts))
 }
 
-pub fn map_to_core_work_schedule(
-    ws: &WorkSchedule,
-    start_time: DateTime<Utc>,
-) -> crate::core::stdcm::WorkSchedule {
-    crate::core::stdcm::WorkSchedule {
-        start_time: elapsed_since_time_ms(&ws.start_date_time, &start_time),
-        end_time: elapsed_since_time_ms(&ws.end_date_time, &start_time),
-        track_ranges: ws
-            .track_ranges
-            .iter()
-            .map(|track| UndirectedTrackRange {
-                track_section: track.track.to_string(),
-                begin: (track.begin * 1000.0) as u64,
-                end: (track.end * 1000.0) as u64,
-            })
-            .collect(),
-    }
-}
-
-fn elapsed_since_time_ms(time: &NaiveDateTime, zero: &DateTime<Utc>) -> u64 {
-    max(0, (Utc.from_utc_datetime(time) - zero).num_milliseconds()) as u64
-}
-
 #[cfg(test)]
 mod tests {
     use axum::http::StatusCode;

editoast/src/views/timetable/path_not_found_handler.rs

@@ -14,7 +14,6 @@ use crate::models::work_schedules::WorkSchedule;
 use crate::views::path::pathfinding::PathfindingResult;
 use crate::views::timetable::stdcm::STDCMResponse;
 
-use super::map_to_core_work_schedule;
 use super::stdcm::build_train_requirements;
 
 pub struct PathNotFoundHandler {
@@ -108,7 +107,7 @@ fn make_work_schedules_request(
 
     let work_schedule_requirements = work_schedules
         .iter()
-        .map(|ws| (ws.id, map_to_core_work_schedule(ws, start_time)))
+        .map(|ws| (ws.id, ws.map_to_core_work_schedule(start_time)))
         .filter(|(_, ws)| ws.end_time > 0 && ws.start_time < search_window_duration)
         .collect();
 

editoast/src/views/timetable/stdcm.rs

@@ -23,7 +23,6 @@ use thiserror::Error;
 use utoipa::IntoParams;
 use utoipa::ToSchema;
 
-use super::map_to_core_work_schedule;
 use super::path_not_found_handler::PathNotFoundHandler;
 use super::stdcm_request_payload::convert_steps;
 use super::stdcm_request_payload::STDCMRequestPayload;
@@ -258,11 +257,12 @@ async fn stdcm(
         time_gap_after: stdcm_request.time_gap_after,
         margin: stdcm_request.margin,
         time_step: Some(2000),
-        work_schedules: filter_stdcm_work_schedules(
-            &work_schedules,
-            earliest_departure_time,
-            latest_simulation_end,
-        ),
+        work_schedules: work_schedules
+            .iter()
+            .filter_map(|ws| {
+                ws.make_stdcm_work_schedule(earliest_departure_time, latest_simulation_end)
+            })
+            .collect(),
         temporary_speed_limits,
         rolling_stock: PhysicsRollingStock::new(rolling_stock.into(), simulation_parameters),
     };
@@ -458,19 +458,6 @@ fn build_single_margin(margin: Option<MarginValue>) -> Margins {
     }
 }
 
-fn filter_stdcm_work_schedules(
-    work_schedules: &[WorkSchedule],
-    start_time: DateTime<Utc>,
-    latest_simulation_end: DateTime<Utc>,
-) -> Vec<crate::core::stdcm::WorkSchedule> {
-    let search_window_duration = (latest_simulation_end - start_time).num_milliseconds() as u64;
-    work_schedules
-        .iter()
-        .map(|ws| map_to_core_work_schedule(ws, start_time))
-        .filter(|ws| ws.end_time > 0 && ws.start_time < search_window_duration)
-        .collect()
-}
-
 /// Return the list of speed limits that are active at any point in a given time range
 async fn build_temporary_speed_limits(
     conn: &mut DbConnection,
@@ -863,8 +850,10 @@ mod tests {
             .to_utc();
 
         // WHEN
-        let filtered =
-            filter_stdcm_work_schedules(&work_schedules, start_time, latest_simulation_end);
+        let filtered: Vec<_> = work_schedules
+            .iter()
+            .filter_map(|ws| ws.make_stdcm_work_schedule(start_time, latest_simulation_end))
+            .collect();
 
         // THEN
         assert!(filtered.is_empty() == filtered_out);