core: add a new updateMap() function to DistanceRangeMap

The old own is now called updateMapIntersection(). The new one inserts the non overlapping ranges without modification. Signed-off-by: Younes Khoudli <[email protected]>
OpenRailAssociation · Nov 28, 2024 · 25026ac · 25026ac
1 parent deac2f4
commit 25026ac
Show file tree

Hide file tree

Showing 5 changed files with 211 additions and 6 deletions.
diff --git a/core/kt-osrd-utils/src/main/kotlin/fr/sncf/osrd/utils/DistanceRangeMap.kt b/core/kt-osrd-utils/src/main/kotlin/fr/sncf/osrd/utils/DistanceRangeMap.kt
@@ -52,12 +52,25 @@ interface DistanceRangeMap<T> : Iterable<DistanceRangeMap.RangeMapEntry<T>> {
 
     /**
      * Updates the map with another one, using a merge function to fuse the values of intersecting
-     * ranges
+     * ranges. Doesn't keep any range from update where there is no intersection.
      */
-    fun <U> updateMap(update: DistanceRangeMap<U>, updateFunction: BiFunction<T, U, T>)
+    fun <U> updateMapIntersection(update: DistanceRangeMap<U>, updateFunction: BiFunction<T, U, T>)
+
+    /**
+     * Updates the map with another one, using a merge function to fuse the values of intersecting
+     * ranges. Calls default on the values of the ranges from update where there is no intersection.
+     */
+    fun updateMap(
+        update: DistanceRangeMap<T>,
+        updateFunction: (T, T) -> T,
+        default: (T) -> T = { it }
+    )
 
     /** Returns true if there is no entry at all */
     fun isEmpty(): Boolean
+
+    /** Clear the map */
+    fun clear()
 }
 
 fun <T> distanceRangeMapOf(

diff --git a/core/kt-osrd-utils/src/main/kotlin/fr/sncf/osrd/utils/DistanceRangeMapImpl.kt b/core/kt-osrd-utils/src/main/kotlin/fr/sncf/osrd/utils/DistanceRangeMapImpl.kt
@@ -190,14 +190,115 @@ data class DistanceRangeMapImpl<T>(
         return bounds.isEmpty()
     }
 
-    override fun <U> updateMap(update: DistanceRangeMap<U>, updateFunction: BiFunction<T, U, T>) {
+    override fun <U> updateMapIntersection(
+        update: DistanceRangeMap<U>,
+        updateFunction: BiFunction<T, U, T>
+    ) {
         for ((updateLower, updateUpper, updateValue) in update) {
             for ((subMapLower, subMapUpper, subMapValue) in this.subMap(updateLower, updateUpper)) {
                 this.put(subMapLower, subMapUpper, updateFunction.apply(subMapValue, updateValue))
             }
         }
     }
 
+    override fun updateMap(
+        update: DistanceRangeMap<T>,
+        updateFunction: (T, T) -> T,
+        default: (T) -> T
+    ) {
+        val resultEntries = mutableListOf<DistanceRangeMap.RangeMapEntry<T>>()
+
+        // Iterate over each range in the update map
+        for ((updateLower, updateUpper, updateValue) in update) {
+            val subMap = this.subMap(updateLower, updateUpper)
+
+            // Track segments that are part of the updated map but not intersecting with `this`
+            if (subMap.isEmpty()) {
+                resultEntries.add(
+                    DistanceRangeMap.RangeMapEntry(updateLower, updateUpper, default(updateValue))
+                )
+            } else {
+                // Handle intersections
+                var lastEnd = updateLower
+
+                for ((subMapLower, subMapUpper, subMapValue) in subMap) {
+                    // Add non-overlapping segment from `update` before the intersection
+                    if (lastEnd < subMapLower) {
+                        resultEntries.add(
+                            DistanceRangeMap.RangeMapEntry(
+                                lastEnd,
+                                subMapLower,
+                                default(updateValue)
+                            )
+                        )
+                    }
+
+                    // Add the overlapping segment with combined value
+                    val combinedValue = updateFunction(subMapValue, updateValue)
+                    resultEntries.add(
+                        DistanceRangeMap.RangeMapEntry(subMapLower, subMapUpper, combinedValue)
+                    )
+
+                    lastEnd = subMapUpper
+                }
+
+                // Add non-overlapping segment from `update` after the intersection
+                if (lastEnd < updateUpper) {
+                    val combinedValue = default(updateValue)
+                    resultEntries.add(
+                        DistanceRangeMap.RangeMapEntry(lastEnd, updateUpper, combinedValue)
+                    )
+                }
+            }
+        }
+
+        // Add non-overlapping segments from `this`
+        // It's the same thing except we never add the overlapping segments
+        // Iterate over each range in the update map
+        for ((thisLower, thisUpper, thisValue) in this) {
+            val subMap = update.subMap(thisLower, thisUpper)
+
+            // Track segments that are part of the updated map but not intersecting with `this`
+            if (subMap.isEmpty()) {
+                val combinedValue = default(thisValue)
+                resultEntries.add(
+                    DistanceRangeMap.RangeMapEntry(thisLower, thisUpper, combinedValue)
+                )
+            } else {
+                // Handle intersections
+                var lastEnd = thisLower
+
+                for ((subMapLower, subMapUpper, subMapValue) in subMap) {
+                    // Add non-overlapping segment from `this` before the intersection
+                    if (lastEnd < subMapLower) {
+                        val combinedValue = default(thisValue)
+                        resultEntries.add(
+                            DistanceRangeMap.RangeMapEntry(lastEnd, subMapLower, combinedValue)
+                        )
+                    }
+                    lastEnd = subMapUpper
+                }
+
+                // Add non-overlapping segment from `this` after the intersection
+                if (lastEnd < thisUpper) {
+                    val combinedValue = default(thisValue)
+                    resultEntries.add(
+                        DistanceRangeMap.RangeMapEntry(lastEnd, thisUpper, combinedValue)
+                    )
+                }
+            }
+        }
+
+        // Clear current map and insert the updated entries
+        this.clear()
+        this.putMany(resultEntries)
+    }
+
+    override fun clear() {
+        bounds.clear()
+        values.clear()
+    }
+
     /** Merges adjacent values, removes 0-length ranges */
     private fun mergeAdjacent() {
         fun remove(i: Int) {

diff --git a/core/kt-osrd-utils/src/test/kotlin/fr/sncf/osrd/utils/TestDistanceRangeMap.kt b/core/kt-osrd-utils/src/test/kotlin/fr/sncf/osrd/utils/TestDistanceRangeMap.kt
@@ -1,9 +1,11 @@
 package fr.sncf.osrd.utils
 
 import fr.sncf.osrd.utils.units.Distance
+import fr.sncf.osrd.utils.units.meters
 import kotlin.test.assertEquals
 import kotlin.time.*
 import kotlin.time.Duration.Companion.seconds
+import org.junit.Assert.*
 import org.junit.Test
 
 class TestDistanceRangeMap {
@@ -310,4 +312,91 @@ class TestDistanceRangeMap {
         assert(!mark2.hasPassedNow())
         assertEquals(mergedEntries, rangeMap.asList())
     }
+
+    @Test
+    fun updateMapIntersection() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 10.0.meters, "A")
+        val updateMap = DistanceRangeMapImpl<String>()
+        updateMap.put(5.0.meters, 15.0.meters, "B")
+        map.updateMapIntersection(updateMap) { old, new -> old + new }
+        assertEquals("AB", map.get(7.5.meters))
+        assertEquals("A", map.get(2.5.meters))
+        assertNull(map.get(12.5.meters))
+    }
+
+    @Test
+    fun updateMap_noOverlap() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 5.0.meters, "A")
+        val update = DistanceRangeMapImpl<String>()
+        update.put(10.0.meters, 15.0.meters, "B")
+        map.updateMap(update, { old, new -> old + new })
+        assertEquals("A", map.get(2.5.meters))
+        assertEquals("B", map.get(12.5.meters))
+        assertNull(map.get(7.5.meters))
+    }
+
+    @Test
+    fun updateMap_partialOverlap() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 10.0.meters, "A")
+        val update = DistanceRangeMapImpl<String>()
+        update.put(5.0.meters, 15.0.meters, "B")
+        map.updateMap(update, { old, new -> old + new })
+        assertEquals("A", map.get(2.5.meters))
+        assertEquals("AB", map.get(7.5.meters))
+        assertEquals("B", map.get(12.5.meters))
+    }
+
+    @Test
+    fun updateMap_fullOverlap() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 10.0.meters, "A")
+        val update = DistanceRangeMapImpl<String>()
+        update.put(0.0.meters, 10.0.meters, "B")
+        map.updateMap(update, { old, new -> old + new })
+        assertEquals("AB", map.get(5.0.meters))
+    }
+
+    @Test
+    fun updateMap_multipleRanges() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 5.0.meters, "A")
+        map.put(10.0.meters, 15.0.meters, "C")
+        val update = DistanceRangeMapImpl<String>()
+        update.put(3.0.meters, 12.0.meters, "B")
+        map.updateMap(update, { old, new -> old + new })
+        assertEquals("A", map.get(1.0.meters))
+        assertEquals("AB", map.get(4.0.meters))
+        assertEquals("B", map.get(8.0.meters))
+        assertEquals("CB", map.get(11.0.meters))
+        assertEquals("C", map.get(14.0.meters))
+    }
+
+    @Test
+    fun updateMapKeepingNonIntersecting_emptyUpdate() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 10.0.meters, "A")
+        val update = DistanceRangeMapImpl<String>()
+        map.updateMap(update, { old, new -> old + new })
+        assertEquals("A", map.get(5.0.meters))
+    }
+
+    @Test
+    fun updateMap_emptyOriginal() {
+        val map = DistanceRangeMapImpl<String>()
+        val update = DistanceRangeMapImpl<String>()
+        update.put(0.0.meters, 10.0.meters, "B")
+        map.updateMap(update, { old, new -> old + new })
+        assertEquals("B", map.get(5.0.meters))
+    }
+
+    @Test
+    fun clear() {
+        val map = DistanceRangeMapImpl<String>()
+        map.put(0.0.meters, 10.0.meters, "A")
+        map.clear()
+        assertNull(map.get(5.0.meters))
+    }
 }
