core: stdcm: cache engineering allowance sequences

Signed-off-by: Eloi Charpentier <[email protected]>

Author: eckter

core/kt-osrd-utils/src/main/kotlin/fr/sncf/osrd/utils/SoftLazy.kt

@@ -0,0 +1,20 @@
+package fr.sncf.osrd.utils
+
+import java.lang.ref.SoftReference
+import kotlin.properties.ReadOnlyProperty
+import kotlin.reflect.KProperty
+
+/**
+ * Usage: `val myValue by SoftLazy { computeValue() }`. The value is evaluated only when it's
+ * actually accessed. It's then stored in a `SoftReference`, which may be cleared if the JVM needs
+ * more RAM. Variables defined that way can then be used transparently, as if they were of type `T`.
+ */
+class SoftLazy<T>(
+    private val computeValue: () -> T,
+) : ReadOnlyProperty<Any?, T> {
+    var cache: SoftReference<T> = SoftReference(null)
+
+    override fun getValue(thisRef: Any?, property: KProperty<*>): T {
+        return cache.get() ?: computeValue().also { cache = SoftReference(it) }
+    }
+}

core/src/main/kotlin/fr/sncf/osrd/stdcm/graph/STDCMEdge.kt

@@ -79,6 +79,7 @@ data class STDCMEdge(
                 null,
                 previousPlannedNodeRelativeTimeDiff,
                 graph.remainingTimeEstimator.invoke(this, null, stepTracker),
+                graph,
             )
         } else {
             // New edge on the same block, after a stop
@@ -103,6 +104,7 @@ data class STDCMEdge(
                 nextStop.originalStep.plannedTimingData,
                 previousPlannedNodeRelativeTimeDiff,
                 graph.remainingTimeEstimator.invoke(this, locationOnEdge, stepTracker),
+                graph,
             )
         }
     }

core/src/main/kotlin/fr/sncf/osrd/stdcm/graph/STDCMNode.kt

@@ -3,7 +3,10 @@ package fr.sncf.osrd.stdcm.graph
 import fr.sncf.osrd.envelope_sim.TrainPhysicsIntegrator.areTimesEqual
 import fr.sncf.osrd.sim_infra.api.Block
 import fr.sncf.osrd.stdcm.PlannedTimingData
+import fr.sncf.osrd.stdcm.graph.engineering_allowance.generatePreviousSimulationSegments
 import fr.sncf.osrd.stdcm.infra_exploration.InfraExplorerWithEnvelope
+import fr.sncf.osrd.utils.SoftLazy
+import fr.sncf.osrd.utils.cacheable
 import fr.sncf.osrd.utils.units.Offset
 import kotlin.math.min
 
@@ -27,6 +30,8 @@ data class STDCMNode(
     val previousPlannedNodeRelativeTimeDiff: Double?,
     // Estimation of the min time it takes to reach the end from this node
     var remainingTimeEstimation: Double,
+    // Reference to the main graph. Only null in some unit tests, where we don't have a full graph
+    val graph: STDCMGraph?,
 ) : Comparable<STDCMNode> {
 
     /**
@@ -183,4 +188,14 @@ data class STDCMNode(
     fun getRealTime(updatedTimeData: TimeData): Double {
         return timeData.getUpdatedEarliestReachableTime(updatedTimeData)
     }
+
+    val previousSimulationSegments by SoftLazy {
+        generatePreviousSimulationSegments(previousEdge, graph).cacheable()
+    }
+    val allowanceOpportunities by SoftLazy {
+        graph!!
+            .allowanceManager
+            .generateAllowanceOpportunities(previousSimulationSegments, speed)
+            .cacheable()
+    }
 }

core/src/main/kotlin/fr/sncf/osrd/stdcm/graph/STDCMPathfinding.kt

@@ -117,7 +117,7 @@ class STDCMPathfinding(
             ConstraintCombiner(initConstraints(fullInfra, listOf(rollingStock)).toMutableList())
 
         assert(steps.last().stop) { "The last stop is supposed to be an actual stop" }
-        starts = getStartNodes(listOf(constraints))
+        starts = getStartNodes(graph, listOf(constraints))
         val path = findPathImpl()
         graph.stdcmSimulations.logWarnings()
         if (path == null) {
@@ -258,7 +258,10 @@ class STDCMPathfinding(
     }
 
     /** Converts start locations into starting nodes. */
-    private fun getStartNodes(constraints: List<PathfindingConstraint<Block>>): Set<STDCMNode> {
+    private fun getStartNodes(
+        graph: STDCMGraph,
+        constraints: List<PathfindingConstraint<Block>>
+    ): Set<STDCMNode> {
         val res = HashSet<STDCMNode>()
         val firstStep = steps[0]
         assert(!firstStep.stop)
@@ -292,6 +295,7 @@ class STDCMPathfinding(
                         firstStep.plannedTimingData,
                         null,
                         graph.bestPossibleTime,
+                        graph,
                     )
                 res.add(node)
             }

core/src/main/kotlin/fr/sncf/osrd/stdcm/graph/engineering_allowance/EngineeringAllowanceManager.kt

@@ -41,16 +41,17 @@ class EngineeringAllowanceManager(
         assert(constDeceleration > 0.0)
 
         val requiredAdditionalTime = expectedStartTime - prevNode.timeData.earliestReachableTime
-        val segments = generatePreviousSimulationSegments(prevNode.previousEdge, graph).cacheable()
-        val opportunities = generateAllowanceOpportunities(segments, prevNode.speed)
+
+        val opportunities = prevNode.allowanceOpportunities
+
         val solution =
             opportunities
                 .takeWhile { it.maxNextAllowanceValue >= requiredAdditionalTime }
                 .firstOrNull { it.addedTime >= requiredAdditionalTime } ?: return null
 
         // We still try to roughly guess the maximum length over which to add the time
         var distance = 0.meters
-        for (segment in segments) {
+        for (segment in prevNode.previousSimulationSegments) {
             if (segment.maxAddedDelay <= requiredAdditionalTime) break
             distance += segment.length
         }

core/src/test/kotlin/fr/sncf/osrd/stdcm/preprocessing/STDCMHeuristicTests.kt

@@ -248,6 +248,7 @@ class STDCMHeuristicTests {
                 null,
                 null,
                 0.0,
+                null
             )
         val defaultEdge =
             STDCMEdge(