handle multiple consists per train schedule

Co-authored-by: Younes Khoudli <younes.khoudli@epita.fr>
Signed-off-by: Loup Federico <16464925+Sh099078@users.noreply.github.com>

Author: Sh099078

core/envelope-sim/src/main/java/fr/sncf/osrd/envelope/EnvelopeBuilder.kt

@@ -47,3 +47,13 @@ class EnvelopeBuilder {
         }
     }
 }
+
+fun concatenateAndShiftEnvelopes(envelopes: List<Envelope>): Envelope {
+    var currentOffset = 0.0
+    val envelopeParts = mutableListOf<EnvelopePart>()
+    for (envelope in envelopes) {
+        envelopeParts.addAll(envelope.map { it.copyAndShift(currentOffset) })
+        currentOffset += envelope.endPos
+    }
+    return Envelope(envelopeParts.toTypedArray())
+}

core/envelope-sim/src/main/java/fr/sncf/osrd/envelope/part/EnvelopePart.kt

@@ -521,16 +521,18 @@ class EnvelopePart(
      */
     fun copyAndShift(
         positionDelta: Double,
-        minPosition: Double,
-        maxPosition: Double,
+        minPosition: Double? = null,
+        maxPosition: Double? = null,
     ): EnvelopePart {
         val newPositions = DoubleArrayList()
         val newSpeeds = DoubleArrayList()
         val newTimeDeltas = DoubleArrayList()
         newPositions.add(positions[0] + positionDelta)
         newSpeeds.add(speeds[0])
         for (i in 1 until positions.size) {
-            val p = max(minPosition, min(maxPosition, positions[i] + positionDelta))
+            var p = positions[i] + positionDelta
+            if (minPosition != null) p = max(minPosition, p)
+            if (maxPosition != null) p = min(maxPosition, p)
             if (newPositions.last().value != p) {
                 // Positions that are an epsilon away may be overlapping after the shift, we only
                 // add the distinct ones

core/src/main/java/fr/sncf/osrd/cli/BenchSTDCM.kt

@@ -175,24 +175,39 @@ class BenchSTDCM : CliCommand {
                 try {
                     val temporarySpeedLimitManager =
                         buildTemporarySpeedLimitManager(infra, request.temporarySpeedLimits)
-                    val rollingStock =
-                        parseRawRollingStock(
-                            request.consistSchedule.values[0].physicsConsist,
-                            request.consistSchedule.values[0].loadingGaugeType,
-                            request.consistSchedule.values[0].supportedSignalingSystems.filter {
-                                // Ignoring ETCS as it is not (yet) supported for STDCM
-                                it != ETCS_LEVEL2.id
-                            },
+                    val consistConfigurations =
+                        request.consistSchedule.values.map { it ->
+                            it.copy(
+                                supportedSignalingSystems =
+                                    it.supportedSignalingSystems.filter {
+                                        // Ignoring ETCS as it is not (yet) supported for
+                                        it != ETCS_LEVEL2.id
+                                    }
+                            )
+                        }
+                    val requestConsistSchedule =
+                        request.consistSchedule.copy(values = consistConfigurations)
+                    val allowedTrackSections =
+                        parseTrackSectionIds(infra, request.allowedTrackSections)
+                    val consistSchedules =
+                        ConsistSchedule(
+                            requestConsistSchedule,
+                            infra,
+                            allowedTrackSections,
+                            request.pathItems.size,
                         )
                     val steps =
-                        parseSteps(infra, request.pathItems, request.startTime, rollingStock.length)
+                        parseSteps(
+                            infra,
+                            request.pathItems,
+                            request.startTime,
+                            consistSchedules.rollingStocks.map { it.length },
+                        )
                     val requirements = getRequirements(request, infra, cacheManager)
-                    val allowedTrackSections =
-                        parseTrackSectionIds(infra, request.allowedTrackSections)
                     path =
                         findPath(
                             infra,
-                            rollingStock,
+                            consistSchedules,
                             request.comfort,
                             0.0,
                             steps,

core/src/main/kotlin/fr/sncf/osrd/api/CommonSchemas.kt

@@ -6,8 +6,10 @@ import fr.sncf.osrd.railjson.schema.common.graph.EdgeDirection
 import fr.sncf.osrd.sim_infra.api.TrackSection
 import fr.sncf.osrd.utils.DistanceRangeMap
 import fr.sncf.osrd.utils.distanceRangeMapOf
+import fr.sncf.osrd.utils.units.Distance
 import fr.sncf.osrd.utils.units.Offset
 import fr.sncf.osrd.utils.units.TimeDelta
+import fr.sncf.osrd.utils.units.meters
 
 data class DirectionalTrackRange(
     @Json(name = "track_section") val trackSection: String,
@@ -51,6 +53,16 @@ data class RangeValues<valueT>(
         }
         return distanceRangeMapOf(rangeMapEntries)
     }
+
+    fun shifted(distance: Distance): RangeValues<valueT> {
+        return this.copy(
+            internalBoundaries =
+                this.internalBoundaries.map {
+                    require(it + distance > Offset(0.meters))
+                    it + distance
+                }
+        )
+    }
 }
 
 class TrackLocation(val track: String, val offset: Offset<TrackSection>)

core/src/main/kotlin/fr/sncf/osrd/api/RollingStockParser.kt

@@ -1,16 +1,23 @@
 package fr.sncf.osrd.api
 
 import fr.sncf.osrd.api.standalone_sim.PhysicsConsistModel
+import fr.sncf.osrd.api.stdcm.RequestConsistSchedule
+import fr.sncf.osrd.envelope_sim.PhysicsRollingStock
 import fr.sncf.osrd.envelope_sim.PhysicsRollingStock.TractiveEffortPoint
 import fr.sncf.osrd.envelope_sim.etcs.toEtcsBrakeParams
+import fr.sncf.osrd.graph.PathfindingConstraint
+import fr.sncf.osrd.pathfinding.constraints.CachedBlockConstraintCombiner
+import fr.sncf.osrd.pathfinding.constraints.initConstraints
 import fr.sncf.osrd.railjson.schema.rollingstock.RJSEffortCurves.*
 import fr.sncf.osrd.railjson.schema.rollingstock.RJSLoadingGaugeType
 import fr.sncf.osrd.railjson.schema.rollingstock.RJSRollingResistance
 import fr.sncf.osrd.railjson.schema.rollingstock.RJSRollingResistance.Davis
 import fr.sncf.osrd.reporting.exceptions.ErrorType
 import fr.sncf.osrd.reporting.exceptions.OSRDError
+import fr.sncf.osrd.sim_infra.api.TrackSectionId
 import fr.sncf.osrd.train.RollingStock
 import fr.sncf.osrd.train.RollingStock.*
+import kotlin.collections.get
 
 /** Parse the rolling stock model into something the backend can work with */
 fun parseRawRollingStock(
@@ -60,6 +67,101 @@ fun parseRawRollingStock(
     )
 }
 
+/**
+ * Associates a list of rolling stocks with their related pathfinding constraints. This class
+ * provides two ways to be built:
+ * - From a list of STDCM query inputs.
+ * - From a list of rolling stocks and their pathfinding constraints. This approach is mostly useful
+ *   for testing purposes.
+ */
+data class ConsistSchedule(
+    val rollingStocks: List<PhysicsRollingStock>,
+    val constraints: List<PathfindingConstraint>?,
+) {
+    init {
+        require(!rollingStocks.isEmpty())
+        require(constraints == null || rollingStocks.size == constraints.size)
+    }
+
+    companion object {
+        operator fun invoke(
+            consistSchedule: RequestConsistSchedule,
+            infra: FullInfra,
+            allowedTrackSections: Set<TrackSectionId>? = null,
+            totalSteps: Int,
+        ): ConsistSchedule {
+            val boundaries = consistSchedule.boundaries
+            val rollingStocks =
+                consistSchedule.values.map {
+                    parseRawRollingStock(
+                        it.physicsConsist,
+                        it.loadingGaugeType,
+                        it.supportedSignalingSystems,
+                    )
+                }
+            return ConsistSchedule(
+                rollingStocks,
+                boundaries,
+                infra,
+                allowedTrackSections,
+                totalSteps,
+            )
+        }
+
+        operator fun invoke(
+            rollingStocks: List<RollingStock>,
+            boundaries: List<Int>,
+            infra: FullInfra,
+            allowedTrackSections: Set<TrackSectionId>? = null,
+            totalSteps: Int,
+        ): ConsistSchedule {
+            // Input validation:
+            when {
+                (rollingStocks.size != boundaries.size + 1) -> {
+                    throw OSRDError(ErrorType.InvalidSTDCMInputs)
+                        .withContext(
+                            "cause",
+                            "${boundaries.size} boundaries and ${rollingStocks.size} consist configurations provided. There should be n-1 boundaries for n consist configurations",
+                        )
+                }
+                (!boundaries.zipWithNext().all { (a, b) -> a < b }) -> {
+                    throw OSRDError(ErrorType.InvalidSTDCMInputs)
+                        .withContext(
+                            "cause",
+                            "Consist change boundaries are not strictly increasing",
+                        )
+                }
+                (!(boundaries.isEmpty() ||
+                    (boundaries.first() != 0 && boundaries.last() != totalSteps - 1))) -> {
+                    throw OSRDError(ErrorType.InvalidSTDCMInputs)
+                        .withContext(
+                            "cause",
+                            "Consist change specified on the first or last step of the path",
+                        )
+                }
+            }
+
+            // Build the rolling stock and constraint for each step:
+            val rollingStocksPerStep = mutableListOf<RollingStock>()
+            val constraints = mutableListOf<PathfindingConstraint>()
+            var previousBoundary = 0
+            for ((index, rollingStock) in rollingStocks.withIndex()) {
+                val boundary = boundaries.getOrNull(index) ?: totalSteps
+                val constraint =
+                    CachedBlockConstraintCombiner(
+                        initConstraints(infra, rollingStock, allowedTrackSections)
+                    )
+                (previousBoundary..<boundary).forEach { _ ->
+                    rollingStocksPerStep.add(rollingStock)
+                    constraints.add(constraint)
+                }
+                previousBoundary = boundary
+            }
+            return ConsistSchedule(rollingStocksPerStep, constraints)
+        }
+    }
+}
+
 private fun parseRollingResistance(rjsRollingResistance: RJSRollingResistance?): Davis {
     if (rjsRollingResistance == null)
         throw OSRDError.newMissingRollingStockFieldError("rolling_resistance")

core/src/main/kotlin/fr/sncf/osrd/api/standalone_sim/SimulationRequest.kt

@@ -58,7 +58,7 @@ enum class AllowanceDistribution {
     }
 }
 
-class PhysicsConsistModel(
+data class PhysicsConsistModel(
     @Json(name = "effort_curves") val effortCurves: EffortCurve,
     @Json(name = "base_power_class") val basePowerClass: String?,
     val length: Length<RollingStock>,