diff --git a/core/src/main/kotlin/fr/sncf/osrd/envelope_sim_infra/EnvelopeTrainPath.kt b/core/src/main/kotlin/fr/sncf/osrd/envelope_sim_infra/EnvelopeTrainPath.kt
@@ -58,7 +58,7 @@ object EnvelopeTrainPath {
         val res = mutableMapOf<String, DistanceRangeMap<Electrification>>()
         for (entry in profileMap.entries) {
             val electrificationMapWithProfiles = electrificationMap.clone()
-            electrificationMapWithProfiles.updateMap(entry.value) {
+            electrificationMapWithProfiles.updateMapIntersection(entry.value) {
                 obj: Electrification,
                 profile: String ->
                 obj.withElectricalProfile(profile)

diff --git a/core/src/main/kotlin/fr/sncf/osrd/utils/EnvelopeTrainPathUtils.kt b/core/src/main/kotlin/fr/sncf/osrd/utils/EnvelopeTrainPathUtils.kt
@@ -12,10 +12,12 @@ import fr.sncf.osrd.utils.units.Distance
 fun buildElectrificationMap(path: PathProperties): DistanceRangeMap<Electrification> {
     val res: DistanceRangeMap<Electrification> = DistanceRangeMapImpl()
     res.put(Distance.ZERO, path.getLength(), NonElectrified())
-    res.updateMap(path.getElectrification()) { _: Electrification?, electrificationMode: String ->
+    res.updateMapIntersection(path.getElectrification()) {
+        _: Electrification?,
+        electrificationMode: String ->
         if (electrificationMode == "") NonElectrified() else Electrified(electrificationMode)
     }
-    res.updateMap(path.getNeutralSections()) {
+    res.updateMapIntersection(path.getNeutralSections()) {
         electrification: Electrification?,
         neutralSection: NeutralSection ->
         Neutral(neutralSection.lowerPantograph, electrification, neutralSection.isAnnouncement)