core: stdcm: memoize simulateBlock

Add computed block envelopes to STDCMGraph. Only compute block envelopes which have not been already computed.
OpenRailAssociation · Oct 11, 2023 · b4ca116 · b4ca116
1 parent 03f1386
commit b4ca116
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 14 deletions.
diff --git a/core/src/main/java/fr/sncf/osrd/stdcm/graph/BlockSimulationParameters.java b/core/src/main/java/fr/sncf/osrd/stdcm/graph/BlockSimulationParameters.java
@@ -0,0 +1,4 @@
+package fr.sncf.osrd.stdcm.graph;
+
+public record BlockSimulationParameters(Integer blockId, double initialSpeed, long start, Long stop) {
+}
diff --git a/core/src/main/java/fr/sncf/osrd/stdcm/graph/STDCMEdgeBuilder.java b/core/src/main/java/fr/sncf/osrd/stdcm/graph/STDCMEdgeBuilder.java
@@ -2,9 +2,7 @@
 
 import fr.sncf.osrd.envelope.Envelope;
 import fr.sncf.osrd.sim_infra.impl.BlockInfraImplKt;
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -164,16 +162,9 @@ STDCMEdge findEdgeSameNextOccupancy(double timeNextOccupancy) {
     private Envelope getEnvelope() {
         if (envelope == null)
             envelope = STDCMSimulations.simulateBlock(
-                    graph.rawInfra,
-                    graph.blockInfra,
-                    blockId,
-                    startSpeed,
-                    startOffset,
-                    graph.rollingStock,
-                    graph.comfort,
-                    graph.timeStep,
-                    STDCMUtils.getStopOnBlock(graph, blockId, startOffset, waypointIndex),
-                    graph.tag
+                    graph,
+                    new BlockSimulationParameters(blockId, startSpeed, startOffset,
+                            STDCMUtils.getStopOnBlock(graph, blockId, startOffset, waypointIndex))
             );
         return envelope;
     }

diff --git a/core/src/main/java/fr/sncf/osrd/stdcm/graph/STDCMGraph.java b/core/src/main/java/fr/sncf/osrd/stdcm/graph/STDCMGraph.java
@@ -13,6 +13,7 @@
 import fr.sncf.osrd.utils.graph.Graph;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 
 /** This is the class that encodes the STDCM problem as a graph on which we can run our pathfinding implementation.
@@ -29,6 +30,7 @@ public class STDCMGraph implements Graph<STDCMNode, STDCMEdge> {
     public final RollingStock rollingStock;
     public final RollingStock.Comfort comfort;
     public final double timeStep;
+    HashMap<BlockSimulationParameters, Envelope> simulatedEnvelopes;
     final List<STDCMStep> steps;
     final DelayManager delayManager;
     final AllowanceManager allowanceManager;
@@ -55,6 +57,7 @@ public STDCMGraph(
         this.rollingStock = rollingStock;
         this.comfort = comfort;
         this.timeStep = timeStep;
+        this.simulatedEnvelopes = new HashMap<>();
         this.steps = steps;
         this.delayManager = new DelayManager(minScheduleTimeStart, maxRunTime, blockAvailability, this);
         this.allowanceManager = new AllowanceManager(this);

diff --git a/core/src/main/java/fr/sncf/osrd/stdcm/graph/STDCMSimulations.java b/core/src/main/java/fr/sncf/osrd/stdcm/graph/STDCMSimulations.java
@@ -49,6 +49,29 @@ static EnvelopeSimContext makeSimContext(
         return EnvelopeSimContextBuilder.build(rollingStock, envelopePath, timeStep, comfort);
     }
 
+    /** Returns the corresponding envelope if the block's envelope has already been computed in the STDCMGraph,
+     * otherwise computes the matching envelope and adds it to the STDCMGraph. **/
+    public static Envelope simulateBlock(STDCMGraph graph, BlockSimulationParameters blockParams) {
+        if (graph.simulatedEnvelopes.containsKey(blockParams)) {
+            return graph.simulatedEnvelopes.get(blockParams);
+        } else {
+            var simulatedEnvelope = simulateBlock(
+                    graph.rawInfra,
+                    graph.blockInfra,
+                    blockParams.blockId(),
+                    blockParams.initialSpeed(),
+                    blockParams.start(),
+                    graph.rollingStock,
+                    graph.comfort,
+                    graph.timeStep,
+                    blockParams.stop(),
+                    graph.tag
+            );
+            graph.simulatedEnvelopes.putIfAbsent(blockParams, simulatedEnvelope);
+            return simulatedEnvelope;
+        }
+    }
+
     /**
      * Returns an envelope matching the given block. The envelope time starts when the train enters the block.
      * stopPosition specifies the position at which the train should stop, may be null (no stop)

diff --git a/core/src/test/java/fr/sncf/osrd/stdcm/PerformanceTests.java b/core/src/test/java/fr/sncf/osrd/stdcm/PerformanceTests.java
@@ -8,7 +8,6 @@
 import com.google.common.collect.Iterables;
 import fr.sncf.osrd.utils.DummyInfra;
 import fr.sncf.osrd.utils.graph.Pathfinding;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import java.util.ArrayList;
 import java.util.Set;
@@ -94,7 +93,6 @@ We create one very long infra (1000 consecutive blocks of 1km each).
         assertNull(res);
     }
 
-    @Disabled("We're not quite there yet")
     @Test
     public void testManyWithDifferentPaths() {
         /*